from __future__ import annotations

import unittest
from datetime import datetime
from pathlib import Path
import tempfile

from datahub.adapters.tushare import TushareAdapter
from datahub.crypto import SecretVault
from datahub.db import HubDB
from datahub.pipeline import Pipeline
from datahub.scheduler import Scheduler
from datahub.settings import Settings
from datahub.timeutil import SHANGHAI
from tests.fixtures import fake_transport


class SchedulerTests(unittest.TestCase):
    def test_skips_eod_on_closed_day(self) -> None:
        tmp = tempfile.TemporaryDirectory()
        db = HubDB(Path(tmp.name) / "hub.db")
        adapter = TushareAdapter("x", transport=fake_transport)
        settings = Settings(encryption_key=SecretVault.generate_key(), scheduler_enabled=False, db_path=db.path)
        pipe = Pipeline(db, adapter, settings)
        pipe.ingest_reference("20240902")
        # 20240907 is closed in fixture
        ran = {"eod_a": 0}

        def fake_eod(_date: str):
            ran["eod_a"] += 1
            return {}

        sched = Scheduler(db, pipe, jobs={"precheck": lambda d: {}, "eod_a": fake_eod, "eod_b": lambda d: {}, "cleanup": lambda d: {}, "backup": lambda d: {}})
        clock = datetime(2024, 9, 7, 16, 0, tzinfo=SHANGHAI)
        fired = sched.tick(clock)
        self.assertNotIn("eod_a", fired)
        self.assertEqual(ran["eod_a"], 0)
        tmp.cleanup()

    def test_fires_eod_on_open_day(self) -> None:
        tmp = tempfile.TemporaryDirectory()
        db = HubDB(Path(tmp.name) / "hub.db")
        adapter = TushareAdapter("x", transport=fake_transport)
        settings = Settings(encryption_key=SecretVault.generate_key(), scheduler_enabled=False, db_path=db.path)
        pipe = Pipeline(db, adapter, settings)
        pipe.ingest_reference("20240902")
        ran = {"eod_a": 0}

        def fake_eod(_date: str):
            ran["eod_a"] += 1
            return {"rows": 1}

        sched = Scheduler(db, pipe, jobs={"precheck": lambda d: {}, "eod_a": fake_eod, "eod_b": lambda d: {}, "cleanup": lambda d: {}, "backup": lambda d: {}})
        clock = datetime(2024, 9, 2, 16, 0, tzinfo=SHANGHAI)
        fired = sched.tick(clock)
        self.assertIn("eod_a", fired)
        self.assertEqual(ran["eod_a"], 1)
        tmp.cleanup()


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