"""Run one UnifiedAgent episode while durably recording what it saw and returned."""

from __future__ import annotations

import hashlib
import json
import os
import time
from dataclasses import asdict, dataclass
from enum import StrEnum
from pathlib import Path
from typing import Any

import unified_agent
from unified_agent import (
    AssistantText,
    CommandOutputDelta,
    CommandRun,
    CommandStarted,
    FileChanged,
    RawEvent,
    Reasoning,
    SessionStarted,
    Task,
    ToolCall,
    ToolProgress,
    ToolResult,
    TurnCompleted,
    UnifiedAgent,
    UnifiedResult,
    collect,
)
from unified_agent.events import AgentEvent

from .identifiers import validate_opaque_id


class AgentRole(StrEnum):
    SUPERVISOR = "supervisor"
    EXECUTOR = "executor"
    VERIFIER = "verifier"


@dataclass(frozen=True)
class EpisodeInput:
    run_id: str
    episode_id: str
    role: AgentRole
    state_revision: int
    task: Task | str
    output_schema: dict[str, Any] | None = None
    max_turns: int | None = None
    task_id: str | None = None
    attempt_id: str | None = None


@dataclass(frozen=True)
class EpisodeTrace:
    input: dict[str, Any]
    events: tuple[dict[str, Any], ...]
    output: dict[str, Any] | None
    truncated_tail: bool


