from __future__ import annotations

import threading
from collections.abc import Callable
from datetime import datetime, time, timedelta
from typing import Any

from datahub.db import HubDB
from datahub.logutil import get_logger
from datahub.pipeline import Pipeline
from datahub.timeutil import isoformat, now_shanghai, yyyymmdd

LOGGER = get_logger()

JobFn = Callable[[str], Any]

EOD_JOB_IDS = {"eod_a", "eod_b", "eod_retry"}


def is_open_day(db: HubDB, day: str) -> bool:
    row = db.fetchone(
        "SELECT is_open FROM trade_calendar WHERE exchange = 'SSE' AND cal_date = ?",
        (day,),
    )
    if row is None:
        return True  # unknown calendar: do not skip reference refresh
    return int(row["is_open"]) == 1


def _hhmm(value: str) -> time:
    return datetime.strptime(value, "%H:%M").time()


class Scheduler:
    """Calendar-driven in-process scheduler. Non-trading days skip EOD fetches.

    EOD datasets that failed to publish (e.g. upstream not ready at 15:05)
    are retried automatically every ``eod_retry_interval_minutes`` between
    ``eod_retry_start`` and ``eod_retry_cutoff``. Progress is persisted in
    ``eod_progress`` so a container restart catches up instead of waiting
    for the next day, and completed days are never re-fetched.
    """

    def __init__(self, db: HubDB, pipeline: Pipeline, jobs: dict[str, JobFn] | None = None) -> None:
        self.db = db
        self.pipeline = pipeline
        self.jobs = jobs or {
            "precheck": self._precheck,
            "eod_a": self._eod_a,
            "eod_b": self._eod_b,
            "eod_retry": self._eod_retry,
            "cleanup": self._cleanup,
            "backup": self._backup,
            "history_backfill": self._history_backfill,
        }
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._fired: set[tuple[str, str, str]] = set()
        self._eod_lock = threading.Lock()

    def start(self, interval_seconds: float = 30.0) -> None:
        if self._thread and self._thread.is_alive():
            return

        def loop() -> None:
            while not self._stop.wait(interval_seconds):
                try:
                    self.tick()
                except Exception:
                    LOGGER.exception("scheduler tick failed")

        self._thread = threading.Thread(target=loop, name="datahub-scheduler", daemon=True)
        self._thread.start()

    def stop(self, timeout: float = 5.0) -> None:
        self._stop.set()
        if self._thread and self._thread is not threading.current_thread():
            self._thread.join(timeout)

    def tick(self, clock: datetime | None = None) -> list[str]:
        now = now_shanghai(clock)
        day = yyyymmdd(now)
        current = now.time()
        ran: list[str] = []
        plan = [
            ("precheck", time(8, 45)),
            ("eod_a", time(15, 5)),
            ("eod_b", time(15, 10)),
            ("cleanup", time(0, 30)),
            ("backup", time(0, 40)),
        ]
        open_day = is_open_day(self.db, day)
        for job_id, at in plan:
            if current < at:
                continue
            key = (job_id, day, at.strftime("%H%M"))
            if key in self._fired:
                continue
            if job_id in {"eod_a", "eod_b"} and not open_day:
                self._fired.add(key)
                continue
            self._fired.add(key)
            if job_id in {"eod_a", "eod_b"}:
                # Record the attempt before running: even a crash must not
                # hide that today's first EOD try already happened.
                self._record_eod_attempt(day, now)
            try:
                self.run_job(job_id, day)
            except Exception:
                if job_id not in {"eod_a", "eod_b"}:
                    raise
                # Keep the tick alive; evening retries take over.
                LOGGER.exception("scheduled job %s failed for %s", job_id, day)
            ran.append(job_id)
            if job_id in {"eod_a", "eod_b"}:
                self._settle_eod(day)
        ran.extend(self._eod_retry_tick(now, day, open_day))
        return ran

    # ------------------------------------------------------------------
    # EOD retry window
    # ------------------------------------------------------------------

    def _eod_retry_tick(self, now: datetime, day: str, open_day: bool) -> list[str]:
        if not open_day:
            return []
        settings = self.pipeline.settings
        current = now.time()
        start = _hhmm(settings.eod_retry_start)
        cutoff = _hhmm(settings.eod_retry_cutoff)
        interval = timedelta(minutes=settings.eod_retry_interval_minutes)
        missing = self.pipeline.missing_official_datasets(day)
        row = self.eod_progress(day)

        if not missing:
            if row is None or row["state"] != "done":
                self._save_eod_progress(day, state="done", finished_at=isoformat(now))
            return []
        if current < start:
            return []
        if row and row["state"] == "cutoff_failed":
            return []
        if current >= cutoff:
            detail = "截止时间已到，缺失数据集: " + ",".join(missing)
            self._save_eod_progress(day, state="cutoff_failed", finished_at=isoformat(now), detail=detail)
            with self.db.write() as connection:
                connection.execute(
                    "INSERT INTO job_runs(job_id, state, started_at, finished_at, error, attempt, detail)"
                    " VALUES ('eod_retry','failed',?,?,?,?,?)",
                    (isoformat(now), isoformat(now), detail, int((row or {}).get("attempts") or 0), "eod cutoff reached"),
                )
            LOGGER.warning(
                "eod retry window closed without data",
                extra={"hub": {"trade_date": day, "missing": missing, "reason": "eod_cutoff"}},
            )
            return []
        last = None
        if row and row["last_attempt_at"]:
            try:
                last = datetime.fromisoformat(str(row["last_attempt_at"]))
            except ValueError:
                last = None
        if last is not None and now_shanghai(last).replace(tzinfo=None) + interval > now.replace(tzinfo=None):
            return []
        if "eod_retry" not in self.jobs:
            return []
        self._record_eod_attempt(day, now)
        ran = []
        try:
            self.run_job("eod_retry", day)
        except Exception:
            # job_runs already carries the failure; the window keeps retrying.
            LOGGER.warning("eod retry failed for %s", day, exc_info=True)
        ran.append("eod_retry")
        self._settle_eod(day)
        return ran

    def _settle_eod(self, day: str) -> None:
        """Flip the day to done as soon as every official dataset is published."""
        if not self.pipeline.missing_official_datasets(day):
            row = self.eod_progress(day)
            if row is None or row["state"] != "done":
                self._save_eod_progress(day, state="done", finished_at=isoformat())

    def eod_progress(self, day: str) -> dict[str, Any] | None:
        return self.db.fetchone("SELECT * FROM eod_progress WHERE trade_date = ?", (day,))

    def eod_status(self, trade_date: str | None = None, clock: datetime | None = None) -> dict[str, Any]:
        """Human/admin facing view: 等待上游 / 下次重试 / 已成功 / 已截止失败."""
        day = yyyymmdd(trade_date or now_shanghai(clock))
        now = now_shanghai(clock)
        row = self.eod_progress(day)
        open_day = is_open_day(self.db, day)
        missing = self.pipeline.missing_official_datasets(day)
        if row and row["state"] == "done":
            state = "done"
        elif not open_day:
            state = "closed_day"
        elif not missing:
            state = "done"
        elif row and row["state"] == "cutoff_failed":
            state = "cutoff_failed"
        elif now.time() < _hhmm("15:05"):
            state = "pending_first_attempt"
        else:
            state = "waiting_upstream"
        return {
            "trade_date": day,
            "is_open_day": open_day,
            "state": state,
            "missing_datasets": missing,
            "attempts": int((row or {}).get("attempts") or 0),
            "last_attempt_at": (row or {}).get("last_attempt_at"),
            "next_retry_at": (row or {}).get("next_retry_at") if state == "waiting_upstream" else None,
            "finished_at": (row or {}).get("finished_at"),
            "detail": (row or {}).get("detail"),
        }

    def _record_eod_attempt(self, day: str, now: datetime) -> None:
        row = self.eod_progress(day)
        attempts = int((row or {}).get("attempts") or 0) + 1
        interval = self.pipeline.settings.eod_retry_interval_minutes
        self._save_eod_progress(
            day,
            state="waiting_upstream",
            attempts=attempts,
            last_attempt_at=isoformat(now),
            next_retry_at=isoformat(now + timedelta(minutes=interval)),
        )

    def _save_eod_progress(self, day: str, **fields: Any) -> None:
        columns = [
            "trade_date", "state", "attempts", "last_attempt_at",
            "next_retry_at", "finished_at", "detail", "updated_at",
        ]
        with self.db.write() as connection:
            existing = connection.execute(
                "SELECT trade_date FROM eod_progress WHERE trade_date = ?",
                (day,),
            ).fetchone()
            if existing is None:
                payload = {name: None for name in columns}
                payload.update({"trade_date": day, "state": "waiting_upstream", "attempts": 0})
                payload.update(fields)
                payload["updated_at"] = isoformat()
                placeholders = ",".join("?" for _ in columns)
                connection.execute(
                    f"INSERT INTO eod_progress({','.join(columns)}) VALUES ({placeholders})",
                    tuple(payload[name] for name in columns),
                )
            else:
                assignments = ", ".join(f"{name} = ?" for name in fields)
                connection.execute(
                    f"UPDATE eod_progress SET {assignments}, updated_at = ? WHERE trade_date = ?",
                    (*fields.values(), isoformat(), day),
                )

    # ------------------------------------------------------------------
    # Job execution
    # ------------------------------------------------------------------

    def run_job(self, job_id: str, trade_date: str) -> dict[str, Any]:
        fn = self.jobs.get(job_id)
        if fn is None:
            raise KeyError(job_id)
        if job_id in EOD_JOB_IDS:
            if not self._eod_lock.acquire(blocking=False):
                return {
                    "job_id": job_id,
                    "state": "skipped",
                    "detail": "another EOD job is already running",
                }
            try:
                return self._run_job(fn, job_id, trade_date)
            finally:
                self._eod_lock.release()
        return self._run_job(fn, job_id, trade_date)

    def _run_job(self, fn: JobFn, job_id: str, trade_date: str) -> dict[str, Any]:
        started = isoformat()
        run_id = None
        with self.db.write() as connection:
            cur = connection.execute(
                "INSERT INTO job_runs(job_id, state, started_at, attempt) VALUES (?,?,?,1)",
                (job_id, "running", started),
            )
            run_id = cur.lastrowid
        try:
            result = fn(trade_date) or {}
            if isinstance(result, dict):
                failures = self.pipeline.eod_failures(result) if job_id in EOD_JOB_IDS else []
                if failures:
                    raise RuntimeError("; ".join(failures))
            with self.db.write() as connection:
                connection.execute(
                    "UPDATE job_runs SET state=?, finished_at=?, rows_out=?, detail=? WHERE id=?",
                    ("ok", isoformat(), result.get("rows") if isinstance(result, dict) else None, str(result)[:2000], run_id),
                )
            return {"job_id": job_id, "result": result, "state": "ok"}
        except Exception as exc:
            with self.db.write() as connection:
                connection.execute(
                    "UPDATE job_runs SET state=?, finished_at=?, error=? WHERE id=?",
                    ("failed", isoformat(), str(exc), run_id),
                )
            raise

    def _precheck(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.ingest_reference(trade_date)

    def _eod_a(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.run_eod_batch_a(trade_date)

    def _eod_b(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.run_eod_batch_b(trade_date)

    def _eod_retry(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.run_eod_missing(trade_date)

    def _history_backfill(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.backfill_history(trade_date)

    def _cleanup(self, trade_date: str) -> dict[str, Any]:
        result = self.pipeline.cleanup()
        if now_shanghai().weekday() == 6:
            self.pipeline.db.vacuum()
            result["vacuum"] = True
        return result

    def _backup(self, trade_date: str) -> dict[str, Any]:
        from pathlib import Path

        dest_dir = Path(self.pipeline.settings.backup_dir)
        dest = dest_dir / f"datahub-{trade_date}.db"
        self.pipeline.db.backup_to(dest)
        keep = int(self.pipeline.settings.quality.get("backup_retain") or 14)
        backups = sorted(dest_dir.glob("datahub-*.db"))
        for old in backups[:-keep]:
            old.unlink(missing_ok=True)
        return {"path": str(dest.name), "kept": min(len(backups), keep)}