"""异步后台任务注册表（T1：后端异步化核心）。

设计原则（对齐 v3 方案 D1-D6）：
- asyncio 原生，无 Celery/Redis/外部队列
- 每个 request_id -> TaskInfo（状态/开始时间/结果或错误）
- 任务写 runs/{id}/workspace.json，SSE 生成器只读该文件（D3：不改 pipeline）
- 定期清理已完成任务（防内存泄漏）
"""
from __future__ import annotations

import asyncio
import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, AsyncGenerator, Dict, Optional

# runs/ 目录（与 pipeline.py 默认一致）
RUNS_DIR = Path("runs")

# 任务状态
STATE_PENDING = "pending"
STATE_RUNNING = "running"
STATE_DONE = "done"
STATE_FAILED = "failed"


@dataclass
class TaskInfo:
    """一个后台任务的状态快照。"""
    request_id: str
    state: str = STATE_PENDING           # pending | running | done | failed
    started_at: float = 0.0             # time.time()
    finished_at: float = 0.0            # time.time()（done/failed 时）
    error: Optional[str] = None          # failed 时错误信息
    # PipelineResult 字段（done 时填充）
    response: Optional[str] = None
    status: Optional[str] = None          # done | fast_path | escalated | failed
    fast_path: bool = False
    rounds_used: int = 0
    api_input_tokens: int = 0
    api_output_tokens: int = 0
    cost_est: float = 0.0
    model_used: Optional[str] = None
    latency_ms: float = 0.0
    workspace_path: Optional[str] = None
    error_detail: Optional[str] = None
    route: list = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "request_id": self.request_id,
            "state": self.state,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "error": self.error,
            "response": self.response,
            "status": self.status,
            "fast_path": self.fast_path,
            "rounds_used": self.rounds_used,
            "api_input_tokens": self.api_input_tokens,
            "api_output_tokens": self.api_output_tokens,
            "cost_est": self.cost_est,
            "model_used": self.model_used,
            "latency_ms": round(self.latency_ms, 2),
            "workspace_path": self.workspace_path,
            "route": self.route,
            "error_detail": self.error_detail,
        }

    def to_result_event(self) -> Dict[str, Any]:
        """终态 SSE result 事件 payload。"""
        return {
            "type": "result",
            "request_id": self.request_id,
            "response": self.response or "",
            "status": self.status,
            "fast_path": self.fast_path,
            "rounds_used": self.rounds_used,
            "api_input_tokens": self.api_input_tokens,
            "api_output_tokens": self.api_output_tokens,
            "cost_est": self.cost_est,
            "model_used": self.model_used,
            "latency_ms": round(self.latency_ms, 2),
            "route": self.route,
            "error": self.error,
        }


