"""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 spawns those child Runs
by reusing the existing single-run execution path (``EvalEngine.run`` with an
``existing_run`` that carries ``campaign_id``). Spawned plan-entry indices are
persisted on the campaign so re-advancing the same clock never double-spawns.

This module also hosts the durable scheduler loop: a thin async shell that,
tick by tick, maps real wall-clock elapsed time (since the campaign's persisted
``started_at``) to a window offset and calls ``advance_campaign``. All authority
lives in the DB (window start, spawned progress, status), so the loop can be
torn down and rebuilt on restart without losing or duplicating work.
"""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import timezone
from typing import Optional

from sqlmodel import Session

from agenteval.evaluation.campaign_scheduler import clock_offset, decide_schedule
from agenteval.evaluation.engine import EvalEngine
from agenteval.models import Campaign, CampaignStatus, EvalRun, RunStatus, RunTrigger
from agenteval.storage.db import get_session, utc_now
from agenteval.storage.repository import (
    CampaignRepository,
    RunRepository,
    ScenarioRepository,
    TargetRepository,
)

_SPAWNED_KEY = "spawned_indices"
_ERRORS_KEY = "errors"

# 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")

# Live loop tasks + cooperative cancel events, keyed by campaign id. Authority
# is the DB; these are only the in-process handles for the running loop.
_tasks: dict[str, asyncio.Task] = {}
_cancel_events: dict[str, asyncio.Event] = {}


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

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


def _elapsed_seconds(started_at) -> float:
    """Real seconds since ``started_at`` (tolerating naive UTC from SQLite)."""
    if started_at.tzinfo is None:
        started_at = started_at.replace(tzinfo=timezone.utc)
    return (utc_now() - started_at).total_seconds()


def _spawned_indices(campaign: Campaign) -> set[int]:
    scheduler = (campaign.summary or {}).get("scheduler", {})
    return set(scheduler.get(_SPAWNED_KEY, []))


def current_window_offset(campaign: Campaign) -> float:
    """The campaign's live window position (seconds), clamped to the window.

    RUNNING campaigns derive it from the persisted ``started_at`` and
    ``time_scale``; not-yet-started campaigns report 0.
    """
    if campaign.started_at is None:
        return 0.0
    offset = clock_offset(
        elapsed_seconds=_elapsed_seconds(campaign.started_at),
        time_scale=campaign.time_scale,
    )
    return min(offset, float(campaign.window_seconds))


def campaign_progress(campaign: Campaign, runs: list[EvalRun]) -> dict:
    """Live progress of a campaign, derived from its child Runs.

    ``completed_runs`` counts only COMPLETED child Runs; failures stay out of
    this field (pass_rate semantics are ADR-0002's concern, not this counter).
    """
    return {
        "current_offset_seconds": current_window_offset(campaign),
        "spawned_runs": len(runs),
        "completed_runs": sum(1 for r in runs if r.status == RunStatus.COMPLETED),
    }


async def _spawn_child_run(campaign: Campaign, scenario_id: str, session: Session) -> str:
    """Create a Run owned by the campaign and drive it through the engine.

    The Run row is created on the caller's ``session``, but the engine runs on
    its own session (obtained internally, closed when the run ends) so one
    child Run's lifecycle never closes the campaign's session — the same
    isolation the single-run background task uses.
    """
    target = TargetRepository(session).get(campaign.target_id)
    scenario = ScenarioRepository(session).get(scenario_id)
    if not target or not scenario:
        raise ValueError(f"campaign target or scenario missing: {campaign.target_id}/{scenario_id}")

    run = RunRepository(session).create(
        EvalRun(
            target_id=campaign.target_id,
            scenario_id=scenario_id,
            scenario_version=scenario.version or 1,
            campaign_id=campaign.id,
            triggered_by=RunTrigger.CAMPAIGN,
            status=RunStatus.PENDING,
        )
    )
    engine = EvalEngine(target=target, scenario=scenario, triggered_by=RunTrigger.CAMPAIGN)
    await engine.run(existing_run=run)
    return run.id or ""


