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

OFFICIAL = {"daily", "valuation", "moneyflow", "auction", "index_daily"}


class DelayedTransport:
    """Upstream that only returns rows for dates it has "published" yet."""

    DATE_APIS = {"daily", "daily_basic", "adj_factor", "moneyflow", "stk_auction", "index_daily"}

    def __init__(self, ready_dates: set[str]) -> None:
        self.ready = set(ready_dates)
        self.calls: list[str] = []

    def __call__(self, api_name: str, params: dict, fields: str):
        self.calls.append(api_name)
        if api_name in self.DATE_APIS:
            trade_date = str(params.get("trade_date") or "")
            if trade_date and trade_date not in self.ready:
                return []
        return fake_transport(api_name, params, fields)


def clock_at(day: str, hh: int, mm: int) -> datetime:
    return datetime(int(day[:4]), int(day[4:6]), int(day[6:8]), hh, mm, tzinfo=SHANGHAI)


class EodRetryTests(unittest.TestCase):
    def _make(self, ready_dates: set[str]):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        db = HubDB(Path(tmp.name) / "hub.db")
        transport = DelayedTransport(ready_dates)
        adapter = TushareAdapter("x", transport=transport)
        settings = Settings(
            encryption_key=SecretVault.generate_key(),
            db_path=db.path,
            backup_dir=Path(tmp.name) / "backups",
        )
        pipe = Pipeline(db, adapter, settings)
        pipe.ingest_reference("20240902")
        sched = Scheduler(db, pipe)
        return db, transport, pipe, sched

    def _job_runs(self, db: HubDB, job_id: str) -> list[dict]:
        return db.fetchall("SELECT * FROM job_runs WHERE job_id = ? ORDER BY id", (job_id,))

    def _batches(self, db: HubDB, day: str) -> list[dict]:
        return db.fetchall("SELECT * FROM batches WHERE trade_date = ?", (day,))

    @staticmethod
    def _batch_ids(db: HubDB, day: str) -> set[str]:
        return {str(row["batch_id"]) for row in db.fetchall("SELECT batch_id FROM batches WHERE trade_date = ?", (day,))}

    @staticmethod
    def _eod_calls(transport: DelayedTransport) -> list[str]:
        return [name for name in transport.calls if name in DelayedTransport.DATE_APIS]

    def _published(self, db: HubDB, day: str) -> set[str]:
        rows = db.fetchall("SELECT dataset FROM publications WHERE trade_date = ?", (day,))
        return {str(row["dataset"]) for row in rows}

    def test_first_empty_then_retry_succeeds(self) -> None:
        day = "20240902"
        db, transport, pipe, sched = self._make(set())

        sched.tick(clock_at(day, 15, 5))   # eod_a: upstream empty -> failed
        sched.tick(clock_at(day, 15, 10))  # eod_b: upstream empty -> failed
        self.assertEqual(self._published(db, day), set())  # quality gate held

        sched.tick(clock_at(day, 15, 20))  # inside window, but <30min since 15:10
        self.assertEqual(self._job_runs(db, "eod_retry"), [])
        status = sched.eod_status(day, clock=clock_at(day, 15, 20))
        self.assertEqual(status["state"], "waiting_upstream")
        self.assertTrue(status["next_retry_at"])
        self.assertEqual(status["missing_datasets"], sorted(OFFICIAL))

        sched.tick(clock_at(day, 15, 40))  # retry #1, still empty
        runs = self._job_runs(db, "eod_retry")
        self.assertEqual(len(runs), 1)
        self.assertEqual(runs[0]["state"], "failed")
        self.assertEqual(self._published(db, day), set())

        transport.ready.add(day)
        sched.tick(clock_at(day, 16, 10))  # retry #2 succeeds
        self.assertEqual(self._published(db, day), OFFICIAL)
        self.assertEqual(sched.eod_status(day, clock=clock_at(day, 16, 10))["state"], "done")
        progress = db.fetchone("SELECT * FROM eod_progress WHERE trade_date = ?", (day,))
        self.assertEqual(progress["state"], "done")
        self.assertEqual(progress["attempts"], 4)  # eod_a + eod_b + 2 retries

        # success stops all further same-day requests
        batches_before = len(self._batches(db, day))
        calls_before = len(transport.calls)
        sched.tick(clock_at(day, 17, 0))
        sched.tick(clock_at(day, 23, 0))
        self.assertEqual(len(self._job_runs(db, "eod_retry")), 2)
        self.assertEqual(len(self._batches(db, day)), batches_before)
        self.assertEqual(len(transport.calls), calls_before)

    def test_never_ready_marks_cutoff_failed_and_stops(self) -> None:
        day = "20240902"
        db, transport, pipe, sched = self._make(set())
        sched.tick(clock_at(day, 15, 5))
        sched.tick(clock_at(day, 15, 10))
        sched.tick(clock_at(day, 15, 40))
        sched.tick(clock_at(day, 16, 10))
        sched.tick(clock_at(day, 23, 29))
        self.assertEqual(len(self._job_runs(db, "eod_retry")), 3)

        sched.tick(clock_at(day, 23, 35))  # past cutoff 23:30
        status = sched.eod_status(day, clock=clock_at(day, 23, 35))
        self.assertEqual(status["state"], "cutoff_failed")
        cutoff_runs = [r for r in self._job_runs(db, "eod_retry") if "截止" in str(r["error"])]
        self.assertEqual(len(cutoff_runs), 1)
        self.assertEqual(self._published(db, day), set())

        attempts = db.fetchone("SELECT attempts FROM eod_progress WHERE trade_date = ?", (day,))["attempts"]
        sched.tick(clock_at(day, 23, 59))
        self.assertEqual(
            db.fetchone("SELECT attempts FROM eod_progress WHERE trade_date = ?", (day,))["attempts"],
            attempts,
        )
        self.assertEqual(len(self._job_runs(db, "eod_retry")), 4)  # 3 retries + 1 cutoff record
        self.assertEqual(self._published(db, day), set())

    def test_restart_catches_up_without_overwriting(self) -> None:
        day = "20240902"
        db, transport, pipe, sched = self._make({day})
        sched.tick(clock_at(day, 15, 5))   # eod_a publishes 4 datasets
        sched.tick(clock_at(day, 15, 10))  # eod_b publishes index
        self.assertEqual(self._published(db, day), OFFICIAL)
        active = db.fetchall("SELECT dataset, active_batch FROM publications WHERE trade_date = ?", (day,))
        active_map = {row["dataset"]: row["active_batch"] for row in active}
        batches_before = self._batch_ids(db, day)
        calls_before = self._eod_calls(transport)

        # container restart: fresh scheduler, missed-time catch-up fires eod_a/eod_b
        sched2 = Scheduler(db, pipe)
        ran = sched2.tick(clock_at(day, 21, 0))
        self.assertIn("eod_a", ran)
        self.assertIn("eod_b", ran)
        self.assertNotIn("eod_retry", ran)
        self.assertEqual(self._published(db, day), OFFICIAL)
        after = db.fetchall("SELECT dataset, active_batch FROM publications WHERE trade_date = ?", (day,))
        self.assertEqual({row["dataset"]: row["active_batch"] for row in after}, active_map)
        self.assertEqual(self._batch_ids(db, day), batches_before)  # no duplicate batches
        self.assertEqual(self._eod_calls(transport), calls_before)  # no duplicate upstream EOD calls
        self.assertEqual(sched2.eod_status(day, clock=clock_at(day, 21, 0))["state"], "done")

    def test_restart_with_partial_publish_only_fetches_missing(self) -> None:
        day = "20240902"
        db, transport, pipe, sched = self._make({day})
        sched.tick(clock_at(day, 15, 5))  # eod_a publishes 4; container "crashes" before eod_b
        self.assertEqual(self._published(db, day), {"daily", "valuation", "moneyflow", "auction"})
        batches_before = self._batch_ids(db, day)

        sched2 = Scheduler(db, pipe)
        ran = sched2.tick(clock_at(day, 15, 20))  # restart: eod_b catch-up, eod_a all skipped
        self.assertIn("eod_b", ran)
        self.assertEqual(self._published(db, day), OFFICIAL)
        new_ids = self._batch_ids(db, day) - batches_before
        new_datasets = {str(b["dataset"]) for b in self._batches(db, day) if str(b["batch_id"]) in new_ids}
        self.assertEqual(new_datasets, {"index_daily"})
        self.assertEqual(sched2.eod_status(day, clock=clock_at(day, 15, 20))["state"], "done")

    def test_closed_day_skips_all_eod_work(self) -> None:
        day = "20240907"  # closed in fixture calendar
        db, transport, pipe, sched = self._make(set())
        for hh, mm in ((15, 5), (15, 10), (15, 40), (16, 10), (20, 0), (23, 40)):
            ran = sched.tick(clock_at(day, hh, mm))
            self.assertNotIn("eod_retry", ran)
        eod_runs = db.fetchall("SELECT * FROM job_runs WHERE job_id LIKE 'eod%'")
        self.assertEqual(eod_runs, [])
        self.assertIsNone(db.fetchone("SELECT * FROM eod_progress WHERE trade_date = ?", (day,)))
        self.assertEqual(self._published(db, day), set())
        self.assertEqual(sched.eod_status(day, clock=clock_at(day, 20, 0))["state"], "closed_day")

    def test_duplicate_and_concurrent_execution_are_safe(self) -> None:
        day = "20240902"
        db, transport, pipe, sched = self._make({day})
        sched.tick(clock_at(day, 15, 5))
        sched.tick(clock_at(day, 15, 10))
        self.assertEqual(self._published(db, day), OFFICIAL)
        batches_before = len(self._batches(db, day))
        calls_before = len(transport.calls)

        out = sched.run_job("eod_retry", day)  # manual duplicate run
        self.assertEqual(out["state"], "ok")
        self.assertEqual(len(self._batches(db, day)), batches_before)
        self.assertEqual(len(transport.calls), calls_before)

        sched._eod_lock.acquire()  # simulate an in-flight EOD job
        try:
            busy = sched.run_job("eod_retry", day)
            self.assertEqual(busy["state"], "skipped")
            busy_a = sched.run_job("eod_a", day)
            self.assertEqual(busy_a["state"], "skipped")
        finally:
            sched._eod_lock.release()
        self.assertEqual(len(self._batches(db, day)), batches_before)


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