from __future__ import annotations

import tempfile
import unittest
from datetime import datetime, timedelta
from pathlib import Path

from datahub.adapters.tushare import TushareAdapter
from datahub.crypto import SecretVault
from datahub.db import HubDB
from datahub.pipeline import EMPTY_BATCH_ERROR, Pipeline, QualityError
from datahub.serving import ApiError, V1API
from datahub.settings import Settings
from datahub.timeutil import SHANGHAI, isoformat
from tests.fixtures import TRADE_DATE, fake_transport

DATASET_API = {
    "daily": "daily",
    "valuation": "daily_basic",
    "moneyflow": "moneyflow",
    "auction": "stk_auction",
    "index_daily": "index_daily",
}


def empty_transport_for(*datasets: str):
    blocked = {DATASET_API[name] for name in datasets}
    if "daily" in datasets:
        blocked.add("adj_factor")

    def transport(api_name, params, fields):
        if api_name in blocked:
            return []
        return fake_transport(api_name, params, fields)

    return transport


def make_pipeline(before_commit=None, clock=None, quality=None) -> tuple[Pipeline, HubDB]:
    tmp = tempfile.TemporaryDirectory()
    db = HubDB(Path(tmp.name) / "hub.db")
    adapter = TushareAdapter("test-token", transport=fake_transport)
    quality_cfg = {
        "daily_row_ratio": 0.98,
        "null_rate_max": 0.01,
        "max_publish_attempts": 3,
        "publication_generations": 3,
        "job_run_retain_days": 90,
        "staging_retain_days": 14,
    }
    if quality:
        quality_cfg.update(quality)
    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_cfg,
        scheduler_enabled=False,
    )
    pipe = Pipeline(db, adapter, settings, before_commit=before_commit, clock=clock)
    pipe._tmp = tmp  # keep alive
    return pipe, db


