"""
Replay — re-run the intelligence pipeline over past calls, in a sandbox.

Live AI windows are the only way the correlator has ever been measured, and
each one costs days of real time and whatever the credits allow: a change
ships, AI goes on, traffic trickles in, someone pulls a dump. Recordings are
kept whether AI is on or not, so the traffic to measure against already
exists. A replay run takes a time range of real calls, feeds them through
the SAME pipeline code the live upload path runs (routers/upload.py
`_extract_and_correlate`) in their original order with the clock pinned to
each call's own end time, and writes everything to
replay_runs/{run_id}/calls|incidents instead of the live collections. The
same range can then be replayed after every change and the runs compared.

Three modes, cheapest last:
  audio        re-transcribe the saved audio (Whisper + correction), then
               extract and correlate. For ranges where AI was off.
  transcripts  reuse the transcript already on each call, re-run extraction
               and correlation.
  reuse        reuse the scenes an earlier run extracted, re-run correlation
               only. Extraction is an LLM call and never returns quite the
               same thing twice, so this is the mode that isolates a
               correlator change from extraction noise.

What never happens in a replay: alerts, summaries, vocabulary learning, and
any write to a live call or incident. The sandbox is enforced by the
ContextVar redirect in app/internal/firestore.py, not by this module
remembering to use different collection names.

One run at a time per process — a run spends real AI credits and its cost is
only estimated, so two concurrent runs would be two unbounded bills.
"""
import asyncio
import os
import statistics
import uuid
from collections import Counter
from datetime import datetime, timedelta, timezone
from typing import Optional

from app.config import settings
from app.internal import clock
from app.internal import firestore as fstore
from app.internal.feature_flags import force_flags, unforce_flags
from app.internal.logger import logger

RUNS = "replay_runs"
MODES = ("audio", "transcripts", "reuse")
MAX_CALLS = 5000
MAX_RANGE_DAYS = 7

# Extraction/transcription run ahead of correlation with this much
# concurrency. They depend only on the call itself; correlation depends on
# every call before it and is kept strictly in order.
PREFETCH = 6

# Rough per-unit AI prices for the pre-run estimate and the running tally.
# Estimates, not a bill — nothing in DRB reads a real invoice (server-26#45).
USD_WHISPER_PER_MIN    = 0.006
USD_PER_EXTRACTION     = 0.0005   # gpt-4o-mini scene extraction + embedding
USD_PER_CORRECTION     = 0.0003   # Gemini flash transcript correction
USD_PER_GEOCODE        = 0.005    # Google geocode, roughly one per located scene
USD_PER_LLM_CORRELATE  = 0.0005   # Gemini flash consensus decision

# Fields the pipeline writes onto a call doc. Stripped when a call is copied
# into the sandbox so the replay recomputes them instead of inheriting the
# live answer. Anything else on the doc (ids, times, talkgroup, srcaddr, audio
# location) is an input and is kept.
_DERIVED = {
    "transcript", "transcript_corrected", "transcript_not_speech",
    "segments", "segments_corrected", "scenes", "incident_id", "incident_ids",
    "tags", "location", "location_coords", "location_mentions", "units",
    "vehicles", "cleared_units", "severity", "incident_type", "type",
    "embedding", "skip_reason", "intelligence_started_at", "reassignment",
    "resolved", "has_updates", "audio_url",
}
_DERIVED_PREFIXES = ("corr_", "chatter_classifier_", "eval_")
_TRANSCRIPT_FIELDS = (
    "transcript", "transcript_corrected", "transcript_not_speech",
    "segments", "segments_corrected",
)

_active_run_id: Optional[str] = None
_active_task: Optional[asyncio.Task] = None
_cancel: set[str] = set()


class ReplayBusy(RuntimeError):
    pass


def sandbox_root(run_id: str) -> str:
    return f"{RUNS}/{run_id}"


