"""圆桌后台作业 dispatcher（Phase 3）：跨 worker 领取 + 崩溃接管 + outbox 补偿。

把「谁来跑这个作业」从「发起请求的那个 worker 的进程内存」收敛到 **DB 租约**：

- **领取（claim）**：后台循环定期 ``list_claimable``（``queued`` 新作业/续跑，或 ``running``
  但租约过期的僵尸），对每个候选用 ``claim_job`` 的**原子条件 UPDATE** 抢占 —— 多 worker
  并发扫到同一候选时只有一个 ``rowcount==1``，其余拿到 None 自动跳过。一致性来自 DB，
  不是进程内锁。
- **心跳**：领取成功后交给 ``RoundtableJobExecutor`` 跑，执行器内部定期 ``renew_lease``；
  worker 崩溃则心跳停止，租约最长一个 TTL 后过期，下一轮循环在其它 worker 重新领取
  （``attempt+1``）—— 崩溃可接管。
- **启动扫描**：本 worker 启动时立即跑一轮 ``dispatch_once``，尽快接走历史遗留的
  queued / 租约过期作业（不用等第一个轮询周期）。
- **outbox 补偿器**：``draft_persist_pending=1`` 的作业（后台研讨结果因乐观锁冲突没能写回
  草稿）在这里被重新投影到草稿 step2.runs，成功后清标记 —— 「研讨结果不丢」由 DB 标记驱动。

所有 worker 都跑同一个 dispatcher；谁先 ``claim_job`` 成功谁执行，天然幂等、互斥。
"""

from __future__ import annotations

import asyncio
import logging
from contextlib import suppress
from datetime import UTC, datetime, timedelta
from typing import Any

from app.gateway.roundtable_job_executor import (
    _LEASE_TTL_SECONDS,
    WORKER_ID,
    params_from_snapshot,
)

logger = logging.getLogger(__name__)

# dispatcher 轮询间隔（秒）。路由 start/resume 成功后会 ``nudge()`` 立即唤醒，故正常路径
# 几乎无感延迟；该间隔只是「没有 nudge 时的兜底」，兼作崩溃接管的发现周期。
_DISPATCH_POLL_INTERVAL = 3.0
# outbox 补偿器轮询间隔（秒）。pending 标记是低频异常路径，慢一点无妨。
_COMPENSATE_INTERVAL = 10.0
# 单轮最多领取的作业数（防止一次扫太多把事件循环占死）。
_CLAIM_BATCH = 20