async def advance_campaign(
    *,
    campaign_id: str,
    elapsed_seconds: float,
    session: Session,
    cancel_event: Optional[asyncio.Event] = None,
) -> 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 skipped and recorded (so one bad entry never wedges the loop),
    but still marked spawned to avoid infinite retries. 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)
    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 or {}).get("scheduler", {}).get(_ERRORS_KEY, []))
    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
        try:
            for _ in range(due.entry.count):
                run_id = await _spawn_child_run(campaign, due.entry.scenario_id, session)
                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)
        spawned.add(due.index)
        # 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.
        scheduler_state: dict = {_SPAWNED_KEY: sorted(spawned)}
        if errors:
            scheduler_state[_ERRORS_KEY] = errors
        summary = dict(campaign.summary or {})
        summary["scheduler"] = scheduler_state
        campaign.summary = summary
        repo.update(campaign)

    return result


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


async def run_campaign_loop(campaign_id: str, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> None:
    """Drive one campaign to completion, ticking on real wall-clock time.

    The loop holds no authority: every tick it reloads the campaign from the DB,
    recomputes elapsed time from the persisted ``started_at``, and advances.
    It exits when the window finishes, the campaign leaves RUNNING (e.g. it was
    cancelled), or the cooperative cancel event fires.
    """
    cancel = _cancel_events.setdefault(campaign_id, asyncio.Event())
    session = get_session()
    try:
        while not cancel.is_set():
            repo = CampaignRepository(session)
            campaign = repo.get(campaign_id)
            if not campaign or campaign.status != CampaignStatus.RUNNING or campaign.started_at is None:
                return

            elapsed = _elapsed_seconds(campaign.started_at)
            result = await advance_campaign(
                campaign_id=campaign_id,
                elapsed_seconds=elapsed,
                session=session,
                cancel_event=cancel,
            )
            if result and result.finished:
                current = repo.get(campaign_id)
                # Only complete if still running (not cancelled meanwhile).
                if current and current.status == CampaignStatus.RUNNING:
                    current.status = CampaignStatus.COMPLETED
                    current.completed_at = utc_now()
                    repo.update(current)
                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()
        _cancel_events.pop(campaign_id, None)
        _tasks.pop(campaign_id, None)


def start_campaign(campaign_id: str, session: Session, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> Optional[asyncio.Task]:
    """Move a campaign into RUNNING (stamping ``started_at`` on first start) and
    launch its loop. Reused for both create-then-start and restart recovery:
    a PLANNED campaign gets a fresh ``started_at``; an already-RUNNING one keeps
    its original window start so recovery resumes at the correct offset.
    """
    repo = CampaignRepository(session)
    campaign = repo.get(campaign_id)
    if not campaign or campaign.status in (CampaignStatus.COMPLETED, CampaignStatus.CANCELLED, CampaignStatus.FAILED):
        return None
    if campaign_id in _tasks:
        return _tasks[campaign_id]

    if campaign.status == CampaignStatus.PLANNED:
        campaign.status = CampaignStatus.RUNNING
        campaign.started_at = utc_now()
        repo.update(campaign)

    _cancel_events.setdefault(campaign_id, asyncio.Event())
    task = asyncio.create_task(
        run_campaign_loop(campaign_id, tick_seconds=tick_seconds),
        name=f"campaign-{campaign_id}",
    )
    _tasks[campaign_id] = task
    return task


def request_cancel(campaign_id: str) -> None:
    """Signal the loop (if live) to stop spawning and exit promptly."""
    event = _cancel_events.get(campaign_id)
    if event is not None:
        event.set()


def resume_running_campaigns(session: Session, *, tick_seconds: float = DEFAULT_TICK_SECONDS) -> int:
    """On startup, relaunch a loop for every campaign left in RUNNING."""
    resumed = 0
    for campaign in CampaignRepository(session).list_all():
        if campaign.status == CampaignStatus.RUNNING and campaign.id:
            start_campaign(campaign.id, session, tick_seconds=tick_seconds)
            resumed += 1
    return resumed


async def shutdown_all() -> None:
    """Gracefully stop all live loops (Web app shutdown)."""
    for event in list(_cancel_events.values()):
        event.set()
    tasks = list(_tasks.values())
    for task in tasks:
        task.cancel()
    for task in tasks:
        # CancelledError is expected here (we just cancelled the task) and does
        # not derive from Exception in 3.8+, so it must be listed explicitly.
        try:
            await task
        except (asyncio.CancelledError, Exception):
            pass
    _tasks.clear()
    _cancel_events.clear()
