"""Campaign runner — the side-effect shell around the pure scheduler.

Given a campaign and a clock position (real elapsed seconds), it asks
``campaign_scheduler.decide_schedule`` what is due, then durably claims each
plan occurrence before reusing the existing single-run execution path
(``EvalEngine.run`` with an ``existing_run``). Child-Run identities are the
primary idempotency authority; the legacy campaign summary remains a fallback
for pre-identity Runs created before this migration.

``CampaignRuntime`` owns the durable scheduler loop and exposes only lifecycle-
level operations. Clock mapping, reconciliation, ticking, claims and process-
local task handles remain implementation details. All authority lives in the
DB, so loops can be rebuilt on restart without losing or duplicating work.
"""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Awaitable, Callable, Optional

from sqlmodel import Session

from agenteval.evaluation.analysis import resolve_analysis_model
from agenteval.evaluation.campaign_lifecycle import complete_campaign
from agenteval.evaluation.campaign_lifecycle import start_campaign as start_campaign_lifecycle
from agenteval.evaluation.campaign_scheduler import (
    TickAction,
    clock_offset,
    decide_schedule,
    decide_tick,
    elapsed_seconds,
    resolve_finalize,
)
from agenteval.evaluation.engine import EvalEngine
from agenteval.evaluation.intelligence_jobs import enqueue_campaign_analysis
from agenteval.models import (
    Campaign,
    CampaignStatus,
    CampaignSummary,
    EvalRun,
    RunStatus,
    RunTrigger,
    Scenario,
    SchedulerState,
)
from agenteval.storage.db import get_session, utc_now
from agenteval.storage.repository import (
    CampaignRepository,
    CampaignRunClaimStatus,
    RunRepository,
    ScenarioRepository,
    TargetRepository,
)
from agenteval.task_registry import TaskRegistry

# Real wall-clock seconds between scheduler ticks. time_scale compresses the
# *window*, not the tick cadence — production 24h campaigns still tick slowly.
DEFAULT_TICK_SECONDS = 1.0

_logger = logging.getLogger("agenteval")


ChildRunExecutor = Callable[[Campaign, Scenario, EvalRun, Session], Awaitable[None]]


@dataclass
class AdvanceResult:
    """Outcome of advancing a campaign's clock once."""

    spawned_run_ids: list[str] = field(default_factory=list)
    finished: bool = False


@dataclass
class CampaignRecoveryResult:
    """Outcome of reconciling child Runs left by a previous process."""

    resumed_run_ids: list[str] = field(default_factory=list)
    failed_run_ids: list[str] = field(default_factory=list)


@dataclass(frozen=True)
class CampaignRuntimeRecovery:
    """Durable Campaign work reconciled during process startup."""

    interrupted_runs: int = 0
    resumed_campaigns: int = 0


def _spawned_indices(campaign: Campaign, session: Optional[Session] = None) -> set[int]:
    """Return plan entries completed by durable identities or legacy progress.

    Once any durable occurrence exists for an entry, identities replace the
    legacy summary for that entry. This repairs summaries written by older
    runners after only part of a multi-occurrence entry was claimed.
    """
    legacy = set(campaign.summary.scheduler.spawned_indices) if campaign.summary else set()
    if not campaign.id or session is None:
        return legacy

    occurrences: dict[int, set[int]] = {}
    for run in RunRepository(session).list_by_campaign(campaign.id):
        if run.campaign_plan_index is None or run.campaign_occurrence_index is None:
            continue
        occurrences.setdefault(run.campaign_plan_index, set()).add(run.campaign_occurrence_index)

    return {
        plan_index
        for plan_index, entry in enumerate(campaign.plan)
        if (plan_index in occurrences and set(range(entry.count)).issubset(occurrences[plan_index]))
        or (plan_index not in occurrences and plan_index in legacy)
    }