def _scenes_coll(run_id: str) -> str:
    # Extracted scenes (embeddings included) live beside the sandbox, not on
    # its call docs, so reading a run's calls for metrics or the incident view
    # doesn't haul every scene's embedding along a second time.
    return f"{sandbox_root(run_id)}/scenes"


def active_run_id() -> Optional[str]:
    if _active_task is not None and not _active_task.done():
        return _active_run_id
    return None


# ---------------------------------------------------------------------------
# Call selection
# ---------------------------------------------------------------------------

def _as_dt(value) -> Optional[datetime]:
    if value is None:
        return None
    if isinstance(value, datetime):
        return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
    try:
        dt = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
    except ValueError:
        return None
    return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)


async def select_calls(
    org_id: str,
    date_from: datetime,
    date_to: datetime,
    system_ids: Optional[list[str]] = None,
    cap: int = MAX_CALLS,
) -> tuple[list[dict], bool]:
    """
    Live calls in [date_from, date_to] for this org, oldest first.

    Pages newest-first because the one composite index on calls that carries
    org_id is (org_id ASC, started_at DESC); ordering the other way would need
    a new index for no gain. Returns (calls, truncated) — truncated means the
    range holds more than `cap` calls and the caller must narrow it rather
    than silently replaying only part of it.
    """
    out: list[dict] = []
    cursor = None
    page = 1000
    while True:
        rows = await fstore.collection_where(
            "calls",
            [("org_id", "==", org_id),
             ("started_at", ">=", date_from),
             ("started_at", "<=", date_to)],
            order_by=[("started_at", "DESCENDING")],
            limit_to=page,
            start_after={"started_at": cursor} if cursor is not None else None,
        )
        for c in rows:
            if c.get("duplicate_of"):
                continue  # another node's copy — live never processes these either
            if system_ids and c.get("system_id") not in system_ids:
                continue
            out.append(c)
        if len(out) > cap:
            return sorted(out[:cap], key=_call_time), True
        if len(rows) < page:
            break
        cursor = rows[-1].get("started_at")
    return sorted(out, key=_call_time), False


def _call_time(call: dict) -> datetime:
    return _as_dt(call.get("started_at")) or datetime.min.replace(tzinfo=timezone.utc)


def _pipeline_time(call: dict) -> datetime:
    """When the live pipeline would have run for this call: at upload, i.e. call end."""
    return _as_dt(call.get("ended_at")) or _call_time(call)


def estimate(calls: list[dict], mode: str) -> dict:
    n = len(calls)
    with_transcript = sum(1 for c in calls if c.get("transcript_corrected") or c.get("transcript"))
    audio_min = sum(float(c.get("duration_s") or 0) for c in calls) / 60
    with_audio = sum(1 for c in calls if c.get("audio_gcs_uri"))
    # Roughly a third of calls carry a geocodable location (09-22 dump: 92/373).
    per_call = USD_PER_EXTRACTION + USD_PER_LLM_CORRELATE + USD_PER_GEOCODE / 3
    if mode == "audio":
        usd = audio_min * USD_WHISPER_PER_MIN + with_audio * (USD_PER_CORRECTION + per_call)
    elif mode == "transcripts":
        usd = with_transcript * per_call
    else:
        usd = n * USD_PER_LLM_CORRELATE
    return {
        "calls": n,
        "calls_with_transcript": with_transcript,
        "calls_with_audio": with_audio,
        "audio_minutes": round(audio_min, 1),
        "est_cost_usd": round(usd, 2),
    }


# ---------------------------------------------------------------------------
# Run lifecycle
# ---------------------------------------------------------------------------

