"""Deep-research job dispatcher.

Mirrors :class:`RoundtableJobDispatcher` exactly in shape — same wake-mechanism
(``asyncio.Event`` + poll), same claim flow (``list_claimable`` → ``claim_job``
CAS → ``executor.start_job``), same startup-scan, same cancel-reaper — minus the
roundtable outbox compensator (deep research has no versioned draft to project
back into; the session row IS the projection target, written directly by the
executor).

All workers run the same dispatcher; whichever ``claim_job`` returns a row
(rowcount==1) executes — consistency comes from the DB, not in-process locks.
"""

from __future__ import annotations

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

from app.gateway.deep_research_job_executor import (
    _LEASE_TTL_SECONDS,
    WORKER_ID,
    DeepResearchJobExecutor,
)

logger = logging.getLogger(__name__)

_DISPATCH_POLL_INTERVAL = 3.0
_CLAIM_BATCH = 20


class DeepResearchJobDispatcher:
    """Cross-worker scheduler for deep-research jobs (one instance per worker)."""

    def __init__(self, store, executor: DeepResearchJobExecutor, *, worker_id: str | None = None) -> None:
        self._store = store
        self._executor = executor
        self._worker_id = worker_id or WORKER_ID
        self._wake = asyncio.Event()
        self._task: asyncio.Task | None = None
        self._stopping = False

    # ── lifecycle ───────────────────────────────────────────────────────────

    def start(self) -> None:
        """Start the background loop (called from lifespan startup). Idempotent."""
        if self._task is not None and not self._task.done():
            return
        self._stopping = False
        self._task = asyncio.create_task(self._loop(), name="deep-research-job-dispatcher")
        logger.info("[deep-research] dispatcher started (worker=%s)", self._worker_id)

    async def stop(self) -> None:
        """Stop the background loop (called from 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 cancel_running(self, job_id: str) -> bool:
        """Stop an in-process executor task. Safe if the job is not on this worker."""
        return self._executor.cancel(job_id)

    def nudge(self) -> None:
        """Wake the loop immediately (called after start/resume to cut latency)."""
        self._wake.set()

    # ── main loop ───────────────────────────────────────────────────────────

    async def _loop(self) -> None:
        # Startup scan: pick up historically queued / expired-lease jobs.
        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()

    async def _safe_dispatch(self) -> None:
        try:
            await self.dispatch_once()
        except Exception:  # noqa: BLE001
            logger.exception("[deep-research] dispatch round failed")

    # ── claim ───────────────────────────────────────────────────────────────

    async def dispatch_once(self) -> int:
        """Scan for claimable jobs and try to claim each.  Returns started count."""
        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("[deep-research] 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:
                peek = await self._store.get_unscoped(job_id)
            except Exception:  # noqa: BLE001
                peek = None
            # Expired-lease ``running`` rows used to be reclaimed and resumed.
            # After a full Gateway restart that re-occupied the per-user 429
            # latch.  Drop them instead of auto-resuming; queued jobs still start.
            if peek is not None and peek.get("status") == "running":
                owner = str(peek.get("user_id") or "")
                with suppress(Exception):
                    if owner:
                        await self._store.request_cancel(job_id, user_id=owner)
                    else:
                        await self._store.request_cancel(job_id)
                with suppress(Exception):
                    await self._store.finalize_cancel(job_id)
                logger.info("[deep-research] dropped stale running job %s instead of reclaiming", job_id)
                continue
            try:
                row = await self._store.claim_job(
                    job_id, lease_owner=self._worker_id, lease_until=lease_until
                )
            except Exception:  # noqa: BLE001
                logger.warning("[deep-research] claim_job failed for %s", job_id, exc_info=True)
                continue
            if row is None:
                continue  # another worker won / status changed

            snapshot = row.get("input_snapshot")
            if not snapshot or not isinstance(snapshot, dict):
                logger.error("[deep-research] job %s claimed but snapshot missing; marking error", job_id)
                with suppress(Exception):
                    await self._store.update_progress(
                        job_id,
                        lease_owner=self._worker_id,
                        status="failed",
                        error_code="SNAPSHOT_MISSING",
                        error_message="input_snapshot 缺失或损坏，无法在后台重建研究",
                    )
                continue

            self._executor.start_job(snapshot, job_id=job_id, lease_owner=self._worker_id)
            started += 1
            logger.info(
                "[deep-research] worker=%s claimed job %s (attempt=%s)",
                self._worker_id,
                job_id,
                row.get("attempt"),
            )
        return started

    async def _reap_cancel_requested(self, now: datetime) -> int:
        """Finalize orphaned cancellations whose lease expired (worker crash)."""
        try:
            candidate_ids = await self._store.list_cancel_reapable(now=now, limit=_CLAIM_BATCH)
        except Exception:  # noqa: BLE001
            return 0

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


__all__ = ["DeepResearchJobDispatcher"]