"""Run one durable PentestGPT Agent trial from the command line."""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import time
import traceback
import uuid
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import unified_agent
from unified_agent import SandboxPolicy, UnifiedAgent

from . import __version__ as pentestgpt_agent_version
from .agents import (
    _EXECUTOR_TURN_BUDGETS,
    EXECUTOR_INSTRUCTIONS,
    EXECUTOR_SCHEMA,
    SUPERVISOR_INSTRUCTIONS,
    SUPERVISOR_SCHEMA,
    Executor,
    Supervisor,
)
from .identifiers import validate_opaque_id
from .loop import PentestLoop
from .memory import (
    AttemptStatus,
    MemoryKernel,
    RunSnapshot,
    RunSpec,
    RunStatus,
    validate_run_spec,
)
from .trace import EpisodeRunner, TraceStore


@dataclass(frozen=True)
class TrialConfig:
    run_id: str
    goal: str
    targets: tuple[str, ...]
    backend: str
    model: str | None
    runs_root: Path
    workspace_root: Path
    effort: str | None = None
    max_decisions: int = 20
    supervisor_max_turns: int = 4
    executor_max_turns: int = 12
    resume: bool = False


def _provider_environment(backend: str) -> dict[str, str]:
    if backend != "claude":
        return {}
    return {
        "CLAUDE_CODE_DISABLE_AUTO_MEMORY": "1",
        "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
        "DISABLE_AUTOUPDATER": "1",
    }


def _build_roles(config: TrialConfig, traces: TraceStore) -> tuple[Supervisor, Executor]:
    role_root = config.workspace_root / config.run_id
    provider_environment = _provider_environment(config.backend)
    # PentestGPT itself runs inside the deployment's isolation boundary. Both
    # roles intentionally receive every provider tool and unrestricted I/O.
    role_sandbox = SandboxPolicy.FULL_ACCESS
    supervisor = Supervisor(
        EpisodeRunner(
            UnifiedAgent(
                config.backend,
                workspace=role_root / "supervisor",
                model=config.model,
                effort=config.effort,
                sandbox=role_sandbox,
                instructions=SUPERVISOR_INSTRUCTIONS,
                extra_env=provider_environment,
            ),
            traces,
        ),
        max_turns=config.supervisor_max_turns,
    )
    executor = Executor(
        EpisodeRunner(
            UnifiedAgent(
                config.backend,
                workspace=role_root / "executor",
                model=config.model,
                effort=config.effort,
                sandbox=role_sandbox,
                instructions=EXECUTOR_INSTRUCTIONS,
                extra_env=provider_environment,
            ),
            traces,
        ),
        max_turns=config.executor_max_turns,
    )
    return supervisor, executor