async def start_run(
    *,
    org_id: str,
    date_from: datetime,
    date_to: datetime,
    mode: str,
    system_ids: Optional[list[str]],
    source_run_id: Optional[str],
    label: str,
    actor: str,
) -> dict:
    global _active_run_id, _active_task
    if active_run_id():
        raise ReplayBusy(f"Replay {active_run_id()} is still running.")
    if mode not in MODES:
        raise ValueError(f"mode must be one of {MODES}")
    if date_to <= date_from:
        raise ValueError("date_to must be after date_from")
    if date_to - date_from > timedelta(days=MAX_RANGE_DAYS):
        raise ValueError(f"Range is capped at {MAX_RANGE_DAYS} days.")
    if mode == "reuse":
        src = await fstore.doc_get(RUNS, source_run_id or "")
        if not src or src.get("org_id") != org_id:
            raise ValueError("reuse mode needs a source_run_id from an earlier run in this org")
        if src.get("status") != "done":
            raise ValueError("The source run did not finish; its scenes are incomplete.")

    calls, truncated = await select_calls(org_id, date_from, date_to, system_ids)
    if truncated:
        raise ValueError(f"Range holds more than {MAX_CALLS} calls — narrow it.")
    if not calls:
        raise ValueError("No calls in that range.")

    run_id = uuid.uuid4().hex[:12]
    now = datetime.now(timezone.utc).isoformat()
    doc = {
        "run_id":        run_id,
        "org_id":        org_id,
        "label":         label or "",
        "mode":          mode,
        "source_run_id": source_run_id if mode == "reuse" else None,
        "date_from":     date_from.isoformat(),
        "date_to":       date_to.isoformat(),
        "system_ids":    system_ids or [],
        "git_sha":       os.getenv("GIT_SHA", "unknown"),
        "created_by":    actor,
        "created_at":    now,
        "status":        "running",
        "estimate":      estimate(calls, mode),
        "progress":      {"total": len(calls), "done": 0, "errors": 0},
        "metrics":       None,
        "errors":        [],
    }
    await fstore.doc_set(RUNS, run_id, doc, merge=False)

    _active_run_id = run_id
    _active_task = asyncio.create_task(_run(run_id, org_id, calls, mode, source_run_id))
    return doc


def request_cancel(run_id: str) -> bool:
    if active_run_id() != run_id:
        return False
    _cancel.add(run_id)
    return True


async def get_run(run_id: str) -> Optional[dict]:
    doc = await fstore.doc_get(RUNS, run_id)
    return await _reconcile(doc) if doc else None


async def list_runs(org_id: str) -> list[dict]:
    docs = await fstore.collection_list(RUNS, org_id=org_id)
    docs = [await _reconcile(d) for d in docs]
    return sorted(docs, key=lambda d: d.get("created_at") or "", reverse=True)


async def _reconcile(doc: dict) -> dict:
    """A run left "running" by a process that restarted (a deploy) never finishes."""
    if doc.get("status") == "running" and doc.get("run_id") != active_run_id():
        doc["status"] = "interrupted"
        await fstore.doc_set(RUNS, doc["run_id"], {"status": "interrupted"})
    return doc


async def delete_run(run_id: str) -> None:
    if active_run_id() == run_id:
        raise ReplayBusy("Cancel the run before deleting it.")
    token = fstore.enter_sandbox(sandbox_root(run_id))
    try:
        for coll, key in (("calls", "call_id"), ("incidents", "incident_id"),
                          (_scenes_coll(run_id), "call_id")):
            for d in await fstore.collection_list(coll):
                if d.get(key):
                    await fstore.doc_delete(coll, d[key])
    finally:
        fstore.exit_sandbox(token)
    await fstore.doc_delete(RUNS, run_id)


async def sandbox_contents(run_id: str) -> tuple[list[dict], list[dict]]:
    token = fstore.enter_sandbox(sandbox_root(run_id))
    try:
        incidents = await fstore.collection_list("incidents")
        calls = await fstore.collection_list("calls")
    finally:
        fstore.exit_sandbox(token)
    return incidents, calls


# ---------------------------------------------------------------------------
# The run itself
# ---------------------------------------------------------------------------

def _flags_for(mode: str) -> dict[str, bool]:
    return {
        "stt_enabled": mode == "audio",
        "transcript_correction_enabled": mode == "audio",
        "correlation_enabled": True,
        "summaries_enabled": False,
        "vocabulary_learning_enabled": False,
    }


