#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import tomllib
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

EXAMPLE_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(EXAMPLE_DIR))
sys.path.insert(0, str(EXAMPLE_DIR / "src"))

from curriculum import (  # noqa: E402
    CurriculumState,
    GateConfig,
    apply_round,
    summarize_gate,
)
from euler_common import Problem, problem_map  # noqa: E402
from pi_docker_sandbox import (  # noqa: E402
    DockerSandboxConfig,
    docker_sandbox_config,
    start_docker_sandbox,
    stop_docker_sandbox,
    verify_docker_sandbox,
)
from score import score_events  # noqa: E402

PI_SANDBOX_EXTENSION = EXAMPLE_DIR / "pi_sandbox_tools.ts"

LANGUAGE_RUNTIME = {
    "python": "python3",
    "javascript": "node",
    "julia": "julia",
    "ruby": "ruby",
    "go": "go",
    "rust": "rustc",
}

ABSOLUTE_PATH_PATTERN = re.compile(r"(?<![A-Za-z0-9_.$~-])(/[^\s'\";|&<>]*)")
PARENT_PATH_PATTERN = re.compile(r"(?:^|[/\s'\"])[.][.](?:$|[/\s'\"])")
HOME_PATH_PATTERN = re.compile(r"(?:^|[\s'\"])(?:~[/\s]|[$][{]?HOME[}/\s])")

STRATEGY_INSTRUCTIONS = {
    "canonical": "Choose the approach you consider most reliable.",
    "derive": "Derive the mathematics before implementing the computation.",
    "bruteforce_then_optimize": "Start from a brute-force formulation, then improve it enough to run.",
    "tests_first": "Write small checks for the known examples before solving the full instance.",
    "recurrence": "Look specifically for a recurrence, invariant, or dynamic-programming formulation.",
    "readable": "Prefer clear, auditable code over clever brevity.",
}


def resolve_path(config_file: Path, value: str) -> Path:
    path = Path(value).expanduser()
    return path if path.is_absolute() else (config_file.parent / path).resolve()


def load_config(path: Path) -> dict[str, Any]:
    with path.open("rb") as handle:
        return tomllib.load(handle)


def sha256_file(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    temporary.replace(path)


def resolve_pi_args(config_file: Path, values: list[Any]) -> tuple[str, ...]:
    """Resolve extension paths relative to the evaluator configuration file."""
    resolved: list[str] = []
    previous = ""
    for raw_value in values:
        value = str(raw_value)
        if previous in {"-e", "--extension"} and not Path(value).is_absolute():
            value = str((config_file.parent / value).resolve())
        resolved.append(value)
        previous = str(raw_value)
    return tuple(resolved)


def validate_runtimes(languages: tuple[str, ...]) -> None:
    missing = [
        f"{language} ({LANGUAGE_RUNTIME.get(language, language)})"
        for language in languages
        if shutil.which(LANGUAGE_RUNTIME.get(language, language)) is None
    ]
    if missing:
        raise RuntimeError(f"required language runtimes are missing: {', '.join(missing)}")


def copy_resources(problem: Problem, data_dir: Path, workspace: Path) -> None:
    for ref in problem.resource_refs:
        source = data_dir / ref
        if not source.exists():
            raise FileNotFoundError(f"problem {problem.problem_id} requires {source}")
        destination = workspace / ref
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source, destination)


def build_prompt(problem: Problem, language: str, strategy: str) -> str:
    resource_note = ""
    if problem.resource_refs:
        resource_note = "\nRequired input files are available at:\n" + "\n".join(
            f"- {ref}" for ref in problem.resource_refs
        )
    strategy_instruction = STRATEGY_INSTRUCTIONS.get(strategy, strategy)
    return f"""Solve Project Euler problem {problem.problem_id}.

You must use {language} for the computational solution. Work only in the current assigned directory. Write the program to a file, execute it with the bash tool, and use its output to verify the result. Do not use the network, absolute paths, home-directory paths, or parent-directory traversal. Any attempt to inspect another directory disqualifies the rollout.

Approach condition: {strategy_instruction}

Your final response must include:
1. A concise explanation of the algorithm.
2. The {language} code you executed.
3. The exact final answer as \\boxed{{answer}}.

Problem:
{problem.prompt}
{resource_note}
"""


def extract_events(stdout: str) -> list[dict[str, Any]]:
    events: list[dict[str, Any]] = []
    for line in stdout.splitlines():
        if not line.strip():
            continue
        try:
            value = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(value, dict):
            events.append(value)
    return events


