from __future__ import annotations

import unittest
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.settings import Settings
from datahub.serving import V1API
from tests.fixtures import TRADE_DATE, fake_transport

GROUP_A = ("daily", "valuation", "moneyflow", "auction")


class GroupTransport:
    """fake_transport with per-API degradation switches for release-group tests."""

    def __init__(self) -> None:
        self.empty: set[str] = set()
        self.keep_rows: dict[str, int] = {}
        self.stocks: list[dict] | None = None
        self.calls: list[str] = []

    def __call__(self, api_name: str, params: dict, fields: str):
        self.calls.append(api_name)
        if api_name in self.empty:
            return []
        if api_name == "stock_basic" and self.stocks is not None:
            return [dict(row) for row in self.stocks]
        rows = fake_transport(api_name, params, fields)
        keep = self.keep_rows.get(api_name)
        if keep is not None:
            return rows[:keep]
        return rows


def make_pipe(transport: GroupTransport, quality_extra: dict | None = None):
    tmp = tempfile.TemporaryDirectory()
    db = HubDB(Path(tmp.name) / "hub.db")
    adapter = TushareAdapter("test-token", transport=transport)
    quality = {
        "daily_row_ratio": 0.98,
        "null_rate_max": 0.01,
        "max_publish_attempts": 2,
        "publication_generations": 3,
    }
    if quality_extra:
        quality.update(quality_extra)
    settings = Settings(
        encryption_key=SecretVault.generate_key(),
        api_token="t" * 32,
        admin_password="admin-pass",
        tushare_token="test-token",
        db_path=db.path,
        quality=quality,
        scheduler_enabled=False,
    )
    pipe = Pipeline(db, adapter, settings)
    pipe._tmp = tmp
    return pipe, db


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