async def run_trial(
    config: TrialConfig,
    *,
    supervisor: Supervisor | None = None,
    executor: Executor | None = None,
) -> dict[str, Any]:
    """Run or explicitly resume a trial and persist a summary even on failure."""
    if (supervisor is None) != (executor is None):
        raise ValueError("supervisor and executor must be supplied together")
    if not config.targets:
        raise ValueError("at least one authorized target is required")
    validate_opaque_id(config.run_id, label="run_id")
    if config.backend not in {"claude", "codex"}:
        raise ValueError("backend must be 'claude' or 'codex'")
    if config.model is not None and len(config.model) > 256:
        raise ValueError("model must be at most 256 characters")
    common_efforts = {"low", "medium", "high", "xhigh"}
    backend_efforts = {
        "claude": common_efforts | {"max"},
        "codex": common_efforts | {"none", "minimal"},
    }
    if config.effort is not None and config.effort not in backend_efforts[config.backend]:
        valid = ", ".join(sorted(backend_efforts[config.backend]))
        raise ValueError(f"effort for {config.backend} must be one of: {valid}")
    if not 1 <= config.max_decisions <= 1_000:
        raise ValueError("max_decisions must be between 1 and 1000")
    if not 1 <= config.supervisor_max_turns <= 100:
        raise ValueError("supervisor_max_turns must be between 1 and 100")
    if not 2 <= config.executor_max_turns <= 100:
        raise ValueError("executor_max_turns must be between 2 and 100")
    spec = RunSpec(
        run_id=config.run_id,
        goal=config.goal,
        allowed_targets=config.targets,
    )
    validate_run_spec(spec)

    run_dir = config.runs_root / config.run_id
    identity = _trial_identity(config)
    identity_path = run_dir / "trial-config.json"
    if config.resume:
        if not run_dir.is_dir() or not identity_path.is_file():
            raise ValueError(f"cannot resume missing trial {config.run_id!r}")
        persisted_identity = json.loads(identity_path.read_text(encoding="utf-8"))
        if persisted_identity != identity:
            raise ValueError(f"persisted trial {config.run_id!r} does not match requested config")
    else:
        run_dir.mkdir(parents=True, exist_ok=False, mode=0o700)
        _write_json(identity_path, identity)
    traces = TraceStore(config.runs_root)
    memory = MemoryKernel(run_dir / "state.sqlite3")
    if supervisor is None or executor is None:
        supervisor, executor = _build_roles(config, traces)

    started_at = time.time()
    try:
        snapshot = await PentestLoop(
            memory=memory,
            supervisor=supervisor,
            executor=executor,
            traces=traces,
            max_decisions=config.max_decisions,
        ).run(spec)
    except Exception as exc:
        summary = _trial_summary(
            config,
            _snapshot_if_present(memory, config.run_id),
            started_at,
            error=f"{type(exc).__name__}: {exc}",
        )
        _write_json(run_dir / "summary.json", summary)
        raise

    summary = _trial_summary(config, snapshot, started_at, error=None)
    _write_json(run_dir / "summary.json", summary)
    return summary


def _snapshot_if_present(memory: MemoryKernel, run_id: str) -> RunSnapshot | None:
    try:
        return memory.snapshot(run_id)
    except Exception:
        return None


def _trial_summary(
    config: TrialConfig,
    snapshot: RunSnapshot | None,
    started_at: float,
    *,
    error: str | None,
) -> dict[str, Any]:
    trace_root = config.runs_root / config.run_id / "traces"
    outputs = []
    if trace_root.exists():
        for path in sorted(trace_root.glob("*/output.json")):
            value = json.loads(path.read_text(encoding="utf-8"))
            if isinstance(value, dict):
                outputs.append(value)
    usage_fields = (
        "input_tokens",
        "cached_input_tokens",
        "output_tokens",
        "reasoning_output_tokens",
    )
    usage = {
        field: sum(int(output.get("usage", {}).get(field, 0) or 0) for output in outputs)
        for field in usage_fields
    }
    canonical_error = error
    if canonical_error is None and snapshot is not None and snapshot.status is RunStatus.FAILED:
        failed_attempts = [
            attempt for attempt in snapshot.attempts if attempt.status is AttemptStatus.ERROR
        ]
        if failed_attempts:
            failure = failed_attempts[-1]
            canonical_error = f"{failure.failure_kind}: {failure.failure_message}"
        else:
            run_failure = next(
                (
                    transition
                    for transition in reversed(snapshot.transitions)
                    if transition.kind in {"supervisor_failed", "decision_limit_reached"}
                ),
                None,
            )
            if run_failure is not None:
                canonical_error = (
                    f"{run_failure.detail.get('failure_kind')}: "
                    f"{run_failure.detail.get('failure_message')}"
                )
    return {
        "schema_version": 1,
        "runtime_policy_revision": 2,
        "pentestgpt_agent_version": pentestgpt_agent_version,
        "run_id": config.run_id,
        "goal": config.goal,
        "targets": list(config.targets),
        "backend": config.backend,
        "model": config.model,
        "effort": config.effort,
        "status": snapshot.status.value if snapshot is not None else "error",
        "revision": snapshot.revision if snapshot is not None else None,
        "tasks": [
            {
                "id": task.id,
                "kind": task.kind.value,
                "target": task.target,
                "status": task.status.value,
            }
            for task in (snapshot.tasks if snapshot is not None else ())
        ],
        "observations": [
            asdict(observation)
            for observation in (snapshot.observations if snapshot is not None else ())
        ],
        "attempts": [
            asdict(attempt) for attempt in (snapshot.attempts if snapshot is not None else ())
        ],
        "transitions": [
            asdict(transition)
            for transition in (snapshot.transitions if snapshot is not None else ())
        ],
        "episodes": len(outputs),
        "cost_usd": sum(float(output.get("cost_usd", 0) or 0) for output in outputs),
        "provider_duration_ms": sum(int(output.get("duration_ms", 0) or 0) for output in outputs),
        "usage": usage,
        "wall_duration_s": time.time() - started_at,
        "error": canonical_error,
    }


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)
    temporary.replace(path)