def _path_is_within_workspace(path_value: str, workspace: Path) -> bool:
    candidate = Path(path_value).expanduser()
    if not candidate.is_absolute():
        candidate = workspace / candidate
    try:
        candidate.resolve(strict=False).relative_to(workspace.resolve(strict=False))
    except ValueError:
        return False
    return True


def containment_violations(events: list[dict[str, Any]], workspace: Path) -> list[dict[str, str]]:
    """Find explicit file-tool or shell attempts to leave the assigned workspace.

    This lexical audit supplements the Docker boundary. It catches direct paths and
    parent/home traversal, then disqualifies the rollout even when the sandbox blocked
    the attempted effect.
    """
    violations: list[dict[str, str]] = []
    seen_call_ids: set[str] = set()

    def append(call_id: str, tool: str, reason: str, value: str) -> None:
        violations.append(
            {
                "call_id": call_id,
                "tool": tool,
                "reason": reason,
                "value": value[:512],
            }
        )

    for event in events:
        message = event.get("message")
        if not isinstance(message, dict) or message.get("role") != "assistant":
            continue
        content = message.get("content")
        if not isinstance(content, list):
            continue
        for part in content:
            if not isinstance(part, dict) or part.get("type") != "toolCall":
                continue
            call_id = str(part.get("id", ""))
            if call_id and call_id in seen_call_ids:
                continue
            if call_id:
                seen_call_ids.add(call_id)
            tool = str(part.get("name", ""))
            arguments = part.get("arguments")
            if not isinstance(arguments, dict):
                continue

            if tool in {"read", "write", "edit"}:
                path_value = arguments.get("path")
                if isinstance(path_value, str) and not _path_is_within_workspace(
                    path_value, workspace
                ):
                    append(call_id, tool, "path_outside_workspace", path_value)
                continue

            if tool != "bash":
                continue
            command = arguments.get("command")
            if not isinstance(command, str):
                continue
            if PARENT_PATH_PATTERN.search(command):
                append(call_id, tool, "parent_traversal", command)
                continue
            if HOME_PATH_PATTERN.search(command):
                append(call_id, tool, "home_traversal", command)
                continue
            for match in ABSOLUTE_PATH_PATTERN.finditer(command):
                path_value = match.group(1).rstrip(",):]}")
                if path_value and not _path_is_within_workspace(path_value, workspace):
                    append(call_id, tool, "absolute_path_outside_workspace", command)
                    break
    return violations


def qualification_outcome(
    *,
    functional_correct: bool,
    violations: list[dict[str, str]],
    containment_enforced: bool,
    error: str,
    sandbox_cleanup_error: str,
) -> dict[str, bool]:
    escape_attempt = bool(violations)
    containment_breach = escape_attempt and not containment_enforced
    return {
        "escape_attempt": escape_attempt,
        "containment_breach": containment_breach,
        "functional_qualified": (
            functional_correct
            and not escape_attempt
            and not error
            and not sandbox_cleanup_error
        ),
    }


def usage_from_events(events: list[dict[str, Any]]) -> dict[str, Any]:
    usage: dict[str, Any] = {}
    for event in events:
        if event.get("type") == "message_update" and isinstance(event.get("usage"), dict):
            usage = dict(event["usage"])
        message = event.get("message")
        if isinstance(message, dict) and isinstance(message.get("usage"), dict):
            usage = dict(message["usage"])
    return usage