async def _execute_child_run(
    campaign: Campaign,
    scenario: Scenario,
    run: EvalRun,
    session: Session,
) -> None:
    target = TargetRepository(session).get(campaign.target_id)
    if not target:
        raise ValueError(f"campaign target missing: {campaign.target_id}")
    engine = EvalEngine(target=target, scenario=scenario, triggered_by=RunTrigger.CAMPAIGN)
    await engine.run(existing_run=run)


async def _spawn_child_run(
    campaign: Campaign,
    scenario_id: str,
    *,
    plan_index: int,
    occurrence_index: int,
    session: Session,
    execute_child_run: ChildRunExecutor = _execute_child_run,
) -> Optional[str]:
    """Claim one occurrence and drive its pending Run through the engine.

    A rejected claim means the Campaign stopped being runnable between the
    scheduler decision and this write; callers must stop spawning. An existing
    pending Run is resumed, while an existing running/completed/failed Run is
    returned without replaying external messages.
    """
    scenario = ScenarioRepository(session).get(scenario_id)
    if not scenario:
        raise ValueError(f"campaign scenario missing: {scenario_id}")

    claim = RunRepository(session).claim_campaign_run(
        campaign_id=campaign.id or "",
        scenario_id=scenario_id,
        scenario_version=scenario.version or 1,
        plan_index=plan_index,
        occurrence_index=occurrence_index,
    )
    if claim.status in (CampaignRunClaimStatus.NOT_FOUND, CampaignRunClaimStatus.CONFLICT):
        return None
    run = claim.run
    if run is None:
        raise RuntimeError(f"Campaign child claim returned no Run: {claim.status.value}")

    if run.status is RunStatus.PENDING:
        await execute_child_run(campaign, scenario, run, session)
    return run.id or ""


async def _advance_campaign(
    *,
    campaign_id: str,
    elapsed_seconds: float,
    session: Session,
    cancel_event: Optional[asyncio.Event] = None,
    execute_child_run: ChildRunExecutor = _execute_child_run,
) -> Optional[AdvanceResult]:
    """Advance the campaign clock to ``elapsed_seconds`` and spawn due Runs.

    ``elapsed_seconds`` is real wall-clock time since the window started; it is
    mapped to a window offset via ``time_scale``. A due entry whose spawning
    raises is recorded but remains due until every occurrence has a durable
    child-Run identity. Returns ``None`` if the campaign does not exist.
    """
    repo = CampaignRepository(session)
    campaign = repo.get(campaign_id)
    if not campaign:
        return None

    offset = clock_offset(elapsed_seconds=elapsed_seconds, time_scale=campaign.time_scale)
    spawned = _spawned_indices(campaign, session)
    decision = decide_schedule(
        plan=campaign.plan,
        window_seconds=campaign.window_seconds,
        clock_offset_seconds=offset,
        spawned_indices=spawned,
    )

    result = AdvanceResult(finished=decision.finished)
    errors: list[dict] = list(campaign.summary.scheduler.errors) if campaign.summary else []
    for due in decision.due:
        # Cancellation stops further spawning; runs already in flight finish.
        if cancel_event is not None and cancel_event.is_set():
            break
        claim_rejected = False
        try:
            for occurrence_index in range(due.entry.count):
                run_id = await _spawn_child_run(
                    campaign,
                    due.entry.scenario_id,
                    plan_index=due.index,
                    occurrence_index=occurrence_index,
                    session=session,
                    execute_child_run=execute_child_run,
                )
                if run_id is None:
                    claim_rejected = True
                    break
                result.spawned_run_ids.append(run_id)
        except Exception as exc:  # one entry failing must not wedge the campaign
            errors.append({"entry_index": due.index, "error": str(exc)})
            _logger.warning("活动 %s 计划条目 %d 派生失败（跳过）: %s", campaign_id, due.index, exc)
        if claim_rejected:
            break
        spawned = _spawned_indices(campaign, session)
        # Persist progress per entry: a failure partway through a multi-entry
        # advance must never lose which entries already spawned, since restart
        # recovery reads this back from the DB. Mutate the existing summary so
        # any unknown top-level keys survive the read-modify-write.
        summary = campaign.summary or CampaignSummary()
        summary.scheduler = SchedulerState(spawned_indices=sorted(spawned), errors=errors)
        campaign.summary = summary
        repo.save_scheduler_state(campaign_id, summary)

    return result


