"""The deterministic one-task-at-a-time penetration-testing loop."""

from __future__ import annotations

from .agents import (
    AgentContractError,
    Executor,
    Supervisor,
    parse_execution_delta,
    parse_supervisor_decision,
    recover_execution_deltas,
)
from .execution import ExecutionValidationError, compile_execution
from .memory import (
    AttemptStatus,
    MemoryKernel,
    RunSnapshot,
    RunSpec,
    RunStatus,
    TaskLease,
)
from .plan import PlanValidationError, TaskStatus, ValidPlan, compile_plan
from .trace import EpisodeTrace, TraceStore, failure_detail, has_action_receipts


class PentestLoop:
    def __init__(
        self,
        *,
        memory: MemoryKernel,
        supervisor: Supervisor,
        executor: Executor,
        traces: TraceStore,
        max_decisions: int = 20,
        max_supervisor_attempts: int = 2,
    ) -> None:
        self.memory = memory
        self.supervisor = supervisor
        self.executor = executor
        self.traces = traces
        self.max_decisions = max_decisions
        self.max_supervisor_attempts = max_supervisor_attempts

    async def run(self, spec: RunSpec) -> RunSnapshot:
        snapshot = self.memory.open_run(spec)
        if snapshot.status is not RunStatus.RUNNING:
            return snapshot

        lease = self._active_lease(snapshot)
        decisions = sum(transition.kind == "plan_committed" for transition in snapshot.transitions)
        while snapshot.status is RunStatus.RUNNING:
            if lease is None:
                if decisions >= self.max_decisions:
                    return self.memory.commit_run_failure(
                        snapshot.run_id,
                        snapshot.revision,
                        failure_kind="decision_limit",
                        failure_message=(
                            f"run reached the limit of {self.max_decisions} Supervisor decisions"
                        ),
                        transition_kind="decision_limit_reached",
                    )
                try:
                    plan = await self._supervisor_plan(snapshot)
                except AgentContractError as exc:
                    return self.memory.commit_run_failure(
                        snapshot.run_id,
                        snapshot.revision,
                        failure_kind="supervisor_contract",
                        failure_message=str(exc),
                    )
                decisions += 1
                plan_commit = self.memory.commit_plan(plan)
                snapshot = self.memory.snapshot(spec.run_id)
                if plan_commit.lease is None:
                    return snapshot
                lease = plan_commit.lease

            snapshot, lease = await self._drive_lease(snapshot, lease)

        return snapshot

    async def _supervisor_plan(self, snapshot: RunSnapshot) -> ValidPlan:
        """Recover or retry one decision, validating before it can mutate canonical state."""
        last_error = "Supervisor did not return a valid plan"
        feedback: str | None = None
        for attempt_number in range(1, self.max_supervisor_attempts + 1):
            episode_id = self._supervisor_episode(snapshot.revision, attempt_number)
            try:
                if self.traces.exists(snapshot.run_id, episode_id):
                    trace = self.traces.recover_terminal_output(snapshot.run_id, episode_id)
                    if (
                        trace.truncated_tail
                        or trace.output is None
                        or trace.output.get("success") is not True
                    ):
                        raise AgentContractError(
                            f"persisted Supervisor episode is not recoverable: {episode_id}"
                        )
                    if (
                        trace.input.get("run_id") != snapshot.run_id
                        or trace.input.get("episode_id") != episode_id
                        or trace.input.get("role") != "supervisor"
                        or trace.input.get("state_revision") != snapshot.revision
                    ):
                        raise AgentContractError(
                            f"persisted Supervisor episode has stale identity: {episode_id}"
                        )
                    decision = parse_supervisor_decision(trace.output.get("structured_output"))
                else:
                    decision = await self.supervisor.decide(
                        snapshot,
                        episode_id=episode_id,
                        feedback=feedback,
                    )
                return compile_plan(decision, snapshot)
            except (AgentContractError, PlanValidationError, OSError, ValueError) as exc:
                last_error = str(exc)
                feedback = last_error
        raise AgentContractError(
            f"Supervisor failed after {self.max_supervisor_attempts} attempts: {last_error}"
        )

    async def _drive_lease(
        self,
        snapshot: RunSnapshot,
        lease: TaskLease,
    ) -> tuple[RunSnapshot, TaskLease | None]:
        """Resolve one lease without ever replaying an existing trace episode."""
        executor_episode = self._executor_episode(snapshot, lease)
        if self.traces.exists(lease.run_id, executor_episode):
            try:
                trace = self.traces.recover_terminal_output(lease.run_id, executor_episode)
            except FileNotFoundError:
                if self.traces.initialization_incomplete(lease.run_id, executor_episode):
                    return self._commit_failure(
                        lease,
                        executor_episode,
                        failure_kind="interrupted",
                        failure_message="Executor episode initialization was interrupted",
                        retryable=True,
                    )
                return self._commit_failure(
                    lease,
                    executor_episode,
                    failure_kind="trace_corrupt",
                    failure_message="Executor trace is missing a required journal file",
                    retryable=False,
                )
            except (OSError, ValueError) as exc:
                return self._commit_failure(
                    lease,
                    executor_episode,
                    failure_kind="trace_corrupt",
                    failure_message=f"Executor trace could not be decoded: {type(exc).__name__}: {exc}",
                    retryable=False,
                )
            return self._settle_trace(snapshot, lease, executor_episode, trace)

        try:
            delta = await self.executor.execute(
                snapshot,
                lease,
                episode_id=executor_episode,
            )
            trace = self.traces.load(lease.run_id, executor_episode)
            execution = compile_execution(
                delta,
                lease,
                trace,
                prior_observations=snapshot.observations,
            )
        except (AgentContractError, ExecutionValidationError) as exc:
            trace = self.traces.load(lease.run_id, executor_episode)
            return self._settle_trace(
                snapshot,
                lease,
                executor_episode,
                trace,
                fallback=str(exc),
            )
        return self.memory.commit_execution(execution), None

    def _settle_trace(
        self,
        snapshot: RunSnapshot,
        lease: TaskLease,
        episode_id: str,
        trace: EpisodeTrace,
        *,
        fallback: str = "Executor episode failed",
    ) -> tuple[RunSnapshot, TaskLease | None]:
        if trace.truncated_tail:
            failure_kind = "trace_corrupt"
            failure_message = "Executor trace has a truncated event record"
        elif trace.output is None:
            failure_kind = "interrupted"
            failure_message = "Executor episode ended without a durable result"
        elif trace.output.get("success") is True:
            try:
                delta = parse_execution_delta(trace.output.get("structured_output"))
                execution = compile_execution(
                    delta,
                    lease,
                    trace,
                    prior_observations=snapshot.observations,
                )
            except (AgentContractError, ExecutionValidationError) as exc:
                failure_kind = "validation"
                failure_message = str(exc)
            else:
                return self.memory.commit_execution(execution), None
        else:
            failure_kind, failure_message = failure_detail(trace, fallback)
            if failure_kind in {"max_structured_output_retries", "max_turns"}:
                try:
                    candidates = recover_execution_deltas(trace)
                except AgentContractError:
                    pass
                else:
                    for delta in candidates:
                        try:
                            execution = compile_execution(
                                delta,
                                lease,
                                trace,
                                recover_structured_transport=True,
                                allow_evidence_fallback=False,
                                prior_observations=snapshot.observations,
                            )
                        except ExecutionValidationError:
                            continue
                        return self.memory.commit_execution(execution), None
                    try:
                        execution = compile_execution(
                            candidates[0],
                            lease,
                            trace,
                            recover_structured_transport=True,
                            prior_observations=snapshot.observations,
                        )
                    except ExecutionValidationError:
                        pass
                    else:
                        return self.memory.commit_execution(execution), None

        retryable = (
            failure_kind in {"provider", "validation", "interrupted"}
            and not trace.truncated_tail
            and not has_action_receipts(trace)
        )
        return self._commit_failure(
            lease,
            episode_id,
            failure_kind=failure_kind,
            failure_message=failure_message,
            retryable=retryable,
        )

    def _commit_failure(
        self,
        lease: TaskLease,
        episode_id: str,
        *,
        failure_kind: str,
        failure_message: str,
        retryable: bool,
    ) -> tuple[RunSnapshot, TaskLease | None]:
        attempt_commit = self.memory.commit_attempt_failure(
            lease,
            trace_episode_id=episode_id,
            failure_kind=failure_kind,
            failure_message=failure_message,
            retryable=retryable,
        )
        return attempt_commit.snapshot, attempt_commit.retry_lease

    @staticmethod
    def _active_lease(snapshot: RunSnapshot) -> TaskLease | None:
        active_attempts = [
            attempt for attempt in snapshot.attempts if attempt.status is AttemptStatus.ACTIVE
        ]
        active_tasks = [task for task in snapshot.tasks if task.status is TaskStatus.ACTIVE]
        if not active_attempts and not active_tasks:
            return None
        if len(active_attempts) != 1 or len(active_tasks) != 1:
            raise RuntimeError("canonical state must contain one active task and attempt")
        attempt = active_attempts[0]
        if attempt.task_id != active_tasks[0].id:
            raise RuntimeError("active attempt does not belong to the active task")
        if attempt.started_revision != snapshot.revision:
            raise RuntimeError("active attempt revision does not match the run revision")
        return TaskLease(
            run_id=snapshot.run_id,
            task_id=attempt.task_id,
            attempt_id=attempt.id,
            revision=attempt.started_revision,
            trace_episode_id=attempt.trace_episode_id,
        )

    @staticmethod
    def _executor_episode(snapshot: RunSnapshot, lease: TaskLease) -> str:
        if lease.trace_episode_id is not None:
            return lease.trace_episode_id
        attempt = next(attempt for attempt in snapshot.attempts if attempt.id == lease.attempt_id)
        return attempt.trace_episode_id or f"executor-{lease.attempt_id}"

    @staticmethod
    def _supervisor_episode(revision: int, attempt_number: int) -> str:
        base = f"supervisor-r{revision}"
        return base if attempt_number == 1 else f"{base}-a{attempt_number}"


__all__ = ["PentestLoop"]