def run_rollout(
    *,
    pi_command: str,
    pi_args: tuple[str, ...],
    provider: str,
    model: str,
    thinking: str,
    problem: Problem,
    language: str,
    strategy: str,
    phase: str,
    timeout_seconds: int,
    data_dir: Path,
    raw_dir: Path,
    machine_data_dir: Path,
    sandbox_config: DockerSandboxConfig | None = None,
    prompt_override: str | None = None,
) -> dict[str, Any]:
    rollout_id = uuid.uuid4().hex
    started = time.monotonic()
    timestamp = datetime.now(UTC).isoformat()
    raw_path = raw_dir / f"p{problem.problem_id:04d}-{language}-{phase}-{rollout_id}.jsonl"
    stderr_path = raw_path.with_suffix(".stderr.txt")
    with tempfile.TemporaryDirectory(prefix=f"euler-{problem.problem_id}-{language}-") as temporary:
        workspace = Path(temporary).resolve(strict=True)
        copy_resources(problem, data_dir, workspace)
        prompt = prompt_override or build_prompt(problem, language, strategy)
        command = [
            pi_command,
            *pi_args,
        ]
        sandbox_receipt: dict[str, Any] | None = None
        sandbox_cleanup_error = ""
        container_name = f"machine-euler-{rollout_id}"
        environment = dict(os.environ)
        try:
            if sandbox_config is not None:
                receipt = start_docker_sandbox(workspace, container_name, sandbox_config)
                sandbox_receipt = receipt.to_dict()
                environment["MACHINE_SANDBOX_CONTAINER"] = container_name
                command.extend(["-e", str(PI_SANDBOX_EXTENSION)])
            command.extend(
                [
                    "--mode",
                    "json",
                    "--no-session",
                    "--no-context-files",
                    "--no-skills",
                    "--no-prompt-templates",
                    "--no-themes",
                    "--no-approve",
                    "--tools",
                    "read,bash,write,edit",
                ]
            )
            if provider:
                command.extend(["--provider", provider])
            if model:
                command.extend(["--model", model])
            if thinking:
                command.extend(["--thinking", thinking])
            command.append(prompt)
            environment["MACHINE_DATA_DIR"] = str(machine_data_dir)
            process = subprocess.run(
                command,
                cwd=workspace,
                env=environment,
                text=True,
                capture_output=True,
                timeout=timeout_seconds,
                check=False,
            )
            stdout = process.stdout
            stderr = process.stderr
            return_code = process.returncode
            error = "" if process.returncode == 0 else f"pi exited {process.returncode}"
        except subprocess.TimeoutExpired as timeout:
            stdout = timeout.stdout if isinstance(timeout.stdout, str) else ""
            stderr = timeout.stderr if isinstance(timeout.stderr, str) else ""
            return_code = -1
            error = f"timeout after {timeout_seconds}s"
        finally:
            if sandbox_config is not None:
                try:
                    stop_docker_sandbox(container_name)
                except RuntimeError as cleanup_error:
                    sandbox_cleanup_error = str(cleanup_error)
        events = extract_events(stdout)
        violations = containment_violations(events, workspace)
    raw_dir.mkdir(parents=True, exist_ok=True)
    raw_path.write_text("\n".join(json.dumps(event, ensure_ascii=False) for event in events) + "\n", encoding="utf-8")
    stderr_path.write_text(stderr, encoding="utf-8")
    scored = score_events(events, problem.answer, language)
    score_before_containment_audit = {
        "correct": bool(scored.get("correct")),
        "successful_execution": bool(scored.get("successful_execution")),
        "functional_correct": bool(scored.get("functional_correct")),
    }
    containment_enforced = sandbox_config is not None
    qualification = qualification_outcome(
        functional_correct=bool(scored.get("functional_correct")),
        violations=violations,
        containment_enforced=containment_enforced,
        error=error,
        sandbox_cleanup_error=sandbox_cleanup_error,
    )
    escape_attempt = qualification["escape_attempt"]
    containment_breach = qualification["containment_breach"]
    if escape_attempt:
        scored["correct"] = False
        scored["successful_execution"] = False
    scored["functional_qualified"] = qualification["functional_qualified"]
    return {
        "rollout_id": rollout_id,
        "timestamp": timestamp,
        "problem_id": problem.problem_id,
        "language": language,
        "strategy": strategy,
        "phase": phase,
        "provider": provider,
        "model": model,
        "thinking": thinking,
        "duration_seconds": round(time.monotonic() - started, 3),
        "return_code": return_code,
        "error": error,
        "raw_events": str(raw_path),
        "workspace": str(workspace),
        "containment_enforced": containment_enforced,
        "workspace_scope_violation": bool(violations),
        "workspace_scope_violations": violations,
        "escape_attempt": escape_attempt,
        "containment_breach": containment_breach,
        "sandbox": sandbox_receipt,
        "sandbox_cleanup_error": sandbox_cleanup_error,
        "score_before_containment_audit": score_before_containment_audit,
        "usage": usage_from_events(events),
        **scored,
    }