class TraceStore:
    """Filesystem trace journal. EpisodeRunner is its only writer."""

    def __init__(self, root: str | Path) -> None:
        self.root = Path(root)

    def load(self, run_id: str, episode_id: str) -> EpisodeTrace:
        episode_dir = self._episode_dir(run_id, episode_id)
        input_data = json.loads((episode_dir / "input.json").read_text(encoding="utf-8"))
        if not isinstance(input_data, dict):
            raise ValueError("trace input must be a JSON object")
        output_path = episode_dir / "output.json"
        output_data: dict[str, Any] | None = None
        if output_path.exists():
            candidate_output = json.loads(output_path.read_text(encoding="utf-8"))
            if not isinstance(candidate_output, dict):
                raise ValueError("trace output must be a JSON object")
            output_data = candidate_output
        lines = (episode_dir / "events.jsonl").read_text(encoding="utf-8").splitlines()
        events: list[dict[str, Any]] = []
        truncated_tail = False
        for index, line in enumerate(lines):
            if not line.strip():
                continue
            try:
                event = json.loads(line)
            except json.JSONDecodeError:
                if index != len(lines) - 1:
                    raise
                truncated_tail = True
                break
            if not isinstance(event, dict):
                raise ValueError("trace event must be a JSON object")
            events.append(event)
        return EpisodeTrace(
            input=input_data,
            events=tuple(events),
            output=output_data,
            truncated_tail=truncated_tail,
        )

    def exists(self, run_id: str, episode_id: str) -> bool:
        return self._episode_dir(run_id, episode_id).exists()

    def initialization_incomplete(self, run_id: str, episode_id: str) -> bool:
        """Return whether the runner crashed before its append-only event journal existed."""
        episode_dir = self._episode_dir(run_id, episode_id)
        return episode_dir.is_dir() and not (episode_dir / "events.jsonl").exists()

    def recover_terminal_output(self, run_id: str, episode_id: str) -> EpisodeTrace:
        """Finish the tiny crash window after a durable terminal event was journaled."""
        trace = self.load(run_id, episode_id)
        if trace.output is not None or trace.truncated_tail or not trace.events:
            return trace
        terminal = trace.events[-1]
        if terminal.get("type") != "turn_completed" or not isinstance(
            terminal.get("success"), bool
        ):
            return trace
        output = {
            "backend": trace.input.get("provider", {}).get("backend"),
            "success": terminal["success"],
            "text": terminal.get("final_text"),
            "structured_output": terminal.get("structured_output"),
            "usage": terminal.get("usage"),
            "cost_usd": terminal.get("cost_usd"),
            "session_id": terminal.get("session_id"),
            "duration_ms": terminal.get("duration_ms"),
            "error": terminal.get("error"),
            "closed_at": time.time(),
            "recovered_from_terminal_event": True,
        }
        self._finish(self._episode_dir(run_id, episode_id), output)
        return EpisodeTrace(
            input=trace.input,
            events=trace.events,
            output=output,
            truncated_tail=False,
        )

    def _begin(self, run_id: str, episode_id: str, input_data: dict[str, Any]) -> Path:
        episode_dir = self._episode_dir(run_id, episode_id)
        episode_dir.mkdir(parents=True, exist_ok=False)
        self._write_json(episode_dir / "input.json", input_data)
        events_path = episode_dir / "events.jsonl"
        events_path.touch(mode=0o600)
        return episode_dir

    def _append(self, episode_dir: Path, event: dict[str, Any]) -> None:
        with (episode_dir / "events.jsonl").open("a", encoding="utf-8") as stream:
            stream.write(json.dumps(event, sort_keys=True, separators=(",", ":")) + "\n")
            stream.flush()
            os.fsync(stream.fileno())

    def _finish(self, episode_dir: Path, output: dict[str, Any]) -> None:
        self._write_json(episode_dir / "output.json", output)

    def _episode_dir(self, run_id: str, episode_id: str) -> Path:
        validate_opaque_id(run_id, label="run_id")
        validate_opaque_id(episode_id, label="episode_id")
        return self.root / run_id / "traces" / episode_id

    @staticmethod
    def _write_json(path: Path, value: dict[str, Any]) -> None:
        temporary = path.with_suffix(path.suffix + ".tmp")
        temporary.write_text(
            json.dumps(value, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        temporary.chmod(0o600)
        os.replace(temporary, path)


class EpisodeRunner:
    """Deep seam around provider invocation and diagnostic trace persistence."""

    def __init__(self, agent: UnifiedAgent, traces: TraceStore) -> None:
        self.agent = agent
        self.traces = traces

    async def run(self, episode: EpisodeInput) -> UnifiedResult:
        rendered_prompt = Task.coerce(episode.task).render(self.agent.name)
        opened_at = time.time()
        provider = {
            "backend": self.agent.name,
            "model": self.agent.model,
            "effort": self.agent.effort,
            "sandbox": self.agent.sandbox.value,
            "unified_agent_version": unified_agent.__version__,
        }
        instructions = self.agent.instructions
        episode_dir = self.traces._begin(
            episode.run_id,
            episode.episode_id,
            {
                "schema_version": 1,
                "run_id": episode.run_id,
                "episode_id": episode.episode_id,
                "role": episode.role.value,
                "state_revision": episode.state_revision,
                "task_id": episode.task_id,
                "attempt_id": episode.attempt_id,
                "provider": provider,
                "signatures": {
                    "prompt_sha256": _sha256_text(rendered_prompt),
                    "instructions_sha256": _sha256_text(instructions or ""),
                    "output_schema_sha256": _sha256_json(episode.output_schema),
                    "provider_sha256": _sha256_json(provider),
                },
                "instructions": instructions,
                "rendered_prompt": rendered_prompt,
                "output_schema": episode.output_schema,
                "max_turns": episode.max_turns,
                "resume": None,
                "opened_at": opened_at,
            },
        )

        events: list[AgentEvent] = []
        trace_sequence = 0
        monitor_run_id = f"{episode.run_id}--{episode.episode_id}"
        monitor = getattr(self.agent, "monitor", None)
        if monitor is None:
            stream = self.agent.stream(
                rendered_prompt,
                output_schema=episode.output_schema,
                resume=None,
                max_turns=episode.max_turns,
            ).__aiter__()
        else:
            stream = self.agent.stream(
                rendered_prompt,
                output_schema=episode.output_schema,
                resume=None,
                max_turns=episode.max_turns,
                run_id=monitor_run_id,
                monitor_metadata={
                    "project": "PentestGPT",
                    "pentest_run_id": episode.run_id,
                    "episode_id": episode.episode_id,
                    "role": episode.role.value,
                    "state_revision": episode.state_revision,
                    "task_id": episode.task_id,
                    "attempt_id": episode.attempt_id,
                },
            ).__aiter__()
        provider_error: str | None = None
        while True:
            try:
                event = await anext(stream)
            except StopAsyncIteration:
                break
            except Exception as exc:  # provider failures become a durable failed result
                provider_error = f"{type(exc).__name__}: {exc}"
                break
            trace_sequence += 1
            event_record = _event_record(event)
            self.traces._append(
                episode_dir,
                {
                    "sequence": trace_sequence,
                    "recorded_at": time.time(),
                    **event_record,
                },
            )
            if not isinstance(event, (Reasoning, RawEvent)):
                events.append(event)

        if provider_error is None and not any(isinstance(event, TurnCompleted) for event in events):
            provider_error = "episode stream ended without TurnCompleted"
        result = collect(
            self.agent.name,
            events,
            error=provider_error,
            monitor_run_id=monitor_run_id if monitor is not None else None,
        )
        self.traces._finish(episode_dir, _result_record(result))
        return result


def failure_detail(trace: EpisodeTrace, fallback: str) -> tuple[str, str]:
    """Prefer the provider's typed terminal failure over a later transport error."""
    for event in reversed(trace.events):
        if event.get("type") != "turn_completed" or event.get("success") is not False:
            continue
        safety_message = _claude_safety_rejection_message(trace, event)
        if safety_message is not None:
            return "provider_safety_block", safety_message
        stop_reason = event.get("stop_reason")
        kind = (
            "max_turns"
            if stop_reason == "error_max_turns"
            else str(stop_reason or "provider").removeprefix("error_")
        )
        message = event.get("error") or stop_reason or fallback
        return kind, str(message)
    output_error = trace.output.get("error") if trace.output else None
    return "provider", str(output_error or fallback)


def _claude_safety_rejection_message(
    trace: EpisodeTrace,
    terminal: dict[str, Any],
) -> str | None:
    if terminal.get("stop_reason") != "success" or terminal.get("error") != "success":
        return None

    raw_error_kinds = [
        str(event.get("kind"))
        for event in trace.events
        if event.get("type") == "raw_event"
        and event.get("backend") == "claude"
        and str(event.get("kind", "")).startswith("assistant_error:")
    ]
    candidate_messages = [
        terminal.get("final_text"),
        *(
            event.get("text")
            for event in reversed(trace.events)
            if event.get("type") == "assistant_text"
        ),
        trace.output.get("text") if trace.output else None,
    ]
    for candidate in candidate_messages:
        if not isinstance(candidate, str):
            continue
        normalized = candidate.casefold()
        if "safeguards flagged this message" in normalized or (
            "cybersecurity topic" in normalized and "flagged" in normalized
        ):
            return candidate
    for kind in raw_error_kinds:
        if "safety" in kind.casefold() or "safeguard" in kind.casefold():
            return f"Claude provider rejected the episode ({kind})"
    return None


def has_action_receipts(trace: EpisodeTrace) -> bool:
    """Return whether replaying the attempt could duplicate an external side effect."""
    for event in trace.events:
        event_type = event.get("type")
        if event_type in {"command_started", "command_run", "file_changed"}:
            return True
        if event_type == "tool_call" and event.get("name") != "StructuredOutput":
            return True
    return False


def is_grounding_receipt(
    event: dict[str, Any],
    structured_output_call_ids: set[object],
) -> bool:
    """Return whether an action result is eligible as exact canonical evidence."""
    if event.get("type") == "command_run":
        exit_code = event.get("exit_code")
        return isinstance(exit_code, int) and not isinstance(exit_code, bool)
    if event.get("type") != "tool_result":
        return False
    return (
        event.get("call_id") not in structured_output_call_ids
        and not bool(event.get("is_error", False))
        and event.get("exit_code") in (None, 0)
    )


def _event_record(event: AgentEvent) -> dict[str, Any]:
    if isinstance(event, Reasoning):
        return {"type": "reasoning", "content_omitted": True}
    if isinstance(event, RawEvent):
        return {
            "type": "raw_event",
            "backend": event.backend,
            "kind": event.kind,
            "data_omitted": True,
        }
    if isinstance(event, SessionStarted):
        return {"type": "session_started", **asdict(event)}
    if isinstance(event, ToolCall):
        return {"type": "tool_call", **asdict(event)}
    if isinstance(event, ToolProgress):
        return {"type": "tool_progress", **asdict(event)}
    if isinstance(event, ToolResult):
        return {"type": "tool_result", **asdict(event)}
    if isinstance(event, CommandStarted):
        return {"type": "command_started", **asdict(event)}
    if isinstance(event, CommandOutputDelta):
        return {"type": "command_output_delta", **asdict(event)}
    if isinstance(event, CommandRun):
        return {"type": "command_run", **asdict(event)}
    if isinstance(event, FileChanged):
        return {"type": "file_changed", **asdict(event)}
    if isinstance(event, AssistantText):
        return {"type": "assistant_text", "text": event.text}
    if isinstance(event, TurnCompleted):
        return {"type": "turn_completed", **asdict(event)}
    raise TypeError(f"unsupported UnifiedAgent event: {type(event).__name__}")


def _result_record(result: UnifiedResult) -> dict[str, Any]:
    return {
        "backend": result.backend,
        "success": result.success,
        "text": result.text,
        "structured_output": result.structured_output,
        "usage": asdict(result.usage),
        "cost_usd": result.cost_usd,
        "session_id": result.session_id,
        "duration_ms": result.duration_ms,
        "error": result.error,
        "closed_at": time.time(),
    }


def _sha256_text(value: str) -> str:
    return hashlib.sha256(value.encode()).hexdigest()


def _sha256_json(value: object) -> str:
    canonical = json.dumps(value, sort_keys=True, separators=(",", ":"))
    return _sha256_text(canonical)


__all__ = [
    "AgentRole",
    "EpisodeInput",
    "EpisodeRunner",
    "EpisodeTrace",
    "TraceStore",
    "failure_detail",
    "has_action_receipts",
    "is_grounding_receipt",
]