class PipelineTests(unittest.TestCase):
    def test_reference_and_daily_publish(self) -> None:
        pipe, db = make_pipeline()
        ref = pipe.ingest_reference(TRADE_DATE)
        self.assertEqual(ref["stocks"], 2)
        self.assertEqual(ref["calendar_from"], "20160101")
        result = pipe.run_dataset("daily", TRADE_DATE)
        self.assertEqual(result["state"], "published")
        self.assertEqual(result["rows"], 2)
        pub = db.fetchone("SELECT * FROM publications WHERE dataset='daily' AND trade_date=?", (TRADE_DATE,))
        self.assertEqual(pub["active_batch"], result["batch_id"])
        rows = db.fetchall("SELECT * FROM eod_bars WHERE batch_id=?", (result["batch_id"],))
        self.assertEqual(len(rows), 2)
        self.assertEqual(rows[0]["amount"] if rows[0]["ts_code"] == "600000.SH" else rows[1]["amount"], 2_000_000.0)

    def test_atomic_publish_abort_leaves_no_half_batch(self) -> None:
        pipe, db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        first = pipe.run_dataset("daily", TRADE_DATE)
        boom = {"n": 0}

        def explode() -> None:
            boom["n"] += 1
            raise RuntimeError("killed")

        pipe.before_commit = explode
        with self.assertRaises(RuntimeError):
            pipe.run_dataset("daily", TRADE_DATE)
        pub = db.fetchone("SELECT * FROM publications WHERE dataset='daily' AND trade_date=?", (TRADE_DATE,))
        self.assertEqual(pub["active_batch"], first["batch_id"])
        visible = db.fetchall(
            "SELECT DISTINCT batch_id FROM eod_bars WHERE trade_date=? AND batch_id=?",
            (TRADE_DATE, pub["active_batch"]),
        )
        self.assertEqual(len(visible), 1)

    def test_rollback_switches_active_batch(self) -> None:
        pipe, _db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        first = pipe.run_dataset("daily", TRADE_DATE)
        second = pipe.run_dataset("daily", TRADE_DATE)
        self.assertNotEqual(first["batch_id"], second["batch_id"])
        rolled = pipe.rollback("daily", TRADE_DATE, actor="test")
        self.assertEqual(rolled["active_batch"], first["batch_id"])
        from datahub.serving import V1API

        api = V1API(pipe.db, pipe, pipe.settings)
        payload = api.handle("/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]})
        self.assertEqual(payload["meta"]["batch_id"], first["batch_id"])

    def test_row_ratio_gate_rejects_short_batch(self) -> None:
        pipe, _db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        original = fake_transport

        def short(api_name, params, fields):
            rows = original(api_name, params, fields)
            if api_name == "daily":
                return rows[:1]
            return rows

        pipe.adapter._transport = short
        with self.assertRaises(QualityError) as ctx:
            pipe.run_dataset("daily", TRADE_DATE)
        self.assertTrue(ctx.exception.report["hard_fail"])
        pub = pipe.db.fetchone("SELECT * FROM publications WHERE dataset='daily'")
        self.assertIsNone(pub)

    def test_wal_mode(self) -> None:
        pipe, db = make_pipeline()
        with db.connect() as connection:
            mode = connection.execute("PRAGMA journal_mode").fetchone()[0]
        self.assertEqual(str(mode).lower(), "wal")

    def test_cleanup_iso_timestamps_respect_retention_on_job_and_src(self) -> None:
        # job_runs.started_at / src_calls.created_at 存 ISO；旧实现用 YYYYMMDD 比较会误删同年记录。
        frozen = datetime(2026, 9, 2, 0, 30, tzinfo=SHANGHAI)
        retain_days = 90
        pipe, db = make_pipeline(clock=lambda: frozen, quality={"job_run_retain_days": retain_days})
        samples = {
            "today": isoformat(frozen),
            "within": isoformat(frozen - timedelta(days=retain_days - 1)),
            "expired": isoformat(frozen - timedelta(days=retain_days + 1)),
        }
        with db.write() as connection:
            for job_id, stamp in samples.items():
                connection.execute(
                    "INSERT INTO job_runs(job_id, state, started_at, attempt) VALUES (?,?,?,1)",
                    (job_id, "ok", stamp),
                )
                connection.execute(
                    "INSERT INTO src_calls(provider, endpoint, ok, latency_ms, error, created_at) VALUES (?,?,?,?,?,?)",
                    ("tushare", job_id, 1, 10, None, stamp),
                )

        pipe.cleanup()

        jobs = {row["job_id"] for row in db.fetchall("SELECT job_id FROM job_runs")}
        calls = {row["endpoint"] for row in db.fetchall("SELECT endpoint FROM src_calls")}
        kept = {"today", "within"}
        self.assertEqual(jobs, kept)
        self.assertEqual(calls, kept)

    def test_empty_index_daily_keeps_previous_official_readable(self) -> None:
        pipe, db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        first = pipe.run_dataset("index_daily", TRADE_DATE)
        self.assertEqual(first["state"], "published")
        self.assertGreater(first["rows"], 0)

        pipe.adapter._transport = empty_transport_for("index_daily")
        with self.assertRaises(QualityError) as ctx:
            pipe.run_dataset("index_daily", TRADE_DATE)
        self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"])

        pub = db.fetchone(
            "SELECT * FROM publications WHERE dataset='index_daily' AND trade_date=?",
            (TRADE_DATE,),
        )
        self.assertEqual(pub["active_batch"], first["batch_id"])
        self.assertEqual(pub["state"], "published")
        empty_batch = db.fetchone(
            "SELECT * FROM batches WHERE batch_id=?",
            (f"{TRADE_DATE}-index_daily-002",),
        )
        self.assertIsNotNone(empty_batch)
        self.assertEqual(empty_batch["state"], "staged")
        self.assertEqual(empty_batch["rows_out"], 0)
        self.assertNotEqual(empty_batch["state"], "failed")
        self.assertIn(EMPTY_BATCH_ERROR, empty_batch["error"] or "")

        api = V1API(pipe.db, pipe, pipe.settings)
        payload = api.handle("/v1/indexes/bars", {"date": [TRADE_DATE], "code": ["000001.SH"]})
        self.assertEqual(payload["meta"]["batch_id"], first["batch_id"])
        self.assertEqual(payload["meta"]["tier"], "official")
        self.assertEqual(payload["meta"]["state"], "published")
        self.assertTrue(payload["data"])

    def test_empty_index_daily_first_batch_is_not_official(self) -> None:
        empty_date = "20260902"
        pipe, db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        pipe.adapter._transport = empty_transport_for("index_daily")
        with self.assertRaises(QualityError) as ctx:
            pipe.run_dataset("index_daily", empty_date)
        self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"])

        batch = db.fetchone("SELECT * FROM batches WHERE batch_id=?", (f"{empty_date}-index_daily-001",))
        self.assertEqual(batch["state"], "staged")
        self.assertEqual(batch["rows_out"], 0)
        self.assertIsNone(
            db.fetchone(
                "SELECT * FROM publications WHERE dataset='index_daily' AND trade_date=?",
                (empty_date,),
            )
        )
        api = V1API(pipe.db, pipe, pipe.settings)
        with self.assertRaises(ApiError) as api_ctx:
            api.handle("/v1/indexes/bars", {"date": [empty_date], "code": ["000001.SH"]})
        self.assertEqual(api_ctx.exception.code, "DATASET_NOT_PUBLISHED")

    def test_nonempty_index_daily_still_publishes(self) -> None:
        pipe, db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        result = pipe.run_dataset("index_daily", TRADE_DATE)
        self.assertEqual(result["state"], "published")
        self.assertEqual(result["rows"], 4)
        pub = db.fetchone(
            "SELECT * FROM publications WHERE dataset='index_daily' AND trade_date=?",
            (TRADE_DATE,),
        )
        self.assertEqual(pub["active_batch"], result["batch_id"])
        rows = db.fetchall("SELECT * FROM eod_index_bars WHERE batch_id=?", (result["batch_id"],))
        self.assertEqual(len(rows), 4)

    def test_empty_batch_guard_covers_all_official_datasets(self) -> None:
        for dataset in ("daily", "valuation", "moneyflow", "auction", "index_daily"):
            with self.subTest(dataset=dataset):
                pipe, db = make_pipeline()
                pipe.ingest_reference(TRADE_DATE)
                first = pipe.run_dataset(dataset, TRADE_DATE)
                pipe.adapter._transport = empty_transport_for(dataset)
                with self.assertRaises(QualityError) as ctx:
                    pipe.run_dataset(dataset, TRADE_DATE)
                self.assertTrue(ctx.exception.report["hard_fail"])
                self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"])
                pub = db.fetchone(
                    "SELECT * FROM publications WHERE dataset=? AND trade_date=?",
                    (dataset, TRADE_DATE),
                )
                self.assertEqual(pub["active_batch"], first["batch_id"])
                self.assertEqual(pub["state"], "published")
                empty_batch = db.fetchone(
                    "SELECT * FROM batches WHERE dataset=? AND trade_date=? AND batch_id != ?",
                    (dataset, TRADE_DATE, first["batch_id"]),
                )
                self.assertEqual(empty_batch["state"], "staged")
                self.assertEqual(empty_batch["rows_out"], 0)

    def test_publish_entry_rejects_empty_staging_without_moving_pointer(self) -> None:
        pipe, db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        first = pipe.run_dataset("moneyflow", TRADE_DATE)
        batch_id = pipe.next_batch_id("moneyflow", TRADE_DATE)
        pipe._set_batch(batch_id, "moneyflow", TRADE_DATE, "publishing", 1, rows_in=0, rows_out=0)
        pipe._stage("moneyflow", batch_id, [])
        with self.assertRaises(QualityError) as ctx:
            pipe.publish("moneyflow", TRADE_DATE, batch_id, state="published")
        self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"])
        pub = db.fetchone(
            "SELECT * FROM publications WHERE dataset='moneyflow' AND trade_date=?",
            (TRADE_DATE,),
        )
        self.assertEqual(pub["active_batch"], first["batch_id"])
        self.assertEqual(pub["state"], "published")
        history = db.fetchall(
            "SELECT batch_id FROM publication_history WHERE dataset='moneyflow' AND trade_date=?",
            (TRADE_DATE,),
        )
        self.assertEqual({row["batch_id"] for row in history}, {first["batch_id"]})


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