def _stored_input(call: dict) -> tuple[Optional[str], list]:
    """
    The transcript + segments live extraction was handed for this call.

    Not simply `transcript_corrected`: live extraction overwrites that field
    with its primary scene's rewrite (intelligence.py), so on a call with
    several scenes it now holds only scene 0's text. The corrector's own
    output survives intact in `segments_corrected`, so rebuild from those
    when they exist; otherwise correction never produced anything and live
    extraction read the raw Whisper transcript.
    """
    if call.get("transcript_not_speech"):
        return None, []  # transcribe_call hands nothing downstream for noise
    corrected = call.get("segments_corrected") or []
    if corrected:
        text = " ".join(str(seg.get("text") or "").strip() for seg in corrected).strip()
        return (text or call.get("transcript")), corrected
    return call.get("transcript"), call.get("segments") or []


def _extraction_fields(sb_call: dict) -> dict:
    """What extraction wrote onto the call doc (tags, units, location,
    embedding, skip_reason, ...), minus everything correlation wrote. A reuse
    run restores these so the orphan sweep, which reads them straight off the
    call doc, sees what it saw in the source run."""
    return {
        k: v for k, v in sb_call.items()
        if (k in _DERIVED or k.startswith("chatter_classifier_"))
        and k not in _TRANSCRIPT_FIELDS
        and k not in ("scenes", "incident_id", "incident_ids", "intelligence_started_at")
    }


def _sandbox_seed(call: dict, mode: str) -> dict:
    keep_transcript = mode in ("transcripts", "reuse")
    seed = {}
    for k, v in call.items():
        if k in _TRANSCRIPT_FIELDS:
            if keep_transcript:
                seed[k] = v
            continue
        if k in _DERIVED or k.startswith(_DERIVED_PREFIXES):
            continue
        seed[k] = v
    # Calls are seeded ahead of the replay clock (see PREFETCH). The orphan
    # re-correlation sweep selects status=="ended" calls by ended_at, so a
    # seeded call keeping its real status would be swept up as an "orphan"
    # before its own turn. Its real status is restored when it is processed.
    seed["status"] = "replay_pending"
    return seed


async def _prepare(call: dict, mode: str, source_scenes: dict[str, dict]) -> dict:
    """
    Everything per call that doesn't depend on other calls: seed the sandbox
    doc, then transcribe and/or extract. Runs ahead of correlation.
    Returns {"transcript", "scenes", "skip"} for the in-order stage.
    """
    from app.internal import intelligence, talkgroups, transcription

    call_id = call["call_id"]
    await fstore.doc_set("calls", call_id, _sandbox_seed(call, mode), merge=False)

    talkgroup_name = await talkgroups.resolve(
        call.get("system_id"), call.get("talkgroup_id"),
        hint=call.get("talkgroup_name"), call_doc=call,
    )

    transcript: Optional[str] = None
    segments: list = []
    if mode == "audio":
        if call.get("audio_gcs_uri"):
            transcript, segments = await transcription.transcribe_call(
                call_id, call["audio_gcs_uri"], talkgroup_name,
                system_id=call.get("system_id"), talkgroup_id=call.get("talkgroup_id"),
            )
    else:
        transcript, segments = _stored_input(call)

    if mode == "reuse":
        src = source_scenes.get(call_id)
        if src is None:
            return {"skip": "not_in_source_run", "talkgroup_name": talkgroup_name}
        if src.get("call_fields"):
            # skip_reason gates upload.py's no-scene fallback and the orphan
            # sweep correlates from tags/units/location on the call doc, so
            # extraction's call-level output comes across with its scenes.
            await fstore.doc_set("calls", call_id, src["call_fields"])
        return {"transcript": transcript, "scenes": src.get("scenes") or [],
                "talkgroup_name": talkgroup_name}

    scenes: list = []
    if transcript:
        scenes = await intelligence.extract_scenes(
            call_id, transcript, talkgroup_name,
            talkgroup_id=call.get("talkgroup_id"), system_id=call.get("system_id"),
            segments=segments, node_id=call.get("node_id"),
        )
    return {"transcript": transcript, "scenes": scenes, "talkgroup_name": talkgroup_name}