class RoundtableJobDispatcher:
    """跨 worker 的圆桌后台作业调度器（每个 worker 进程一个实例，lifespan 起停）。"""

    def __init__(
        self,
        store,
        executor,
        *,
        draft_store=None,
        task_draft_store=None,
        worker_id: str | None = None,
    ) -> None:
        self._store = store
        self._executor = executor
        # outbox 补偿器投影草稿用（个人草稿 / task 草稿两个存储）。
        self._draft_store = draft_store
        self._task_draft_store = task_draft_store
        self._worker_id = worker_id or WORKER_ID
        self._wake = asyncio.Event()
        self._task: asyncio.Task | None = None
        self._stopping = False
        self._last_compensate = datetime.now(UTC)

    # ── 生命周期 ─────────────────────────────────────────────────────────────

    def start(self) -> None:
        """启动后台循环（lifespan startup 调用）。幂等。"""
        if self._task is not None and not self._task.done():
            return
        self._stopping = False
        self._task = asyncio.create_task(self._loop(), name="roundtable-job-dispatcher")
        logger.info("[roundtable-jobs] dispatcher started (worker=%s)", self._worker_id)

    async def stop(self) -> None:
        """停止后台循环（lifespan shutdown 调用）。"""
        self._stopping = True
        self._wake.set()
        if self._task is not None:
            self._task.cancel()
            with suppress(asyncio.CancelledError, Exception):
                await self._task
            self._task = None

    def nudge(self) -> None:
        """立即唤醒一轮领取（路由 start/resume 成功后调用，降低启动延迟）。"""
        self._wake.set()

    # ── 主循环 ───────────────────────────────────────────────────────────────

    async def _loop(self) -> None:
        # 启动即扫一轮：尽快接管历史遗留 queued / 租约过期（僵尸）作业。
        await self._safe_dispatch()
        while not self._stopping:
            try:
                await asyncio.wait_for(self._wake.wait(), timeout=_DISPATCH_POLL_INTERVAL)
            except TimeoutError:
                pass
            self._wake.clear()
            if self._stopping:
                break
            await self._safe_dispatch()
            # outbox 补偿器按自己的节奏跑（低频）。
            if (datetime.now(UTC) - self._last_compensate).total_seconds() >= _COMPENSATE_INTERVAL:
                self._last_compensate = datetime.now(UTC)
                await self._safe_compensate()

    async def _safe_dispatch(self) -> None:
        try:
            await self.dispatch_once()
        except Exception:  # noqa: BLE001 — 单轮失败不能杀死循环，下轮重试
            logger.exception("[roundtable-jobs] dispatch round failed")

    async def _safe_compensate(self) -> None:
        try:
            await self.compensate_once()
        except Exception:  # noqa: BLE001
            logger.exception("[roundtable-jobs] compensate round failed")

    # ── 领取 ─────────────────────────────────────────────────────────────────

    async def dispatch_once(self) -> int:
        """扫一轮可领取作业并逐个条件抢占。返回本轮成功启动的作业数。"""
        now = datetime.now(UTC)
        await self._reap_cancel_requested(now)
        try:
            candidate_ids = await self._store.list_claimable(now=now, limit=_CLAIM_BATCH)
        except Exception:  # noqa: BLE001
            logger.warning("[roundtable-jobs] list_claimable failed", exc_info=True)
            return 0
        if not candidate_ids:
            return 0

        started = 0
        lease_until = now + timedelta(seconds=_LEASE_TTL_SECONDS)
        for job_id in candidate_ids:
            try:
                row = await self._store.claim_job(
                    job_id, lease_owner=self._worker_id, lease_until=lease_until
                )
            except Exception:  # noqa: BLE001
                logger.warning("[roundtable-jobs] claim_job failed for %s", job_id, exc_info=True)
                continue
            if row is None:
                # 别的 worker 先领了 / 状态已变 —— 条件 UPDATE rowcount=0，跳过即可。
                continue
            params = params_from_snapshot(
                row.get("input_snapshot"), job_id=job_id, lease_owner=self._worker_id
            )
            if params is None:
                # 快照缺失/损坏：无法重建执行参数 —— 标记 error 并释放活跃键，避免僵死占位。
                logger.error(
                    "[roundtable-jobs] job %s claimed but snapshot missing/corrupt; marking error",
                    job_id,
                )
                with suppress(Exception):
                    await self._store.update_progress(
                        job_id,
                        lease_owner=self._worker_id,
                        status="error",
                        error="input_snapshot 缺失或损坏，无法在后台重建研讨",
                    )
                continue
            self._executor.start_job(params)
            started += 1
            logger.info(
                "[roundtable-jobs] worker=%s claimed job %s (attempt=%s, resume=%s)",
                self._worker_id, job_id, row.get("attempt"), params.is_resume,
            )
        return started

    async def _reap_cancel_requested(self, now: datetime) -> int:
        """收口持租 worker 崩溃后遗留的 ``cancel_requested`` 作业。

        有效租约仍在时，必须让持租 worker 的取消看门狗收口；无租约或租约过期时，已经
        不存在可信的持租者，此处通过 ``finalize_cancel`` 的条件 UPDATE 安全地回收。
        """
        try:
            candidate_ids = await self._store.list_cancel_reapable(
                now=now, limit=_CLAIM_BATCH
            )
        except Exception:  # noqa: BLE001
            logger.warning("[roundtable-jobs] list_cancel_reapable failed", exc_info=True)
            return 0

        finalized = 0
        for job_id in candidate_ids:
            try:
                row = await self._store.finalize_cancel(job_id)
            except Exception:  # noqa: BLE001
                logger.warning(
                    "[roundtable-jobs] finalize orphaned cancellation failed for %s",
                    job_id,
                    exc_info=True,
                )
                continue
            if row is not None:
                finalized += 1
                logger.info(
                    "[roundtable-jobs] reaped orphaned cancellation for job %s", job_id
                )
        return finalized

    # ── outbox 补偿器 ────────────────────────────────────────────────────────

    def _draft_store_for(self, task_id: str | None):
        """task 作业 → task-draft 存储；普通作业 → 个人草稿存储（与执行器同规则）。"""
        return self._task_draft_store if task_id else self._draft_store

    async def compensate_once(self) -> int:
        """把 ``draft_persist_pending=1`` 的作业结果投影回草稿，成功后清标记。

        投影逻辑与执行器 ``_write_run_to_draft`` 的终态写一致：读最新草稿 → 按
        ``run_id=job-<jobId>`` upsert 一个 run（用作业行里的 dialogues/step3/status）→
        带 ``expected_version`` 乐观锁写回。写成功（或草稿已含该 run 的终态）即清标记。
        返回本轮成功投影的作业数。
        """
        try:
            pending = await self._store.list_draft_persist_pending(limit=_CLAIM_BATCH)
        except Exception:  # noqa: BLE001
            logger.warning("[roundtable-jobs] list_draft_persist_pending failed", exc_info=True)
            return 0
        if not pending:
            return 0

        projected = 0
        for job in pending:
            job_id = job.get("id")
            draft_id = job.get("draft_id")
            draft_store = self._draft_store_for(job.get("task_id"))
            if not draft_id or draft_store is None:
                # 无草稿可投影（如纯 task 作业但 task_store 未配）——直接清标记，避免死循环。
                with suppress(Exception):
                    await self._store.clear_draft_persist_pending(job_id)
                continue
            try:
                ok = await self._project_job_to_draft(draft_store, job)
            except Exception:  # noqa: BLE001
                logger.warning(
                    "[roundtable-jobs] compensate project failed for job %s", job_id, exc_info=True
                )
                continue
            if ok:
                with suppress(Exception):
                    await self._store.clear_draft_persist_pending(job_id)
                projected += 1
                logger.info("[roundtable-jobs] compensated draft write for job %s", job_id)
        return projected

    async def _project_job_to_draft(self, draft_store: Any, job: dict[str, Any]) -> bool:
        """把单个作业的最新状态投影成草稿 step2.runs 里的一个 run。成功返回 True。

        与执行器不同：这里没有 ``params``（作业可能跑在别的 worker / 已结束），只能用作业行
        里已落库的 ``dialogues``/``step3``/``status`` 重建 run —— 这些正是「研讨结果」本体，
        足够让前端加载草稿时看到完整的后台研讨记录与报告。
        """
        from deerflow.persistence.roundtable_drafts.sql import DraftConcurrentWriteError

        job_id = job.get("id")
        draft_id = job.get("draft_id")
        user_id = job.get("user_id")
        run_id = f"job-{job_id}"
        status = job.get("status") or "running"
        dialogues = job.get("dialogues") or []
        step3 = job.get("step3")

        # task 草稿不分权：get/update 不传 user_id；个人草稿按 user。
        is_task = bool(job.get("task_id"))

        for _attempt in range(3):
            try:
                draft = (
                    await draft_store.get_draft(draft_id)
                    if is_task
                    else await draft_store.get_draft(draft_id, user_id)
                )
            except Exception:  # noqa: BLE001
                draft = None
            if draft is None:
                # 草稿已被删除 —— 无处投影，视作完成（清标记）。
                return True
            expected_version = draft.get("version")

            step2 = dict(draft.get("step2") or {})
            runs = list(step2.get("runs") or [])
            idx = next((i for i, r in enumerate(runs) if r.get("id") == run_id), -1)
            created = runs[idx].get("createdAt") if idx >= 0 else (job.get("created_at") or datetime.now(UTC).isoformat())
            prev_step3 = runs[idx].get("step3") if idx >= 0 else None
            effective_step3 = step3 if step3 is not None else prev_step3

            # 从 dialogues 里提取席位信息（尽量还原 selectedAgents）。
            agents_seen: list[dict[str, str]] = []
            seen_ids: set[str] = set()
            for d in dialogues:
                aid = str(d.get("agentId") or "")
                if aid and aid not in seen_ids and d.get("role") == "seat":
                    seen_ids.add(aid)
                    agents_seen.append({"agent_id": aid, "name": str(d.get("name") or aid)})

            run = {
                "id": run_id,
                "title": "后台研讨",
                "createdAt": created,
                "source": job.get("orchestration_mode") or "recommend",
                "selectedAgents": agents_seen,
                "threadIds": job.get("thread_ids"),
                "coordinatorName": job.get("coordinator_name"),
                "step2RoundtableDialogues": dialogues,
                "lastLeaderContent": next(
                    (d.get("content") for d in reversed(dialogues) if d.get("role") == "leader"), ""
                ) or "",
                "hasConsensus": status == "done",
                "consensusPercentage": int(job.get("consensus_percentage") or (100 if status == "done" else 0)),
                "budgetLimit": 0,
                "orchestrationMode": job.get("orchestration_mode") or "recommend",
                "chain": job.get("chain"),
                "orchestrationPlan": job.get("orchestration_plan"),
                "step3": effective_step3,
                "jobId": job_id,
                "jobStatus": status,
            }
            runs = [run if i == idx else r for i, r in enumerate(runs)] if idx >= 0 else [*runs, run]
            step2["runs"] = runs
            step2["activeRunId"] = run_id
            has_report = effective_step3 is not None
            update_kwargs: dict[str, Any] = {
                "step2": step2,
                "furthest_step": 3 if (status == "done" and has_report) else 2,
            }
            if has_report:
                update_kwargs["step3"] = effective_step3

            try:
                if is_task:
                    await draft_store.update_draft(
                        draft_id, expected_version=expected_version, **update_kwargs
                    )
                else:
                    await draft_store.update_draft(
                        draft_id, user_id, expected_version=expected_version, **update_kwargs
                    )
                return True
            except DraftConcurrentWriteError:
                # 并发写冲突 —— 重读最新草稿再合并重试（与执行器同款乐观锁重试）。
                continue
            except Exception:  # noqa: BLE001
                logger.warning(
                    "[roundtable-jobs] compensate update_draft failed for %s", draft_id, exc_info=True
                )
                return False
        # 乐观锁重试耗尽：保留 pending 标记，等下一轮补偿器再试。
        return False