class JobStore:
    """asyncio 后台任务注册表。

    线程安全（asyncio 事件循环单线程，不需要额外锁）。
    任务由 register() 注册、由 _task_done() 填充结果。
    SSE 生成器调用 watch_file() 轮询 workspace.json 变化。
    """

    def __init__(self, max_age_seconds: float = 3600.0, max_running: int = 50):
        """
        Args:
            max_age_seconds: 已完成任务在内存中保留时间（秒），超时后自动清理
            max_running: 最大同时运行任务数，超出后拒绝新任务
        """
        self._tasks: Dict[str, TaskInfo] = {}
        self._asyncio_tasks: Dict[str, "asyncio.Task[None]"] = {}  # request_id -> asyncio.Task
        self._cancel_events: Dict[str, asyncio.Event] = {}        # request_id -> cancel Event
        self.max_age_seconds = max_age_seconds
        self.max_running = max_running
        self._poll_interval = 0.3  # workspace.json 轮询间隔（秒）

    # ─────────────────────────────────────────────────────────────────
    # 公共 API
    # ─────────────────────────────────────────────────────────────────

    def register(self, request_id: str) -> tuple[bool, str]:
        """注册一个 pending 任务。返回 (True, "") 成功，(False, reason) 容量满。"""
        running = [t for t in self._tasks.values() if t.state == STATE_RUNNING]
        if len(running) >= self.max_running:
            return False, f"同时运行任务已达上限 {self.max_running}，请稍后重试"
        if request_id in self._tasks:
            return False, f"任务 {request_id} 已存在"
        self._tasks[request_id] = TaskInfo(
            request_id=request_id,
            state=STATE_PENDING,
            started_at=time.time(),
        )
        return True, ""

    def get(self, request_id: str) -> Optional[TaskInfo]:
        """查询任务状态。"""
        return self._tasks.get(request_id)

    def list_all(self) -> Dict[str, TaskInfo]:
        """列出所有任务（含已完成）。"""
        return dict(self._tasks)

    def list_running(self) -> Dict[str, TaskInfo]:
        """只列出运行中任务。"""
        return {k: v for k, v in self._tasks.items() if v.state == STATE_RUNNING}

    def submit(
        self,
        request_id: str,
        coro,  # type: asyncio.coroutine
    ) -> "asyncio.Task[None]":
        """提交协程到后台运行；内部注册 asyncio.Task。"""
        task = asyncio.create_task(self._run_wrapper(request_id, coro))
        self._asyncio_tasks[request_id] = task
        return task

    def new_cancel(self, request_id: str) -> asyncio.Event:
        """为 SSE 连接创建一个取消事件。"""
        evt = asyncio.Event()
        self._cancel_events[request_id] = evt
        return evt

    def get_cancel(self, request_id: str) -> Optional[asyncio.Event]:
        """获取已有取消事件。"""
        return self._cancel_events.get(request_id)

    # ─────────────────────────────────────────────────────────────────
    # 内部：任务运行包装
    # ─────────────────────────────────────────────────────────────────

    async def _run_wrapper(self, request_id: str, coro):
        """把用户协程包装成可追踪的任务：更新状态、捕获异常、清理。"""
        info = self._tasks.get(request_id)
        if info is None:
            return
        info.state = STATE_RUNNING
        try:
            await coro
        except Exception as exc:  # pragma: no cover
            if info:
                info.state = STATE_FAILED
                info.finished_at = time.time()
                info.error = f"unhandled:{type(exc).__name__}:{exc}"
        finally:
            # 清理 asyncio task 引用
            self._asyncio_tasks.pop(request_id, None)
            # 定期 GC 已完成任务
            self._cleanup_aged()

    def _task_done(
        self,
        request_id: str,
        result: "PipelineResult",  # from pipeline.PipelineResult
    ) -> None:
        """任务正常完成时由调用方调用，写入结果。"""
        info = self._tasks.get(request_id)
        if info is None:
            return
        info.state = STATE_DONE
        info.finished_at = time.time()
        info.response = result.response
        info.status = result.status
        info.fast_path = result.fast_path
        info.rounds_used = result.rounds_used
        info.api_input_tokens = result.api_input_tokens
        info.api_output_tokens = result.api_output_tokens
        info.cost_est = result.cost_est
        info.model_used = result.model_used
        info.latency_ms = result.latency_ms
        info.workspace_path = result.workspace_path
        info.route = list(result.route) if result.route else []
        if result.error:
            info.error_detail = result.error

    def _task_failed(
        self,
        request_id: str,
        error: str,
        detail: Optional[str] = None,
    ) -> None:
        """任务异常结束时调用。"""
        info = self._tasks.get(request_id)
        if info is None:
            return
        info.state = STATE_FAILED
        info.finished_at = time.time()
        info.error = error
        info.error_detail = detail

    # ─────────────────────────────────────────────────────────────────
    # SSE 专用：workspace.json 文件轮询
    # ─────────────────────────────────────────────────────────────────

    async def watch_workspace(
        self,
        request_id: str,
        cancel_event: asyncio.Event,
    ) -> AsyncGenerator[Dict[str, Any], None]:
        """监视 runs/{id}/workspace.json，yield 单条事件。

        事件类型：
        - {"type": "status", "value": <state>, "request_id": <id>}
        - {"type": "workspace", "version": <n>, "workspace": <dict>}
        - {"type": "error", "detail": <str>}

        当任务进入 done/failed 状态或 cancel_event 被 set 时停止。
        """
        ws_path = RUNS_DIR / request_id / "workspace.json"
        seen_mtime: float = 0.0
        seen_size: int = 0

        while not cancel_event.is_set():
            info = self.get(request_id)
            # 检查终态
            if info and info.state in (STATE_DONE, STATE_FAILED):
                if info.state == STATE_FAILED:
                    yield {"type": "error", "detail": info.error or "任务失败"}
                # result 由 /runs/{id}/status 提供，此处只推送终态 status
                yield {"type": "status", "value": info.state, "request_id": request_id}
                break

            # 读文件变化（mtime + size 双检）
            if ws_path.exists():
                try:
                    stat = ws_path.stat()
                    if stat.st_mtime != seen_mtime or stat.st_size != seen_size:
                        raw = ws_path.read_text(encoding="utf-8")
                        data = json.loads(raw)
                        version = int(data.get("meta", {}).get("round", 0))
                        seen_mtime = stat.st_mtime
                        seen_size = stat.st_size
                        yield {
                            "type": "workspace",
                            "version": version,
                            "state": info.state if info else STATE_RUNNING,
                            "request_id": request_id,
                            "workspace": data,
                        }
                except (json.JSONDecodeError, OSError):
                    pass  # 文件正在写入，忽略

            await asyncio.sleep(self._poll_interval)

        # 连接关闭前最后推一次终态
        info = self.get(request_id)
        if info:
            yield {"type": "status", "value": info.state, "request_id": request_id}

    # ─────────────────────────────────────────────────────────────────
    # 内存清理
    # ─────────────────────────────────────────────────────────────────

    def _cleanup_aged(self) -> None:
        """删除超过 max_age_seconds 的已完成任务引用。"""
        now = time.time()
        to_remove = [
            rid for rid, info in self._tasks.items()
            if info.state in (STATE_DONE, STATE_FAILED)
            and (now - info.finished_at) > self.max_age_seconds
        ]
        for rid in to_remove:
            self._tasks.pop(rid, None)

    def cancel(self, request_id: str) -> bool:
        """取消运行中的任务。返回 True 找到并取消，False 未找到。"""
        task = self._asyncio_tasks.get(request_id)
        if task is None:
            return False
        task.cancel()
        info = self.get(request_id)
        if info:
            info.state = STATE_FAILED
            info.finished_at = time.time()
            info.error = "cancelled_by_user"
        return True


# ─────────────────────────────────────────────────────────────────────────────
# 全局单例（gateway 进程内共享）
# ─────────────────────────────────────────────────────────────────────────────
_store: Optional[JobStore] = None


def get_job_store() -> JobStore:
    global _store
    if _store is None:
        _store = JobStore(max_age_seconds=3600.0, max_running=50)
    return _store


def reset_job_store() -> None:
    """测试用：重置全局注册表。"""
    global _store
    _store = None