async def _sweeps_until(t: datetime, state: dict) -> None:
    """Run the live periodic sweeps (idle auto-resolve, orphan re-correlation) at every tick up to t."""
    from app.internal import recorrelation_sweep, summarizer

    interval = timedelta(minutes=settings.summary_interval_minutes)
    if state["last_sweep"] is None:
        state["last_sweep"] = t
        return
    while state["last_sweep"] + interval <= t:
        state["last_sweep"] += interval
        tok = clock.pin(state["last_sweep"])
        try:
            await summarizer._resolve_stale_incidents()
            await recorrelation_sweep._run_sweep_pass()
        finally:
            clock.unpin(tok)


async def _run(run_id: str, org_id: str, calls: list[dict], mode: str,
               source_run_id: Optional[str]) -> None:
    from app.routers.upload import _extract_and_correlate

    global _active_run_id
    progress = {"total": len(calls), "done": 0, "errors": 0, "skipped": 0,
                "extractions": 0, "audio_minutes": 0.0}
    errors: list[str] = []
    status = "done"

    source_scenes: dict[str, dict] = {}
    if mode == "reuse" and source_run_id:
        rows = await fstore.collection_list(_scenes_coll(source_run_id))
        source_scenes = {r["call_id"]: r for r in rows if r.get("call_id")}

    sb_token = fstore.enter_sandbox(sandbox_root(run_id))
    fl_token = force_flags(_flags_for(mode))
    try:
        sem = asyncio.Semaphore(PREFETCH)

        async def prep(call: dict):
            async with sem:
                tok = clock.pin(_pipeline_time(call))
                try:
                    return await _prepare(call, mode, source_scenes)
                finally:
                    clock.unpin(tok)

        pending: dict[int, asyncio.Task] = {}
        sweep_state = {"last_sweep": None}
        last_t = None
        for i, call in enumerate(calls):
            for j in range(i, min(i + PREFETCH * 2, len(calls))):
                if j not in pending:
                    pending[j] = asyncio.create_task(prep(calls[j]))
            if run_id in _cancel:
                status = "cancelled"
                break

            t = _pipeline_time(call)
            last_t = t
            try:
                prepared = await pending.pop(i)
                await _sweeps_until(t, sweep_state)
                if prepared.get("skip"):
                    progress["skipped"] += 1
                else:
                    tok = clock.pin(t)
                    try:
                        await fstore.doc_set("calls", call["call_id"], {
                            "status": call.get("status") or "ended",
                            "intelligence_started_at": t.isoformat(),
                        })
                        _, _, scenes = await _extract_and_correlate(
                            call_id=call["call_id"],
                            node_id=call.get("node_id"),
                            system_id=call.get("system_id"),
                            talkgroup_id=call.get("talkgroup_id"),
                            talkgroup_name=prepared["talkgroup_name"],
                            transcript=prepared["transcript"],
                            scenes=prepared["scenes"],
                        )
                    finally:
                        clock.unpin(tok)
                    # Kept whole (embedding included) so a later "reuse" run
                    # can correlate from exactly these scenes.
                    sb_call = await fstore.doc_get("calls", call["call_id"]) or {}
                    await fstore.doc_set(_scenes_coll(run_id), call["call_id"], {
                        "call_id": call["call_id"],
                        "scenes": scenes,
                        "call_fields": _extraction_fields(sb_call),
                    }, merge=False)
                    if prepared["transcript"] and mode != "reuse":
                        progress["extractions"] += 1
                if mode == "audio":
                    progress["audio_minutes"] += float(call.get("duration_s") or 0) / 60
            except Exception as e:
                progress["errors"] += 1
                if len(errors) < 20:
                    errors.append(f"{call.get('call_id')}: {type(e).__name__}: {e}"[:300])
                logger.warning(f"Replay {run_id}: call {call.get('call_id')} failed: {e}")
            progress["done"] = i + 1
            if (i + 1) % 25 == 0:
                await fstore.doc_set(RUNS, run_id, {"progress": dict(progress), "errors": errors})

        for task in pending.values():
            task.cancel()

        if status == "done" and last_t is not None:
            # Let every incident age out exactly as it would have live.
            await _sweeps_until(
                last_t + timedelta(minutes=settings.incident_auto_resolve_minutes
                                   + 2 * settings.summary_interval_minutes),
                sweep_state,
            )

        incidents = await fstore.collection_list("incidents")
        sb_calls = await fstore.collection_list("calls")
        metrics = compute_metrics(incidents, sb_calls)
        metrics["est_cost_usd"] = _running_cost(progress, metrics, mode)
    except Exception as e:
        status = "failed"
        errors.append(f"run: {type(e).__name__}: {e}"[:300])
        metrics = None
        logger.error(f"Replay {run_id} failed: {e}")
    finally:
        unforce_flags(fl_token)
        fstore.exit_sandbox(sb_token)
        _cancel.discard(run_id)
        _active_run_id = None

    await fstore.doc_set(RUNS, run_id, {
        "status": status,
        "progress": progress,
        "errors": errors,
        "metrics": metrics,
        "finished_at": datetime.now(timezone.utc).isoformat(),
    })
    logger.info(f"Replay {run_id} {status}: {progress}")


