# ============================================================
# 轻量异步任务队列（moduleId: core-async-jobs, 可重生 ✅）
# 矩阵 88 行剩余：敏感性 / 蒙特卡洛重算放后台（提交 / 轮询 / 取消 / 结果）
# 线程模型：daemon worker 线程执行；fn 签名 fn(record)，可检查
#   record.cancel_requested 提前退出（或抛 JobCancelled）实现可取消长任务。
# 与网关解耦：本模块只做队列与状态机，业务闭包由调用方注入。
# ============================================================
from __future__ import annotations

import itertools
import threading
import time
import uuid
from collections.abc import Callable
from datetime import UTC, datetime
from typing import Any, Literal

from pydantic import BaseModel, Field

from server.timeutil import fmt_dt

JobStatus = Literal["pending", "running", "done", "failed", "cancelled"]
_TERMINAL = ("done", "failed", "cancelled")


class JobCancelled(Exception):
    """任务主动取消：worker 捕获后状态归为 cancelled。"""


def _now() -> str:
    # DTZ005：统一取本地时区的 aware 时间，墙钟与仓库其余模块一致
    return fmt_dt(datetime.now(UTC).astimezone())


class JobRecord(BaseModel):
    """任务记录：状态机 pending → running → done/failed/cancelled。"""

    job_id: str
    kind: str
    status: JobStatus = "pending"
    progress: int = 0
    params: dict[str, Any] = Field(default_factory=dict)
    actor: str = "planner"
    result: Any = None
    error: str | None = None
    created_at: str = ""
    started_at: str | None = None
    finished_at: str | None = None
    cancel_requested: bool = False


class JobQueue:
    """轻量异步任务队列：提交 / 轮询 / 取消 / 结果 / 统计。"""

    def __init__(self, *, max_finished: int = 200) -> None:
        self._lock = threading.RLock()
        self._jobs: dict[str, JobRecord] = {}
        self._seq = itertools.count(1)
        self.max_finished = max_finished

    # ---------------- 提交 ----------------
    def submit(
        self,
        kind: str,
        fn: Callable[[JobRecord], Any],
        *,
        params: dict[str, Any] | None = None,
        actor: str = "planner",
    ) -> str:
        """提交任务：立即起 daemon 线程执行，返回 job_id。"""
        job_id = f"job-{next(self._seq):04d}-{uuid.uuid4().hex[:6]}"
        record = JobRecord(
            job_id=job_id, kind=kind, params=params or {}, actor=actor, created_at=_now(),
        )
        with self._lock:
            self._jobs[job_id] = record
        thread = threading.Thread(
            target=self._worker, args=(job_id, fn),
            name=f"aps-job-{job_id}", daemon=True,
        )
        thread.start()
        return job_id

    def _worker(self, job_id: str, fn: Callable[[JobRecord], Any]) -> None:
        with self._lock:
            record = self._jobs[job_id]
            record.status = "running"
            record.started_at = _now()
            cancel = record.cancel_requested
        if cancel:
            self._finish(job_id, "cancelled")
            return
        try:
            result = fn(record)
            with self._lock:
                cancelled = self._jobs[job_id].cancel_requested
            if cancelled:
                self._finish(job_id, "cancelled")
            else:
                with self._lock:
                    self._jobs[job_id].result = result
                    self._jobs[job_id].progress = 100
                self._finish(job_id, "done")
        except JobCancelled:
            self._finish(job_id, "cancelled")
        except Exception as exc:  # noqa: BLE001 - worker 兜底任意业务异常，转 failed 状态
            with self._lock:
                self._jobs[job_id].error = str(exc)
            self._finish(job_id, "failed")
        finally:
            self._prune()

    def _finish(self, job_id: str, status: JobStatus) -> None:
        with self._lock:
            record = self._jobs[job_id]
            record.status = status
            record.finished_at = _now()
            if status == "cancelled":
                record.result = None
                record.error = "cancelled"

    # ---------------- 轮询 / 结果 ----------------
    def poll(self, job_id: str) -> dict[str, Any] | None:
        """查询任务快照（含 result/error）；不存在返回 None。"""
        with self._lock:
            record = self._jobs.get(job_id)
            return record.model_dump(mode="json") if record else None

    def wait(self, job_id: str, timeout: float = 15.0) -> dict[str, Any]:
        """阻塞轮询至终态（done/failed/cancelled）。"""
        deadline = time.time() + timeout
        while True:
            record = self.poll(job_id)
            if record is None:
                raise KeyError(job_id)
            if record["status"] in _TERMINAL:
                return record
            if time.time() >= deadline:
                raise TimeoutError(f"任务 {job_id} 等待超时（{timeout}s）")
            time.sleep(0.05)

    # ---------------- 取消 ----------------
    def cancel(self, job_id: str) -> bool:
        """取消任务：pending 立即取消；running 置取消标记（worker 收尾时落 cancelled）；
        已终态或不存在返回 False。"""
        with self._lock:
            record = self._jobs.get(job_id)
            if record is None or record.status in _TERMINAL:
                return False
            record.cancel_requested = True
            if record.status == "pending":
                record.status = "cancelled"
                record.finished_at = _now()
            return True

    # ---------------- 列表 / 统计 ----------------
    def list(self, status: str | None = None) -> list[dict[str, Any]]:
        with self._lock:
            items = [r.model_dump(mode="json") for r in self._jobs.values()]
        if status:
            items = [i for i in items if i["status"] == status]
        items.sort(key=lambda i: i["created_at"] or "", reverse=True)
        return items

    def stats(self) -> dict[str, int]:
        with self._lock:
            counts: dict[str, int] = {
                "total": len(self._jobs), "pending": 0, "running": 0,
                "done": 0, "failed": 0, "cancelled": 0,
            }
            for r in self._jobs.values():
                counts[r.status] = counts.get(r.status, 0) + 1
        return counts

    def _prune(self) -> None:
        """只保留最多 max_finished 条终态任务（按完成时间淘汰最旧）。"""
        with self._lock:
            finished = sorted(
                (r for r in self._jobs.values() if r.status in _TERMINAL),
                key=lambda r: r.finished_at or "",
            )
            overflow = len(finished) - self.max_finished
            for r in finished[:overflow]:
                self._jobs.pop(r.job_id, None)

    def clear(self) -> None:
        with self._lock:
            self._jobs.clear()


_queue: JobQueue | None = None


def get_job_queue() -> JobQueue:
    global _queue
    if _queue is None:
        _queue = JobQueue()
    return _queue


def reset_job_queue() -> None:
    """测试用：丢弃单例，下次重建。"""
    global _queue
    _queue = None