def _auto_start_analysis(campaign: Campaign, session: Session) -> None:
    """Enqueue 智能分析 when a 正式线 campaign completes.

    Accelerated campaigns, an unresolvable analysis model, and any enqueue
    failure all skip silently — the analysis is an enhancement and must never
    block or break campaign completion.
    """
    if campaign.time_scale != 1 or not campaign.id:
        return
    try:
        if resolve_analysis_model(campaign, session) is None:
            return
        enqueue_campaign_analysis(campaign.id, triggered_by="auto")
    except Exception as exc:
        _logger.warning("活动 %s 自动分析触发失败（已跳过）: %s", campaign.id, exc)


async def _reconcile_campaign_child_runs(
    campaign_id: str,
    session: Session,
    *,
    cancel_event: Optional[asyncio.Event] = None,
    execute_child_run: ChildRunExecutor = _execute_child_run,
) -> Optional[CampaignRecoveryResult]:
    """Resume safe pending claims and fail running Runs without replaying.

    A pending claim has not crossed the Engine's send seam and can be resumed.
    A running Run may already have reached the target, so recovery records an
    explicit interruption instead of executing it again.
    """
    campaign = CampaignRepository(session).get(campaign_id)
    if campaign is None:
        return None
    if campaign.status is not CampaignStatus.RUNNING:
        return CampaignRecoveryResult()

    run_repo = RunRepository(session)
    result = CampaignRecoveryResult(failed_run_ids=run_repo.mark_campaign_running_interrupted(campaign_id))
    for run in run_repo.list_pending_campaign_children(campaign_id):
        if cancel_event is not None and cancel_event.is_set():
            break
        plan_index = run.campaign_plan_index
        occurrence_index = run.campaign_occurrence_index
        valid_identity = (
            plan_index is not None
            and occurrence_index is not None
            and 0 <= plan_index < len(campaign.plan)
            and 0 <= occurrence_index < campaign.plan[plan_index].count
            and run.scenario_id == campaign.plan[plan_index].scenario_id
        )
        if not valid_identity:
            failed = run_repo.mark_pending_campaign_child_interrupted(run.id or "")
            if failed and failed.id:
                result.failed_run_ids.append(failed.id)
            continue

        try:
            resumed_id = await _spawn_child_run(
                campaign,
                run.scenario_id,
                plan_index=plan_index,
                occurrence_index=occurrence_index,
                session=session,
                execute_child_run=execute_child_run,
            )
            if resumed_id is None:
                break
            result.resumed_run_ids.append(resumed_id)
        except Exception as exc:
            fresh = run_repo.get(run.id or "")
            if fresh is not None and fresh.status is RunStatus.PENDING:
                failed = run_repo.mark_pending_campaign_child_interrupted(run.id or "")
                if failed and failed.id:
                    result.failed_run_ids.append(failed.id)
            elif fresh is not None and fresh.status is RunStatus.FAILED and fresh.id:
                result.failed_run_ids.append(fresh.id)
            _logger.warning("活动 %s 子运行 %s 恢复失败: %s", campaign_id, run.id, exc)

    return result


# ── durable scheduler loop ──────────────────────────────────────────────────


