from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from datahub.adapters.tushare import TushareAdapter
from datahub.crypto import SecretVault
from datahub.db import HubDB
from datahub.pipeline import Pipeline, QualityError
from datahub.settings import Settings
from tests.fixtures import TRADE_DATE, fake_transport


def make_pipeline(before_commit=None) -> tuple[Pipeline, HubDB]:
    tmp = tempfile.TemporaryDirectory()
    db = HubDB(Path(tmp.name) / "hub.db")
    adapter = TushareAdapter("test-token", transport=fake_transport)
    settings = Settings(
        encryption_key=SecretVault.generate_key(),
        api_token="t" * 32,
        admin_password="admin-pass",
        tushare_token="test-token",
        db_path=db.path,
        quality={"daily_row_ratio": 0.98, "null_rate_max": 0.01, "max_publish_attempts": 3, "publication_generations": 3},
        scheduler_enabled=False,
    )
    pipe = Pipeline(db, adapter, settings, before_commit=before_commit)
    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)
        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")


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