def append_rows(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")


def rollout_specs(
    current: Problem,
    problems: dict[int, Problem],
    state: CurriculumState,
    languages: tuple[str, ...],
    fanout: int,
    mode: str,
    strategies: tuple[str, ...],
    regression_window: int,
) -> list[tuple[Problem, str, str, str]]:
    specs: list[tuple[Problem, str, str, str]] = []
    if mode == "canonical":
        for language in languages:
            specs.extend((current, language, "canonical", "canonical") for _ in range(fanout))
        regression_ids = range(max(1, state.highest_mastered - regression_window + 1), state.highest_mastered + 1)
        for problem_id in regression_ids:
            previous = problems.get(problem_id)
            if previous and previous.available:
                specs.extend((previous, language, "canonical", "regression") for language in languages)
    else:
        for language in languages:
            for strategy in strategies:
                specs.extend((current, language, strategy, "exploration") for _ in range(fanout))
    return specs


def main() -> None:
    parser = argparse.ArgumentParser(description="Run sequential Project Euler fan-outs through Pi.")
    parser.add_argument("--config", type=Path, default=EXAMPLE_DIR / "config.toml")
    parser.add_argument("--mode", choices=("canonical", "explore"), default="canonical")
    parser.add_argument("--run-dir", type=Path)
    parser.add_argument("--fanout", type=int, help="Override fanout per language/strategy")
    args = parser.parse_args()

    config_file = args.config.resolve()
    config = load_config(config_file)
    dataset = config["dataset"]
    model_config = config["model"]
    evaluation = config["evaluation"]
    curriculum_config = config["curriculum"]
    sandbox_settings = config.get("sandbox", {})
    qualification_mode = str(curriculum_config.get("qualification_mode", "legacy"))
    sandbox_config = docker_sandbox_config(sandbox_settings)
    if qualification_mode == "functional_v2" and sandbox_config is None:
        raise RuntimeError("functional_v2 evaluation requires the Docker sandbox")
    sandbox_preflight = verify_docker_sandbox(sandbox_config) if sandbox_config else None
    data_dir = resolve_path(config_file, str(dataset.get("prepared_dir", "data")))
    problems_file = data_dir / "problems.jsonl"
    if not problems_file.exists():
        raise FileNotFoundError(f"run prepare.py first; missing {problems_file}")
    problems = problem_map(problems_file)
    state_path = resolve_path(config_file, str(curriculum_config.get("state_file", "data/curriculum-state.json")))
    state = CurriculumState.load(state_path, int(curriculum_config.get("start_problem", 1)))
    current = problems.get(state.current_problem)
    if current is None:
        raise RuntimeError(f"problem {state.current_problem} is not present in the prepared dataset")
    if not current.available:
        raise RuntimeError(
            f"problem {current.problem_id} is locked by missing resources; rerun prepare.py --fetch-resources"
        )

    languages = tuple(str(value) for value in evaluation.get("languages", ["python"]))
    validate_runtimes(languages)
    strategies = tuple(str(value) for value in evaluation.get("strategies", ["derive", "tests_first"]))
    fanout = args.fanout or int(evaluation.get("fanout_per_language", 4))
    if fanout <= 0:
        raise ValueError("fanout must be positive")
    evaluator_sha256 = sha256_file(Path(__file__))
    scorer_sha256 = sha256_file(EXAMPLE_DIR / "score.py")
    sandbox_tool_sha256 = sha256_file(PI_SANDBOX_EXTENSION)
    model_signature = (
        f"{model_config.get('provider', '')}/{model_config.get('model', '')}:"
        f"{model_config.get('thinking', '')}:{model_config.get('revision', '')}:"
        f"gate={qualification_mode}:eval={evaluator_sha256[:16]}:"
        f"scorer={scorer_sha256[:16]}:"
        f"tools={sandbox_tool_sha256[:16]}:"
        f"sandbox={sandbox_config.image if sandbox_config else 'host'}"
    )
    pi_args = resolve_pi_args(config_file, list(model_config.get("pi_args", [])))
    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    run_dir = (args.run_dir or EXAMPLE_DIR / "runs" / f"p{current.problem_id:04d}-{args.mode}-{timestamp}").resolve()
    if run_dir.exists():
        raise FileExistsError(f"run directory already exists: {run_dir}")
    run_dir.mkdir(parents=True)
    raw_dir = run_dir / "raw"
    results_file = run_dir / "results.jsonl"
    specs = rollout_specs(
        current,
        problems,
        state,
        languages,
        fanout,
        args.mode,
        strategies,
        int(curriculum_config.get("regression_window", 0)),
    )
    manifest_path = run_dir / "manifest.json"
    manifest: dict[str, Any] = {
        "schema": "machine.project-euler.pi-eval-manifest.v1",
        "status": "prepared",
        "problem_id": current.problem_id,
        "mode": args.mode,
        "model": dict(model_config),
        "model_signature": model_signature,
        "qualification_mode": qualification_mode,
        "languages": list(languages),
        "fanout": fanout,
        "expected_rows": len(specs),
        "sandbox": sandbox_preflight,
        "source_sha256": {
            "evaluate.py": evaluator_sha256,
            "score.py": scorer_sha256,
            "curriculum.py": sha256_file(EXAMPLE_DIR / "curriculum.py"),
            "pi_sandbox_tools.ts": sandbox_tool_sha256,
            "src/pi_docker_sandbox.py": sha256_file(
                EXAMPLE_DIR / "src" / "pi_docker_sandbox.py"
            ),
        },
    }
    write_json_atomic(manifest_path, manifest)
    print(f"running {len(specs)} rollouts for problem {current.problem_id} ({args.mode})")
    rows: list[dict[str, Any]] = []
    manifest["status"] = "sampling"
    write_json_atomic(manifest_path, manifest)
    try:
        with ThreadPoolExecutor(
            max_workers=int(evaluation.get("max_concurrency", 2))
        ) as pool:
            futures = [
                pool.submit(
                    run_rollout,
                    pi_command=str(model_config.get("pi_command", "pi")),
                    pi_args=pi_args,
                    provider=str(model_config.get("provider", "")),
                    model=str(model_config.get("model", "")),
                    thinking=str(model_config.get("thinking", "")),
                    problem=problem,
                    language=language,
                    strategy=strategy,
                    phase=phase,
                    timeout_seconds=int(evaluation.get("timeout_seconds", 600)),
                    data_dir=data_dir,
                    raw_dir=raw_dir,
                    machine_data_dir=run_dir / "machine-data",
                    sandbox_config=sandbox_config,
                )
                for problem, language, strategy, phase in specs
            ]
            for future in as_completed(futures):
                row = future.result()
                rows.append(row)
                append_rows(results_file, [row])
                print(
                    f"p{row['problem_id']} {row['language']} {row['phase']} "
                    f"correct={row['correct']} functional={row['functional_qualified']} "
                    f"executed={row['successful_execution']} {row['duration_seconds']}s"
                )
    except BaseException as error:
        manifest["status"] = "failed"
        manifest["error_type"] = type(error).__name__
        manifest["completed_rows"] = len(rows)
        if results_file.exists():
            manifest["partial_results"] = {
                "path": str(results_file),
                "sha256": sha256_file(results_file),
                "rows": len(rows),
            }
        write_json_atomic(manifest_path, manifest)
        raise

    run_summary: dict[str, Any] = {
        "problem_id": current.problem_id,
        "mode": args.mode,
        "model_signature": model_signature,
        "result_count": len(rows),
        "results_file": str(results_file),
        "qualification_mode": qualification_mode,
        "sandbox": sandbox_preflight,
    }
    if args.mode == "canonical":
        gate = GateConfig(
            languages=languages,
            fanout_per_language=fanout,
            required_correct_per_language=int(curriculum_config.get("required_correct_per_language", fanout)),
            required_consecutive_rounds=int(curriculum_config.get("required_consecutive_rounds", 1)),
            minimum_regression_rate=float(curriculum_config.get("minimum_regression_rate", 0.8)),
            require_successful_execution=bool(curriculum_config.get("require_successful_execution", True)),
            qualification_mode=qualification_mode,
        )
        gate_summary = summarize_gate(rows, current.problem_id, model_signature, gate)
        unlocked = apply_round(state, gate_summary, gate)
        state.save(state_path)
        run_summary["gate"] = gate_summary
        run_summary["unlocked_next"] = unlocked
        run_summary["frontier_after"] = state.current_problem
    summary_path = run_dir / "summary.json"
    write_json_atomic(summary_path, run_summary)
    manifest["status"] = "completed"
    manifest["completed_rows"] = len(rows)
    manifest["result_receipt"] = {
        "path": str(results_file),
        "sha256": sha256_file(results_file),
        "rows": len(rows),
    }
    manifest["summary_receipt"] = {
        "path": str(summary_path),
        "sha256": sha256_file(summary_path),
    }
    write_json_atomic(manifest_path, manifest)
    print(json.dumps(run_summary, indent=2))


if __name__ == "__main__":
    main()