def _trial_identity(config: TrialConfig) -> dict[str, Any]:
    def digest(value: object) -> str:
        encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
        return hashlib.sha256(encoded).hexdigest()

    return {
        "schema_version": 2,
        "run_id": config.run_id,
        "goal": config.goal,
        "targets": list(config.targets),
        "backend": config.backend,
        "model": config.model,
        "effort": config.effort,
        "max_decisions": config.max_decisions,
        "max_attempts_per_task": 2,
        "supervisor_max_turns": config.supervisor_max_turns,
        "executor_max_turns": config.executor_max_turns,
        "executor_turn_budgets": {
            kind.value: budget for kind, budget in _EXECUTOR_TURN_BUDGETS.items()
        },
        "executor_result_turns_reserved": 1,
        "supervisor_contract_sha256": digest(
            {"instructions": SUPERVISOR_INSTRUCTIONS, "schema": SUPERVISOR_SCHEMA}
        ),
        "executor_contract_sha256": digest(
            {"instructions": EXECUTOR_INSTRUCTIONS, "schema": EXECUTOR_SCHEMA}
        ),
        "supervisor_sandbox": SandboxPolicy.FULL_ACCESS.value,
        "executor_sandbox": SandboxPolicy.FULL_ACCESS.value,
        "provider_environment": _provider_environment(config.backend),
        "unified_agent_version": unified_agent.__version__,
    }


def _parse_args(argv: list[str] | None = None) -> TrialConfig:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--goal", required=True)
    parser.add_argument("--target", action="append", dest="targets", required=True)
    parser.add_argument("--backend", choices=("claude", "codex"), default="claude")
    parser.add_argument("--model")
    parser.add_argument("--effort")
    parser.add_argument("--run-id", default=f"trial-{uuid.uuid4().hex[:12]}")
    parser.add_argument("--runs-root", type=Path, default=Path("runs"))
    parser.add_argument("--workspace-root", type=Path, default=Path("agent-workspaces"))
    parser.add_argument("--max-decisions", type=int, default=20)
    parser.add_argument("--supervisor-max-turns", type=int, default=4)
    parser.add_argument("--executor-max-turns", type=int, default=12)
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args(argv)
    return TrialConfig(
        run_id=args.run_id,
        goal=args.goal,
        targets=tuple(args.targets),
        backend=args.backend,
        model=args.model,
        runs_root=args.runs_root,
        workspace_root=args.workspace_root,
        effort=args.effort,
        max_decisions=args.max_decisions,
        supervisor_max_turns=args.supervisor_max_turns,
        executor_max_turns=args.executor_max_turns,
        resume=args.resume,
    )


def main(argv: list[str] | None = None) -> int:
    config = _parse_args(argv)
    try:
        summary = asyncio.run(run_trial(config))
    except Exception:
        traceback.print_exc()
        return 1
    print("TRIAL_RESULT=" + json.dumps(summary, sort_keys=True))
    return 0 if summary.get("status") == RunStatus.COMPLETED.value else 1


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())


__all__ = ["TrialConfig", "main", "run_trial"]