def _running_cost(progress: dict, metrics: dict, mode: str) -> float:
    usd = progress["audio_minutes"] * USD_WHISPER_PER_MIN
    if mode == "audio":
        usd += progress["extractions"] * USD_PER_CORRECTION
    usd += progress["extractions"] * (USD_PER_EXTRACTION + USD_PER_GEOCODE / 3)
    usd += metrics.get("llm_decisions", 0) * USD_PER_LLM_CORRELATE
    return round(usd, 2)


# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------

def compute_metrics(incidents: list[dict], calls: list[dict]) -> dict:
    """
    The numbers that say whether incidents are being tracked, from one run's
    sandbox. Same questions every correlation review has asked by hand, so two
    runs over the same range compare directly.
    """
    sizes = [len(i.get("call_ids") or []) for i in incidents]
    resolved_via = Counter(
        (i.get("resolved_via") or ("unknown" if i.get("status") == "resolved" else "still_active"))
        for i in incidents
    )
    corr_path: Counter = Counter()
    consensus: Counter = Counter()
    for c in calls:
        scenes = c.get("scenes") or {}
        records = [s.get("corr_debug") or {} for s in scenes.values()] if scenes else [c]
        for r in records:
            corr_path[r.get("corr_path") or "none"] += 1
            consensus[r.get("corr_consensus") or "none"] += 1
    linked = sum(1 for c in calls if c.get("incident_ids"))
    llm = sum(n for k, n in consensus.items() if k not in ("none", "rules_only"))
    return {
        "calls": len(calls),
        "calls_linked": linked,
        "calls_orphaned": len(calls) - linked,
        "incidents": len(incidents),
        "single_call_incidents": sum(1 for s in sizes if s == 1),
        "single_call_pct": round(100 * sum(1 for s in sizes if s == 1) / len(sizes), 1) if sizes else None,
        "median_calls_per_incident": statistics.median(sizes) if sizes else None,
        "max_calls_in_incident": max(sizes) if sizes else None,
        "incidents_with_units_cleared": sum(1 for i in incidents if i.get("units_cleared")),
        "incidents_with_coords": sum(1 for i in incidents if i.get("location_coords")),
        "resolved_via": dict(resolved_via),
        "corr_path": dict(corr_path),
        "corr_consensus": dict(consensus),
        "llm_decisions": llm,
    }