from __future__ import annotations

import unittest
from datetime import datetime
from tempfile import TemporaryDirectory

from sdk import OrderItem, PositionItem
from strategy.zt.state import BUYING, DONE, SELLING, SOLD, TState, TStateItem


class ZTStateTests(unittest.TestCase):
    def test_reconcile_marks_sell_and_buy_orders_completed(self):
        with TemporaryDirectory() as directory:
            state = TState.for_strategy(directory, "zt", "A")
            state.set(TStateItem("A", 1000, 10, "2026-08-31", SELLING, "sell-1", 500, 11))
            position = PositionItem(stock_code="A", volume=500, open_price=10)
            state.reconcile([position], [OrderItem("1", "A", "SELL", "", "56", datetime.now(), 500, "sell-1")], "2026-08-31")
            self.assertEqual(state.get("A").phase, SOLD)

            item = state.get("A")
            item.phase, item.buy_order_id = BUYING, "buy-1"
            state.set(item)
            state.reconcile([position], [OrderItem("2", "A", "BUY", "", "56", datetime.now(), 500, "buy-1")], "2026-08-31")
            self.assertEqual(state.get("A").phase, DONE)

    def test_new_position_becomes_dcm_base_state(self):
        with TemporaryDirectory() as directory:
            state = TState.for_strategy(directory, "zt", "A")
            state.reconcile([PositionItem(stock_code="A", volume=800, open_price=12.5)], [], "2026-08-31")
            item = state.get("A")
            self.assertEqual((item.base_qty, item.base_cost), (800, 12.5))


if __name__ == "__main__":
    unittest.main()