#!/usr/bin/env python3
"""Evaluate Pi-shaped tool use on frozen non-Euler transfer tasks."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from collections import defaultdict
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
SRC_DIR = EXAMPLE_DIR / "src"
sys.path.insert(0, str(EXAMPLE_DIR))
sys.path.insert(0, str(SRC_DIR))

from euler_common import Problem, load_jsonl  # noqa: E402
from evaluate import (  # noqa: E402
    append_rows,
    load_config,
    resolve_pi_args,
    run_rollout,
    validate_runtimes,
)
from pi_docker_sandbox import (  # noqa: E402
    docker_sandbox_config,
    verify_docker_sandbox,
)
from transfer_tasks import (  # noqa: E402
    EXTENSIONS,
    LANGUAGES,
    TransferTask,
    build_transfer_tasks,
    materialize_task_resources,
    transfer_task_digest,
)

DEFAULT_TRANSFER_DIGEST = "5857a0999a1d4b14636b3e7a9ad2afb587174eefef370c193dd0960d46e12de1"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    digest.update(path.read_bytes())
    return digest.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 build_transfer_prompt(
    task: TransferTask,
    language: str,
    resource_refs: tuple[str, ...],
) -> str:
    instruction = task.instruction.replace("<ext>", EXTENSIONS[language])
    resources = "\n".join(f"- `{ref}`" for ref in resource_refs)
    return f"""Complete transfer task `{task.task_id}`.

Use {language}. Work only with the files in the assigned workspace. You must create or edit a program file, execute it with the bash tool, and use the program's output to verify the result. Do not solve the task only in prose.

Task:
{instruction}

Available files:
{resources}