class ReleaseGroupSwitchTests(unittest.TestCase):
    def setUp(self) -> None:
        self.transport = GroupTransport()
        self.pipe, self.db = make_pipe(self.transport)
        self.pipe.ingest_reference(TRADE_DATE)

    def test_whole_group_switches_in_one_publish_instant(self) -> None:
        results = self.pipe.run_eod_batch_a(TRADE_DATE)
        self.assertEqual(set(results), {*GROUP_A, "stocks"})
        self.assertEqual({item["state"] for item in results.values()}, {"published"})
        pubs = self.db.fetchall("SELECT * FROM publications WHERE trade_date = ?", (TRADE_DATE,))
        self.assertEqual(len(pubs), 5)
        self.assertEqual(len({row["published_at"] for row in pubs}), 1)
        # official rows copied and serving resolves the new batches
        api = V1API(self.db, self.pipe, self.pipe.settings)
        payload = api.handle("/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]})
        self.assertEqual(payload["meta"]["batch_id"], results["daily"]["batch_id"])
        stocks = api.handle("/v1/stocks", {})
        self.assertEqual(stocks["meta"]["batch_id"], results["stocks"]["batch_id"])

    def test_any_member_failure_blocks_entire_group(self) -> None:
        self.transport.empty = {"daily_basic"}  # valuation upstream returns nothing
        results = self.pipe.run_eod_batch_a(TRADE_DATE)
        self.assertEqual(results["valuation"]["state"], "failed")
        self.assertEqual(results["moneyflow"]["state"], "aborted")
        self.assertEqual(results["auction"]["state"], "aborted")
        self.assertEqual(results["daily"]["state"], "failed")  # staged fine, then abandoned
        # nothing became visible, and the reason is recorded
        self.assertEqual(publications_map(self.db, TRADE_DATE), {})
        abandoned = self.db.fetchall(
            "SELECT * FROM batches WHERE trade_date = ? AND state = 'failed'",
            (TRADE_DATE,),
        )
        self.assertTrue(any("release group not switched" in str(row["error"] or "") for row in abandoned))
        audit = self.db.fetchone(
            "SELECT * FROM audit_log WHERE action = 'release-group' ORDER BY id DESC"
        )
        self.assertIn("valuation", str(audit["detail"]))
        # still missing → evening retries keep trying
        self.assertIn("daily", self.pipe.missing_official_datasets(TRADE_DATE))

    def test_failure_keeps_previous_complete_version_serving(self) -> None:
        first = self.pipe.run_dataset("daily", TRADE_DATE)
        self.transport.empty = {"daily_basic"}
        results = self.pipe.run_eod_missing(TRADE_DATE)
        # incomplete A-group restages daily with the others; valuation fails → no A switch
        self.assertEqual(results["daily"]["state"], "failed")
        self.assertEqual(results["valuation"]["state"], "failed")
        # the already-published daily batch is untouched and keeps serving
        self.assertEqual(self.pipe.active_batch("daily", TRADE_DATE), first["batch_id"])
        pubs = publications_map(self.db, TRADE_DATE)
        self.assertEqual(pubs["daily"], first["batch_id"])
        self.assertNotIn("valuation", pubs)
        self.assertNotIn("moneyflow", pubs)
        self.assertNotIn("auction", pubs)
        # B-group is an independent boundary and may still publish
        self.assertEqual(results["index_daily"]["state"], "published")
        payload = V1API(self.db, self.pipe, self.pipe.settings).handle(
            "/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]}
        )
        self.assertEqual(payload["meta"]["batch_id"], first["batch_id"])

    def test_partial_group_retry_does_not_mix_batches(self) -> None:
        """Already-published A members must be restaged with missing ones."""
        first_daily = self.pipe.run_dataset("daily", TRADE_DATE)
        first_moneyflow = self.pipe.run_dataset("moneyflow", TRADE_DATE)
        results = self.pipe.run_eod_missing(TRADE_DATE)
        # A-group switched as one boundary; B-group (index) also published
        for name in (*GROUP_A, "stocks"):
            self.assertEqual(results[name]["state"], "published", name)
        self.assertEqual(results["index_daily"]["state"], "published")
        pubs = self.db.fetchall(
            "SELECT dataset, active_batch, published_at FROM publications WHERE trade_date = ?",
            (TRADE_DATE,),
        )
        by_ds = {str(row["dataset"]): row for row in pubs}
        # old partial batches replaced — no cross-batch mix of the first wave
        self.assertNotEqual(by_ds["daily"]["active_batch"], first_daily["batch_id"])
        self.assertNotEqual(by_ds["moneyflow"]["active_batch"], first_moneyflow["batch_id"])
        a_times = {by_ds[name]["published_at"] for name in (*GROUP_A, "stocks")}
        self.assertEqual(len(a_times), 1)
        # serving resolves the new complete A-group batches
        api = V1API(self.db, self.pipe, self.pipe.settings)
        daily = api.handle("/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]})
        self.assertEqual(daily["meta"]["batch_id"], results["daily"]["batch_id"])
        self.assertEqual(daily["meta"]["batch_id"], by_ds["daily"]["active_batch"])

    def test_reads_during_switch_see_old_state_until_commit(self) -> None:
        snapshots: list[dict] = []

        def watcher() -> None:
            with self.db.connect() as connection:
                rows = connection.execute(
                    "SELECT dataset, active_batch FROM publications WHERE trade_date = ?",
                    (TRADE_DATE,),
                ).fetchall()
                snapshots.append({str(row["dataset"]): row["active_batch"] for row in rows})

        self.pipe.before_commit = watcher
        self.pipe.run_eod_batch_a(TRADE_DATE)
        # inside the switch transaction the group was still invisible
        self.assertEqual(snapshots[0], {})
        after = publications_map(self.db, TRADE_DATE)
        self.assertEqual(set(after), {*GROUP_A, "stocks"})

    def test_switch_crash_rolls_back_whole_group(self) -> None:
        def explode() -> None:
            raise RuntimeError("killed mid-switch")

        self.pipe.before_commit = explode
        with self.assertRaises(RuntimeError):
            self.pipe.run_eod_batch_a(TRADE_DATE)
        self.assertEqual(publications_map(self.db, TRADE_DATE), {})
        for table in ("eod_bars", "eod_valuation", "eod_moneyflow", "eod_auction", "eod_stocks"):
            rows = self.db.fetchall(f"SELECT * FROM {table} WHERE trade_date = ?", (TRADE_DATE,))
            self.assertEqual(rows, [], table)

    def test_duplicate_runs_are_idempotent(self) -> None:
        self.pipe.run_eod_batch_a(TRADE_DATE)
        self.pipe.run_eod_batch_b(TRADE_DATE)
        batches_before = {
            str(row["batch_id"])
            for row in self.db.fetchall("SELECT batch_id FROM batches WHERE trade_date = ?", (TRADE_DATE,))
        }
        calls_before = len(self.transport.calls)
        again = self.pipe.run_eod_missing(TRADE_DATE)
        self.assertEqual({item["state"] for item in again.values()}, {"skipped"})
        self.assertEqual({item["reason"] for item in again.values()}, {"already_published"})
        batches_after = {
            str(row["batch_id"])
            for row in self.db.fetchall("SELECT batch_id FROM batches WHERE trade_date = ?", (TRADE_DATE,))
        }
        self.assertEqual(batches_after, batches_before)
        self.assertEqual(len(self.transport.calls), calls_before)
        self.assertEqual(self.pipe.missing_official_datasets(TRADE_DATE), [])

    def test_cross_gate_failure_blocks_switch(self) -> None:
        transport = GroupTransport()
        pipe, db = make_pipe(
            transport,
            quality_extra={"cross_gates": [
                {"left": "daily", "right": "moneyflow", "min_key_overlap": 1.0},
            ]},
        )
        pipe.ingest_reference(TRADE_DATE)
        transport.keep_rows["moneyflow"] = 1  # moneyflow covers only half the market
        results = pipe.run_eod_batch_a(TRADE_DATE)
        self.assertEqual(results["moneyflow"]["state"], "failed")
        self.assertIn("cross gate", str(results["moneyflow"]["error"]))
        self.assertEqual(publications_map(db, TRADE_DATE), {})

    def test_stocks_master_and_snapshot_switch_together_or_not_at_all(self) -> None:
        original = [
            {"ts_code": "600000.SH", "symbol": "600000", "name": "浦发银行", "area": "上海",
             "industry": "银行", "market": "主板", "list_status": "L", "list_date": "19991110"},
            {"ts_code": "920071.BJ", "symbol": "920071", "name": "N金钛", "area": "辽宁",
             "industry": "小金属", "market": "北交所", "list_status": "L", "list_date": "20240901"},
        ]
        renamed = [dict(original[0]), {**original[1], "name": "金钛股份"}]
        self.transport.stocks = renamed
        self.pipe.run_eod_batch_a(TRADE_DATE)
        master = self.db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
        self.assertEqual(master["name"], "金钛股份")
        stocks_pub = self.db.fetchone(
            "SELECT active_batch FROM publications WHERE dataset = 'stocks' AND trade_date = ?",
            (TRADE_DATE,),
        )
        self.assertIsNotNone(stocks_pub)

        # failure path: rename staged but the group is blocked → master stays untouched
        transport = GroupTransport()
        transport.stocks = original
        pipe, db = make_pipe(
            transport,
            quality_extra={"cross_gates": [
                {"left": "daily", "right": "moneyflow", "min_key_overlap": 1.0},
            ]},
        )
        pipe.ingest_reference(TRADE_DATE)  # master seeded with "N金钛"
        transport.stocks = renamed
        transport.keep_rows["moneyflow"] = 1
        results = pipe.run_eod_batch_a(TRADE_DATE)
        self.assertEqual(results["stocks"]["state"], "failed")
        master = db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
        self.assertEqual(master["name"], "N金钛")  # rename not applied
        stocks_pub = db.fetchone(
            "SELECT active_batch FROM publications WHERE dataset = 'stocks' AND trade_date = ?",
            (TRADE_DATE,),
        )
        self.assertIsNone(stocks_pub)


class StocksRefreshAtomicTests(unittest.TestCase):
    def setUp(self) -> None:
        self.transport = GroupTransport()
        self.pipe, self.db = make_pipe(self.transport)
        self.pipe.ingest_reference(TRADE_DATE)
        self.transport.stocks = [
            {"ts_code": "600000.SH", "symbol": "600000", "name": "浦发银行", "area": "上海",
             "industry": "银行", "market": "主板", "list_status": "L", "list_date": "19991110"},
            {"ts_code": "920071.BJ", "symbol": "920071", "name": "N金钛", "area": "辽宁",
             "industry": "小金属", "market": "北交所", "list_status": "L", "list_date": "20240901"},
        ]
        first = self.pipe.refresh_stocks(TRADE_DATE)
        self.assertEqual(first["state"], "published")
        self.first_batch = first["batch_id"]

    def test_refresh_keeps_master_when_snapshot_publish_fails(self) -> None:
        self.transport.stocks = [
            {"ts_code": "600000.SH", "symbol": "600000", "name": "浦发银行", "area": "上海",
             "industry": "银行", "market": "主板", "list_status": "L", "list_date": "19991110"},
            {"ts_code": "920071.BJ", "symbol": "920071", "name": "金钛股份", "area": "辽宁",
             "industry": "小金属", "market": "北交所", "list_status": "L", "list_date": "20240901"},
        ]

        def explode() -> None:
            raise RuntimeError("snapshot switch killed")

        self.pipe.before_commit = explode
        with self.assertRaises(RuntimeError):
            self.pipe.refresh_stocks(TRADE_DATE)
        master = self.db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
        self.assertEqual(master["name"], "N金钛")  # rename not applied
        self.assertEqual(self.pipe.active_batch("stocks", TRADE_DATE), self.first_batch)
        audit = self.db.fetchone(
            "SELECT * FROM audit_log WHERE action = 'stocks-refresh' ORDER BY id DESC"
        )
        self.assertIn("failed", str(audit["detail"]))
        self.assertIn("snapshot switch killed", str(audit["detail"]))

    def test_refresh_keeps_master_when_quality_gate_rejects(self) -> None:
        self.transport.stocks = []  # empty → hard fail before publish
        with self.assertRaises(Exception):
            self.pipe.refresh_stocks(TRADE_DATE)
        master = self.db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
        self.assertEqual(master["name"], "N金钛")
        self.assertEqual(self.pipe.active_batch("stocks", TRADE_DATE), self.first_batch)
        audit = self.db.fetchone(
            "SELECT * FROM audit_log WHERE action = 'stocks-refresh' ORDER BY id DESC"
        )
        self.assertIn("failed", str(audit["detail"]))


class ForceBoundaryEntryTests(unittest.TestCase):
    """CLI force / admin backfill must rebuild the full A/B boundary."""

    def setUp(self) -> None:
        self.transport = GroupTransport()
        self.pipe, self.db = make_pipe(self.transport)
        self.pipe.ingest_reference(TRADE_DATE)
        self.first = self.pipe.run_eod_batch_a(TRADE_DATE)
        self.pipe.run_eod_batch_b(TRADE_DATE)

    def test_force_republish_valuation_rebuilds_whole_a_group(self) -> None:
        before = publications_map(self.db, TRADE_DATE)
        results = self.pipe.force_republish_boundary("valuation", TRADE_DATE)
        self.assertEqual({item["state"] for item in results.values()}, {"published"})
        after = publications_map(self.db, TRADE_DATE)
        for name in (*GROUP_A, "stocks"):
            self.assertNotEqual(after[name], before[name], name)
            self.assertEqual(after[name], results[name]["batch_id"], name)
        # B-group left alone
        self.assertEqual(after["index_daily"], before["index_daily"])
        pubs = self.db.fetchall(
            "SELECT dataset, published_at FROM publications WHERE trade_date = ?",
            (TRADE_DATE,),
        )
        a_times = {row["published_at"] for row in pubs if row["dataset"] in {*GROUP_A, "stocks"}}
        self.assertEqual(len(a_times), 1)

    def test_force_republish_index_rebuilds_only_b_group(self) -> None:
        before = publications_map(self.db, TRADE_DATE)
        results = self.pipe.force_republish_boundary("index_daily", TRADE_DATE)
        self.assertEqual(results["index_daily"]["state"], "published")
        after = publications_map(self.db, TRADE_DATE)
        self.assertNotEqual(after["index_daily"], before["index_daily"])
        for name in GROUP_A:
            self.assertEqual(after[name], before[name], name)

    def test_admin_backfill_official_dataset_uses_boundary(self) -> None:
        from datahub.admin_api import AdminAPI
        from datahub.auth import AuthService
        from datahub.crypto import SecretVault
        from datahub.scheduler import Scheduler
        from datahub.serving import ApiError

        vault = SecretVault(self.pipe.settings.encryption_key)
        auth = AuthService(self.db, vault, self.pipe.settings.api_token, "StartPass1")
        admin = AdminAPI(self.db, self.pipe, Scheduler(self.db, self.pipe), auth)
        before = publications_map(self.db, TRADE_DATE)
        result = admin.backfill("moneyflow", TRADE_DATE, "StartPass1", f"moneyflow:{TRADE_DATE}", "tester")
        self.assertEqual(result["moneyflow"]["state"], "published")
        after = publications_map(self.db, TRADE_DATE)
        for name in (*GROUP_A, "stocks"):
            self.assertNotEqual(after[name], before[name], name)
        # bad password / wrong confirm still rejected
        with self.assertRaises(ApiError):
            admin.backfill("daily", TRADE_DATE, "wrong", f"daily:{TRADE_DATE}", "tester")

    def test_admin_backfill_switch_crash_is_failed_precondition(self) -> None:
        from datahub.admin_api import AdminAPI
        from datahub.auth import AuthService
        from datahub.crypto import SecretVault
        from datahub.scheduler import Scheduler
        from datahub.serving import ApiError

        vault = SecretVault(self.pipe.settings.encryption_key)
        auth = AuthService(self.db, vault, self.pipe.settings.api_token, "StartPass1")
        admin = AdminAPI(self.db, self.pipe, Scheduler(self.db, self.pipe), auth)
        before = publications_map(self.db, TRADE_DATE)

        def explode() -> None:
            raise RuntimeError("killed mid-switch")

        self.pipe.before_commit = explode
        with self.assertRaises(ApiError) as ctx:
            admin.backfill("valuation", TRADE_DATE, "StartPass1", f"valuation:{TRADE_DATE}", "tester")
        self.assertEqual(ctx.exception.code, "FAILED_PRECONDITION")
        self.assertIn("killed mid-switch", ctx.exception.message)
        # previous complete A/B versions keep serving
        self.assertEqual(publications_map(self.db, TRADE_DATE), before)


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