async def _run_campaign_loop(
    campaign_id: str,
    cancel: asyncio.Event,
    *,
    session_factory: Callable[[], Session],
    now: Callable[[], datetime],
    tick_seconds: float,
    execute_child_run: ChildRunExecutor,
) -> None:
    """Drive one campaign to completion, ticking on real wall-clock time.

    The loop holds no authority and no decision logic: every tick it reloads the
    campaign, asks the pure ``decide_tick`` what to do, and only performs I/O.
    It exits when the window finishes, the campaign leaves RUNNING (e.g. it was
    cancelled), or the cooperative cancel event fires.
    """
    session = session_factory()
    try:
        await _reconcile_campaign_child_runs(
            campaign_id,
            session,
            cancel_event=cancel,
            execute_child_run=execute_child_run,
        )
        while not cancel.is_set():
            repo = CampaignRepository(session)
            campaign = repo.get(campaign_id)
            if not campaign:
                return

            current_time = now()
            decision = decide_tick(
                now=current_time,
                started_at=campaign.started_at,
                status=campaign.status,
                window_seconds=campaign.window_seconds,
                time_scale=campaign.time_scale,
                plan=campaign.plan,
                spawned_indices=_spawned_indices(campaign, session),
            )
            if decision.action is TickAction.STOP:
                return

            elapsed = elapsed_seconds(now=current_time, started_at=campaign.started_at)
            await _advance_campaign(
                campaign_id=campaign_id,
                elapsed_seconds=elapsed,
                session=session,
                cancel_event=cancel,
                execute_child_run=execute_child_run,
            )

            if decision.action is TickAction.COMPLETE:
                current = repo.get(campaign_id)
                # The lifecycle CAS is the authoritative cancel-race guard.
                if current and resolve_finalize(current.status) is TickAction.COMPLETE:
                    try:
                        completed = complete_campaign(campaign_id=campaign_id, session=session)
                    except Exception as exc:
                        _logger.warning("活动 %s 完成迁移失败: %s", campaign_id, exc)
                    else:
                        _auto_start_analysis(completed, session)
                        return

            try:
                await asyncio.wait_for(cancel.wait(), timeout=tick_seconds)
            except asyncio.TimeoutError:
                pass
    except Exception as exc:  # a crashing loop must not take down the app
        _logger.warning("活动调度循环 %s 异常退出: %s", campaign_id, exc)
    finally:
        session.close()


class CampaignRuntime:
    """Own durable Campaign execution behind one lifecycle-level interface."""

    def __init__(
        self,
        *,
        session_factory: Callable[[], Session] = get_session,
        now: Callable[[], datetime] = utc_now,
        tick_seconds: float = DEFAULT_TICK_SECONDS,
        execute_child_run: ChildRunExecutor = _execute_child_run,
    ) -> None:
        self._session_factory = session_factory
        self._now = now
        self._tick_seconds = tick_seconds
        self._execute_child_run = execute_child_run
        self._registry = TaskRegistry()

    def start(self, campaign_id: str) -> bool:
        """Start or resume one durable Campaign, idempotently."""
        session = self._session_factory()
        try:
            started = start_campaign_lifecycle(session, campaign_id)
        finally:
            session.close()
        if started is None:
            return False
        self._launch(campaign_id)
        return True

    def cancel(self, campaign_id: str) -> None:
        """Signal the process-local loop after durable cancellation commits."""
        self._registry.cancel(campaign_id)

    def recover(self) -> CampaignRuntimeRecovery:
        """Repair interrupted Runs and relaunch every running Campaign."""
        session = self._session_factory()
        try:
            interrupted_runs = RunRepository(session).mark_orphans_failed()
            running_ids = [
                campaign.id
                for campaign in CampaignRepository(session).list_all()
                if campaign.status is CampaignStatus.RUNNING and campaign.id
            ]
        finally:
            session.close()
        for campaign_id in running_ids:
            self._launch(campaign_id)
        return CampaignRuntimeRecovery(
            interrupted_runs=interrupted_runs,
            resumed_campaigns=len(running_ids),
        )

    async def shutdown(self) -> None:
        """Gracefully stop all process-local Campaign loops."""
        await self._registry.shutdown_all()

    def _launch(self, campaign_id: str) -> None:
        self._registry.launch(
            campaign_id,
            lambda cancel: _run_campaign_loop(
                campaign_id,
                cancel,
                session_factory=self._session_factory,
                now=self._now,
                tick_seconds=self._tick_seconds,
                execute_child_run=self._execute_child_run,
            ),
        )


campaign_runtime = CampaignRuntime()