Finish with a short explanation and state the verified integer result clearly. `FINAL_ANSWER: value` is preferred, but presentation format is scored separately from functional correctness.
"""


def aggregate(rows: list[dict[str, Any]]) -> dict[str, Any]:
    buckets: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        buckets["overall"].append(row)
        buckets[f"language:{row['language']}"].append(row)
        buckets[f"task:{row['strategy']}"].append(row)

    def summarize(values: list[dict[str, Any]]) -> dict[str, Any]:
        attempts = len(values)
        return {
            "attempts": attempts,
            "functional_qualified": sum(
                bool(value.get("functional_qualified")) for value in values
            ),
            "execution_output_correct": sum(
                bool(value.get("execution_output_correct")) for value in values
            ),
            "response_answer_correct": sum(
                bool(value.get("response_answer_correct")) for value in values
            ),
            "format_valid": sum(bool(value.get("format_valid")) for value in values),
            "workspace_scope_violations": sum(
                bool(value.get("workspace_scope_violation")) for value in values
            ),
            "containment_breaches": sum(
                bool(value.get("containment_breach")) for value in values
            ),
            "errors": sum(bool(value.get("error")) for value in values),
            "tokens": sum(
                int(value.get("usage", {}).get("totalTokens", 0)) for value in values
            ),
        }

    return {
        "overall": summarize(buckets["overall"]),
        "by_language": {
            language: summarize(buckets[f"language:{language}"]) for language in LANGUAGES
        },
        "by_task": {
            task.task_id: summarize(buckets[f"task:{task.task_id}"])
            for task in build_transfer_tasks()
        },
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--run-dir", type=Path)
    parser.add_argument("--fanout", type=int, default=2)
    parser.add_argument("--languages", default=",".join(LANGUAGES))
    parser.add_argument(
        "--tasks",
        default=",".join(task.task_id for task in build_transfer_tasks()),
    )
    parser.add_argument("--expected-task-digest", default=DEFAULT_TRANSFER_DIGEST)
    return parser


def main() -> None:
    args = build_parser().parse_args()
    config_file = args.config.resolve()
    config = load_config(config_file)
    model_config = config["model"]
    evaluation = config["evaluation"]
    languages = tuple(value.strip() for value in args.languages.split(",") if value.strip())
    validate_runtimes(languages)
    all_tasks = build_transfer_tasks()
    task_digest = transfer_task_digest(all_tasks)
    if task_digest != args.expected_task_digest:
        raise RuntimeError(
            f"transfer-task digest mismatch: expected {args.expected_task_digest}, "
            f"generated {task_digest}"
        )
    requested_tasks = {value.strip() for value in args.tasks.split(",") if value.strip()}
    unknown_tasks = requested_tasks - {task.task_id for task in all_tasks}
    if unknown_tasks:
        raise ValueError(f"unknown transfer tasks: {', '.join(sorted(unknown_tasks))}")
    tasks = [task for task in all_tasks if task.task_id in requested_tasks]
    sandbox_config = docker_sandbox_config(config.get("sandbox", {}))
    if sandbox_config is None:
        raise RuntimeError("transfer evaluation requires the Docker sandbox")
    sandbox_preflight = verify_docker_sandbox(sandbox_config)

    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    run_dir = (
        args.run_dir or EXAMPLE_DIR / "runs" / f"transfer-{timestamp}"
    ).resolve()
    if run_dir.exists():
        raise FileExistsError(f"run directory already exists: {run_dir}")
    run_dir.mkdir(parents=True)
    task_data = run_dir / "task-data"
    raw_dir = run_dir / "raw"
    results_file = run_dir / "results.jsonl"
    pi_args = resolve_pi_args(config_file, list(model_config.get("pi_args", [])))

    conditions: list[tuple[Problem, TransferTask, str, str]] = []
    private_tasks: list[dict[str, Any]] = []
    for task_index, task in enumerate(tasks, start=9001):
        for language in languages:
            refs = materialize_task_resources(task, language, task_data)
            problem = Problem(
                problem_id=task_index,
                prompt=task.instruction,
                answer=task.answer,
                available=True,
                resource_refs=refs,
            )
            prompt = build_transfer_prompt(task, language, refs)
            conditions.extend((problem, task, language, prompt) for _ in range(args.fanout))
            private_tasks.append(
                {
                    "problem_id": task_index,
                    "task_id": task.task_id,
                    "language": language,
                    "answer": task.answer,
                    "resource_refs": list(refs),
                }
            )

    manifest = {
        "schema": "machine.project-euler.transfer-eval-manifest.v1",
        "status": "prepared",
        "model": dict(model_config),
        "fanout": args.fanout,
        "languages": list(languages),
        "selected_tasks": [task.task_id for task in tasks],
        "task_digest": task_digest,
        "expected_task_digest": args.expected_task_digest,
        "sandbox": sandbox_preflight,
        "tasks": private_tasks,
        "source_sha256": {
            "evaluate_transfer.py": sha256_file(Path(__file__)),
            "evaluate.py": sha256_file(EXAMPLE_DIR / "evaluate.py"),
            "score.py": sha256_file(EXAMPLE_DIR / "score.py"),
            "src/transfer_tasks.py": sha256_file(SRC_DIR / "transfer_tasks.py"),
            "pi_sandbox_tools.ts": sha256_file(EXAMPLE_DIR / "pi_sandbox_tools.ts"),
        },
    }
    manifest_path = run_dir / "manifest.json"
    write_json_atomic(manifest_path, manifest)

    rows: list[dict[str, Any]] = []
    manifest["status"] = "sampling"
    manifest["expected_rows"] = len(conditions)
    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=task.task_id,
                    phase="transfer",
                    timeout_seconds=int(evaluation.get("timeout_seconds", 300)),
                    data_dir=task_data,
                    raw_dir=raw_dir,
                    machine_data_dir=run_dir / "machine-data",
                    sandbox_config=sandbox_config,
                    prompt_override=prompt,
                )
                for problem, task, language, prompt in conditions
            ]
            for future in as_completed(futures):
                row = future.result()
                rows.append(row)
                append_rows(results_file, [row])
                print(
                    f"{row['strategy']} {row['language']} "
                    f"functional={row['functional_qualified']} {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(load_jsonl(results_file)),
            }
        write_json_atomic(manifest_path, manifest)
        raise

    if len(rows) != len(conditions):
        manifest["status"] = "failed"
        manifest["error_type"] = "IncompleteResults"
        manifest["completed_rows"] = len(rows)
        write_json_atomic(manifest_path, manifest)
        raise RuntimeError(
            f"transfer evaluation completed {len(rows)} of {len(conditions)} expected rows"
        )

    summary = {
        "schema": "machine.project-euler.transfer-eval-summary.v1",
        "status": "completed",
        "task_digest": task_digest,
        "model_revision": str(model_config.get("revision", "")),
        "results_file": str(results_file),
        **aggregate(rows),
    }
    summary_path = run_dir / "summary.json"
    write_json_atomic(summary_path, summary)
    manifest["status"] = "completed"
    manifest["completed_rows"] = len(rows)
    manifest["result_receipt"] = {
        "path": str(results_file),
        "sha256": sha256_file(results_file),
        "rows": len(load_jsonl(results_file)),
    }
    manifest["summary_receipt"] = {
        "path": str(summary_path),
        "sha256": sha256_file(summary_path),
    }
    write_json_atomic(manifest_path, manifest)
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()