"""ReplayEngine — Phase 4: re-dispatch historical runs under candidate policy.

A replay takes a historical ``correlation_id`` (a completed
``invoke_skill`` run) and re-executes the same skill under a DIFFERENT
policy_id (a candidate reflex). The utility delta between the original
and the replay is one piece of evidence the PromotionArena needs.

Honest limitation: replays run against CURRENT vault state, not frozen
historical state. The env_fingerprint of the replay will usually differ
from the original — we record both so reviewers can see the drift. A
hermetic replay would require vault snapshot isolation, which is a
non-trivial infra lift and not part of Phase 4.

This is Phase 4 scope. The WorkQueue handler ``evolution_replay`` accepts
``{"candidate_id": ..., "correlation_id_original": ...}`` and writes a
``ReplayJob`` row. The job is processed via ``replay_job()`` below.
"""

from __future__ import annotations

import logging
import time
from typing import TYPE_CHECKING, Any

from agntspace.activity.decisions import (
    install_scope_mutations,
    new_correlation_id,
    reset_policy_id,
    reset_run_id,
    reset_scope_mutations,
    set_policy_id,
    set_run_id,
)
from agntspace.evolution.manifest import build_run_manifest
from agntspace.evolution.schemas import POLICY_INCUMBENT_DEFAULT

if TYPE_CHECKING:
    from agntspace.evolution.ledger import EvolutionLedger
    from agntspace.evolution.service import EvolutionService

log = logging.getLogger(__name__)


