import sqlite3
import tempfile
import unittest
from contextlib import closing
from pathlib import Path

from libs.state import FLAG_BUY, FLAG_SELL, State
from sdk import DealItem, PositionItem


class SnapshotArchiveTests(unittest.TestCase):
    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.store = State(Path(tmp.name) / 'state.db')
        self.code = '600000.SH'

    def deal(self, identity, qty, flag=FLAG_BUY, code=None):
        return DealItem(
            stock_code=code or self.code, order_sys_id=identity,
            remark=f'zt-base-{identity}|zt', offset_flag=flag,
            volume=qty, price=12, trade_amount=qty * 12,
            trade_date='2026-09-12', trade_time='100000',
        )

    def snapshot(self, qty):
        self.store.sync_state([PositionItem(stock_code=self.code, volume=qty, open_price=10)])

    def test_matching_partial_fills_only_mark_and_survive_restart(self):
        self.snapshot(100)
        saved = dict(self.store.state[self.code])
        self.store.sync_deals([self.deal('one', 40), self.deal('two', 60)])
        self.store.archiving()
        self.assertEqual(self.store.state[self.code], saved)
        self.assertTrue(all(d['is_arch'] == 1 for d in self.store.deals.values()))
        restarted = State(self.store.path)
        restarted.archiving()
        self.assertEqual(restarted.state[self.code], saved)

    def test_total_includes_added_holdings(self):
        self.snapshot(100)
        with closing(self.store._connect()) as db, db:
            db.execute("UPDATE state SET added_qty=50, added_price=9, status='CUSTOM'")
        self.store.load()
        saved = dict(self.store.state[self.code])
        self.store.sync_deals([self.deal('one', 150)])
        self.store.archiving()
        self.assertEqual(self.store.state[self.code], saved)
        self.assertEqual(self.store.deals['one']['is_arch'], 1)

    def test_nonmatching_and_other_stock_are_incremental(self):
        self.snapshot(100)
        self.store.sync_deals([self.deal('one', 40), self.deal('other', 60, code='600001.SH')])
        self.store.archiving()
        self.assertEqual(self.store.state[self.code]['base_qty'], 140)
        self.assertEqual(self.store.state['600001.SH']['base_qty'], 60)

    def test_matching_sell_still_liquidates(self):
        self.snapshot(100)
        self.store.sync_deals([self.deal('sell', 100, FLAG_SELL)])
        self.store.archiving()
        self.assertNotIn(self.code, self.store.state)
        self.assertEqual(self.store.deals['sell']['is_arch'], 1)

    def test_mixed_batch_with_matching_gross_volume_is_not_skipped(self):
        self.snapshot(100)
        self.store.sync_deals([self.deal('buy', 40), self.deal('sell', 60, FLAG_SELL)])
        self.store.archiving()
        self.assertEqual(self.store.state[self.code]['base_qty'], 80)
        self.assertTrue(all(d['is_arch'] == 1 for d in self.store.deals.values()))

    def test_failed_mark_rolls_back_entire_stock_and_retries(self):
        self.snapshot(100)
        saved = dict(self.store.state[self.code])
        self.store.sync_deals([self.deal('one', 40), self.deal('two', 60)])
        with closing(self.store._connect()) as db, db:
            db.execute("""CREATE TRIGGER fail_mark BEFORE UPDATE OF is_arch ON deals
                WHEN OLD.order_sys_id='two'
                BEGIN SELECT RAISE(ABORT, 'test failure'); END""")
        with self.assertLogs(level='WARNING'):
            self.store.archiving()
        self.assertEqual(self.store.state[self.code], saved)
        self.assertTrue(all(d['is_arch'] == 0 for d in self.store.deals.values()))
        with closing(self.store._connect()) as db, db:
            db.execute('DROP TRIGGER fail_mark')
        self.store.archiving()
        self.assertEqual(self.store.state[self.code], saved)
        self.assertTrue(all(d['is_arch'] == 1 for d in self.store.deals.values()))


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