from __future__ import annotations

import json
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.revision import revision_datasets
from datahub.timeutil import isoformat, now_shanghai, yyyymmdd

LOGGER = get_logger()

JobFn = Callable[[str], Any]

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


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 previous_open_day(db: HubDB, day: str) -> str | None:
    row = db.fetchone(
        """
        SELECT cal_date FROM trade_calendar
        WHERE exchange = 'SSE' AND is_open = 1 AND cal_date < ?
        ORDER BY cal_date DESC LIMIT 1
        """,
        (day,),
    )
    if row is None:
        return None
    return str(row["cal_date"])


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,
            "eod_revise": self._eod_revise,
            "stocks_refresh": self._stocks_refresh,
            "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)),
        ]
        for refresh_at in self.pipeline.settings.stocks_refresh_times:
            plan.append(("stocks_refresh", _hhmm(refresh_at)))
        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", "stocks_refresh"} 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", "stocks_refresh"}:
                    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))
        ran.extend(self._revision_review_tick(now, day, open_day))
        ran.extend(self._revision_catchup_tick(now, 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 revision_progress(self, day: str) -> dict[str, Any] | None:
        return self.db.fetchone("SELECT * FROM revision_progress WHERE trade_date = ?", (day,))

    def revision_status(self, trade_date: str | None = None, clock: datetime | None = None) -> dict[str, Any]:
        """等待复核 / 复核失败 / 已追平 / 已截止."""
        day = yyyymmdd(trade_date or now_shanghai(clock))
        row = self.revision_progress(day)
        open_day = is_open_day(self.db, day)
        published = self._revision_ready(day)
        if row and row["state"] in {"aligned", "review_failed", "cutoff", "waiting_review"}:
            state = str(row["state"])
        elif not open_day:
            state = "closed_day"
        elif not published:
            state = "pending_publish"
        else:
            state = "waiting_review"
        return {
            "trade_date": day,
            "is_open_day": open_day,
            "state": state,
            "datasets": list(revision_datasets(self.pipeline.settings.quality)),
            "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 in {"waiting_review", "review_failed"} else None,
            "finished_at": (row or {}).get("finished_at"),
            "catchup_done": bool(int((row or {}).get("catchup_done") or 0)),
            "detail": (row or {}).get("detail"),
            "window": f"{self.pipeline.settings.revision_review_start}-{self.pipeline.settings.revision_review_cutoff}",
        }

    def _revision_ready(self, day: str) -> bool:
        return all(
            self.pipeline.active_batch(dataset, day)
            for dataset in revision_datasets(self.pipeline.settings.quality)
        )

    def _revision_due(self, now: datetime, row: dict[str, Any] | None) -> bool:
        if row is None or not row.get("last_attempt_at"):
            return True
        try:
            last = datetime.fromisoformat(str(row["last_attempt_at"]))
        except ValueError:
            return True
        interval = timedelta(minutes=self.pipeline.settings.revision_review_interval_minutes)
        return now_shanghai(last).replace(tzinfo=None) + interval <= now.replace(tzinfo=None)

    def _revision_review_tick(self, now: datetime, day: str, open_day: bool) -> list[str]:
        if not open_day or not self._revision_ready(day):
            return []
        settings = self.pipeline.settings
        current = now.time()
        start = _hhmm(settings.revision_review_start)
        cutoff = _hhmm(settings.revision_review_cutoff)
        row = self.revision_progress(day)
        if current < start:
            if row is None:
                self._save_revision_progress(day, state="waiting_review")
            return []
        if current >= cutoff:
            if row is None or row["state"] not in {"aligned", "cutoff"}:
                detail = "复核窗口已截止"
                self._save_revision_progress(
                    day, state="cutoff", 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_revise','failed',?,?,?,?,?)",
                        (
                            isoformat(now), isoformat(now), detail,
                            int((row or {}).get("attempts") or 0), "revision cutoff reached",
                        ),
                    )
            elif row["state"] == "aligned" and not row.get("finished_at"):
                self._save_revision_progress(day, finished_at=isoformat(now))
            return []
        if not self._revision_due(now, row):
            return []
        if "eod_revise" not in self.jobs:
            return []
        return self._run_revision_job(day, now, catchup=False)

    def _revision_catchup_tick(self, now: datetime, day: str) -> list[str]:
        prev = previous_open_day(self.db, day)
        if prev is None or prev >= day:
            return []
        if not self._revision_ready(prev):
            return []
        row = self.revision_progress(prev)
        if row and int(row.get("catchup_done") or 0):
            return []
        if not self._revision_due(now, row):
            return []
        if "eod_revise" not in self.jobs:
            return []
        return self._run_revision_job(prev, now, catchup=True)

    def _run_revision_job(self, day: str, now: datetime, catchup: bool) -> list[str]:
        attempts = int((self.revision_progress(day) or {}).get("attempts") or 0) + 1
        interval = self.pipeline.settings.revision_review_interval_minutes
        self._save_revision_progress(
            day,
            state="waiting_review",
            attempts=attempts,
            last_attempt_at=isoformat(now),
            next_retry_at=isoformat(now + timedelta(minutes=interval)),
        )
        ran: list[str] = []
        try:
            out = self.run_job("eod_revise", day)
        except Exception as exc:
            LOGGER.warning("revision review failed for %s: %s", day, exc)
            self._save_revision_progress(
                day,
                state="review_failed",
                detail="复核失败，保留上一完整版本",
            )
            ran.append("eod_revise")
            return ran
        ran.append("eod_revise")
        if out.get("state") == "skipped":
            return ran
        result = out.get("result") if isinstance(out.get("result"), dict) else {}
        failed = [
            name for name, item in result.items()
            if isinstance(item, dict) and item.get("state") == "failed"
        ]
        review = result.get("review") if isinstance(result.get("review"), dict) else None
        watched = [
            result[name]
            for name in revision_datasets(self.pipeline.settings.quality)
            if isinstance(result.get(name), dict)
        ]
        diff_blob = None
        if review and review.get("diffs"):
            diff_blob = json.dumps(review.get("diffs"), ensure_ascii=False)
        else:
            for item in watched:
                if item.get("diffs"):
                    diff_blob = json.dumps(item.get("diffs"), ensure_ascii=False)
                    break
        revised = bool(review and review.get("reason") == "revised")
        matched = any(item.get("reason") == "unchanged" or item.get("state") == "aligned" for item in watched)
        if failed:
            self._save_revision_progress(
                day,
                state="review_failed",
                detail="复核失败，保留上一完整版本",
                last_diff=diff_blob,
            )
        elif revised or matched:
            fields: dict[str, Any] = {
                "state": "aligned",
                "finished_at": isoformat(now),
                "detail": "已追平" if revised else "已追平（无变化）",
                "last_diff": diff_blob,
            }
            if catchup:
                fields["catchup_done"] = 1
            self._save_revision_progress(day, **fields)
        return ran

    def _save_revision_progress(self, day: str, **fields: Any) -> None:
        columns = [
            "trade_date", "state", "attempts", "last_attempt_at",
            "next_retry_at", "finished_at", "catchup_done", "last_diff", "detail", "updated_at",
        ]
        with self.db.write() as connection:
            existing = connection.execute(
                "SELECT trade_date FROM revision_progress WHERE trade_date = ?",
                (day,),
            ).fetchone()
            if existing is None:
                payload = {name: None for name in columns}
                payload.update({
                    "trade_date": day,
                    "state": "waiting_review",
                    "attempts": 0,
                    "catchup_done": 0,
                })
                payload.update(fields)
                payload["updated_at"] = isoformat()
                placeholders = ",".join("?" for _ in columns)
                connection.execute(
                    f"INSERT INTO revision_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 revision_progress SET {assignments}, updated_at = ? WHERE trade_date = ?",
                    (*fields.values(), isoformat(), day),
                )

    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 _eod_revise(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.review_published_revisions(trade_date)

    def _stocks_refresh(self, trade_date: str) -> dict[str, Any]:
        return self.pipeline.refresh_stocks(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)}