class ReplayEngine:
    """Re-dispatches historical runs under candidate policies.

    Wired from the WorkQueue handler ``evolution_replay`` registered in
    main.py. Each replay produces a ``ReplayJob`` row with
    ``utility_delta`` populated, which the PromotionArena reads to check
    the ``min_replays >= 3`` gate.
    """

    def __init__(
        self,
        *,
        ledger: EvolutionLedger,
        service: EvolutionService,
        orchestrator: Any | None = None,
    ) -> None:
        self._ledger = ledger
        self._service = service
        self._orchestrator = orchestrator

    async def enqueue_replay(
        self, *, candidate_id: str, correlation_id_original: str,
    ) -> int:
        """Create a pending replay job. Returns job_id.

        The job runs later (either on the heartbeat schedule or when the
        PromotionArena triggers it via the WorkQueue).
        """
        return await self._ledger.enqueue_replay(
            incumbent_id=POLICY_INCUMBENT_DEFAULT,
            candidate_id=candidate_id,
            correlation_id_original=correlation_id_original,
        )

    async def replay_job(self, payload: dict[str, Any]) -> dict[str, Any]:
        """Execute one replay. WorkQueue handler entry point.

        ``payload`` shape: ``{"candidate_id": str, "correlation_id_original": str}``
        or ``{"job_id": int}`` to pick up a pre-enqueued pending job.
        """
        candidate_id = str(payload.get("candidate_id", ""))
        correlation_id_original = str(payload.get("correlation_id_original", ""))
        if not candidate_id or not correlation_id_original:
            return {"status": "failed", "reason": "missing required fields"}

        # Defensive: orchestrator wiring is required for live re-dispatch.
        if self._orchestrator is None:
            return {"status": "skipped", "reason": "no_orchestrator"}

        # Fetch the original run so we can reconstruct the invocation.
        # We query by correlation_id, not run_id, because replays reference
        # the causal chain (correlation_id is the stable join key).
        original_runs = await self._ledger.list_runs(limit=500)
        original = next(
            (r for r in original_runs if r.get("correlation_id") == correlation_id_original),
            None,
        )
        if original is None:
            return {"status": "failed", "reason": f"original run not found: {correlation_id_original}"}

        skill_name = original["skill_name"]
        agent_name = original["agent_name"]
        caller = original.get("caller", "replay")

        # Build a fresh correlation_id + run_id for the replay — its
        # decisions + outcome are attributed to this new scope, not the
        # original. The ``parent_run_id`` field on the manifest links
        # back to the source run for auditability.
        replay_corr = new_correlation_id(f"replay-{candidate_id}")

        # Open our own correlation + mutation + policy scope. We ALSO
        # skip the full Agent.invoke_skill path and dispatch directly —
        # we don't want invoke_skill's evolution prolog/epilog to run
        # (that would create ANOTHER run row and muddy attribution).
        cid_token = None
        mut_counter = None
        mut_token = None
        run_token = None
        policy_token = None

        manifest = build_run_manifest(
            skill_name=skill_name,
            agent_name=agent_name,
            caller=f"{caller}.replay",
            correlation_id=replay_corr,
            parent_run_id=original["run_id"],
            policy_id=candidate_id,
            vault_reader=getattr(self._service, "_vault_reader", None),
        )
        started_perf = time.perf_counter()

        try:
            from agntspace.activity.decisions import reset_correlation_id, set_correlation_id
            cid_token = set_correlation_id(replay_corr)
            mut_counter, mut_token = install_scope_mutations()
            run_token = set_run_id(manifest.run_id)
            policy_token = set_policy_id(candidate_id)

            # Persist the replay's manifest up-front, same as invoke_skill
            await self._service.persist_run(manifest)
            await self._service.persist_selector_telemetry(
                run_id=manifest.run_id,
                policy_id=candidate_id,
                eligible_ids=[candidate_id, POLICY_INCUMBENT_DEFAULT],
                top_samples=[{
                    "reflex_id": candidate_id,
                    "sampled_score": 1.0,
                    "posterior_mean": 0.0,
                    "chosen": True,
                    "replay": True,
                }],
                chosen_reason=f"replay:for:{correlation_id_original}",
            )

            # Actually dispatch — this is the live re-execution
            try:
                task = (
                    f"Replay of original run {original['run_id']}. "
                    f"Re-execute skill '{skill_name}' under candidate policy "
                    f"'{candidate_id}' for A/B comparison. "
                    "Call real tools; do not hallucinate."
                )
                await self._orchestrator.dispatch(
                    agent_name=agent_name,
                    task=task,
                    context=f"Replay context — original correlation: {correlation_id_original}",
                    skill_scope=skill_name,
                    policy_id=candidate_id,
                    run_id=manifest.run_id,
                )
                replay_failed = False
                replay_error: str | None = None
            except Exception as exc:
                replay_failed = True
                replay_error = f"{type(exc).__name__}: {exc}"
                log.info("replay dispatch raised: %s", replay_error)

            # Compute outcome + persist for the replay
            duration_ms = int((time.perf_counter() - started_perf) * 1000)
            await self._service.persist_outcome_for_run(
                manifest=manifest,
                succeeded=not replay_failed and not mut_counter.is_zero_effect(),
                zero_effect=mut_counter.is_zero_effect(),
                blocked=False,
                error=replay_error,
                duration_ms=duration_ms,
                correlation_id=replay_corr,
            )

            # Compute utility delta
            original_utility = float(
                (original.get("outcome") or {}).get("utility", 0.0),
            )
            replay_run = await self._service.get_run(manifest.run_id)
            replay_utility = float(
                (replay_run or {}).get("outcome", {}).get("utility", 0.0),
            )
            utility_delta = replay_utility - original_utility

            # Record the replay job result
            await self._ledger.complete_replay(
                candidate_id=candidate_id,
                correlation_id_original=correlation_id_original,
                correlation_id_replay=replay_corr,
                utility_delta=utility_delta,
            )

            return {
                "status": "done",
                "candidate_id": candidate_id,
                "correlation_id_original": correlation_id_original,
                "correlation_id_replay": replay_corr,
                "original_utility": round(original_utility, 4),
                "replay_utility": round(replay_utility, 4),
                "utility_delta": round(utility_delta, 4),
                "replay_failed": replay_failed,
                "replay_error": replay_error,
            }
        finally:
            # Reset contextvars in reverse order of setting — strict
            # scope discipline so a replay can't leak state into the
            # next work-queue job on the same worker.
            if policy_token is not None:
                reset_policy_id(policy_token)
            if run_token is not None:
                reset_run_id(run_token)
            if mut_token is not None:
                reset_scope_mutations(mut_token)
            if cid_token is not None:
                try:
                    from agntspace.activity.decisions import reset_correlation_id
                    reset_correlation_id(cid_token)
                except Exception:
                    pass
