#!/usr/bin/env python3
"""No-spend preflight for the hidden-variant tool-RL environment."""

from __future__ import annotations

import argparse
import asyncio
import json
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any

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

from euler_variants import (  # noqa: E402
    FAMILIES,
    LANGUAGES,
    SANDBOX_LANGUAGES,
    VariantTask,
    assert_split_isolation,
    assess_sandbox_response,
    build_sandbox_program,
    build_variant_tasks,
    ensure_loopback_sandbox_url,
    task_set_digest,
)

REFERENCE_SOLUTIONS = {
    "sum_multiples": {
        "python": """def solve(limit, divisor_a, divisor_b):
    return sum(value for value in range(limit) if value % divisor_a == 0 or value % divisor_b == 0)
""",
        "javascript": """function solve(limit, divisorA, divisorB) {
  let total = 0;
  for (let value = 0; value < limit; value++) {
    if (value % divisorA === 0 || value % divisorB === 0) total += value;
  }
  return total;
}
""",
        "julia": """function solve(limit, divisor_a, divisor_b)
    return sum(value for value in 0:(limit - 1) if value % divisor_a == 0 || value % divisor_b == 0)
end
""",
    },
    "even_fibonacci": {
        "python": """def solve(limit):
    total = 0
    left, right = 1, 2
    while left < limit:
        if left % 2 == 0:
            total += left
        left, right = right, left + right
    return total
""",
        "javascript": """function solve(limit) {
  let total = 0;
  let left = 1;
  let right = 2;
  while (left < limit) {
    if (left % 2 === 0) total += left;
    [left, right] = [right, left + right];
  }
  return total;
}
""",
        "julia": """function solve(limit)
    total = 0
    left, right = 1, 2
    while left < limit
        if left % 2 == 0
            total += left
        end
        left, right = right, left + right
    end
    return total
end
""",
    },
}


LOCAL_COMMANDS = {
    "python": ("python3", "-c"),
    "javascript": ("node", "-e"),
    "julia": ("julia", "-e"),
}


def _fake_response(return_code: int, stdout: str, stderr: str) -> dict[str, Any]:
    return {
        "status": "Success" if return_code == 0 else "RuntimeError",
        "compile_result": None,
        "run_result": {
            "status": "Finished",
            "return_code": return_code,
            "stdout": stdout,
            "stderr": stderr,
        },
    }


def run_trusted_local(task: VariantTask, language: str) -> dict[str, Any]:
    executable, flag = LOCAL_COMMANDS[language]
    resolved = shutil.which(executable)
    if resolved is None:
        raise RuntimeError(f"required local runtime is missing: {executable}")
    program = build_sandbox_program(
        task, language, REFERENCE_SOLUTIONS[task.family][language]
    )
    completed = subprocess.run(
        [resolved, flag, program],
        capture_output=True,
        text=True,
        timeout=20,
        check=False,
    )
    assessment = assess_sandbox_response(
        task,
        transport_success=completed.returncode == 0,
        response=_fake_response(
            completed.returncode, completed.stdout, completed.stderr
        ),
    )
    if not assessment.passed:
        raise RuntimeError(
            f"trusted local {task.family}/{language} reference failed: "
            f"{assessment.status}: {assessment.diagnostic}"
        )
    return {
        "family": task.family,
        "language": language,
        "status": assessment.status,
        "passed_cases": assessment.passed_cases,
        "total_cases": assessment.total_cases,
    }


async def run_trusted_sandbox(
    tasks: list[VariantTask],
    sandbox_url: str,
    timeout: int,
) -> list[dict[str, Any]]:
    from tinker_cookbook.sandbox import SandboxFusionClient

    client = SandboxFusionClient(url=ensure_loopback_sandbox_url(sandbox_url))
    receipts: list[dict[str, Any]] = []
    try:
        for task in tasks:
            for language in LANGUAGES:
                success, response = await client.run(
                    code=build_sandbox_program(
                        task,
                        language,
                        REFERENCE_SOLUTIONS[task.family][language],
                    ),
                    files={},
                    timeout=timeout,
                    language=SANDBOX_LANGUAGES[language],
                )
                assessment = assess_sandbox_response(
                    task,
                    transport_success=success,
                    response=response,
                )
                if not assessment.passed:
                    raise RuntimeError(
                        f"SandboxFusion {task.family}/{language} reference failed: "
                        f"{assessment.status}: {assessment.diagnostic}"
                    )
                receipts.append(
                    {
                        "family": task.family,
                        "language": language,
                        "status": assessment.status,
                        "passed_cases": assessment.passed_cases,
                        "total_cases": assessment.total_cases,
                    }
                )
    finally:
        await client.close()
    return receipts


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seed", type=int, default=1729)
    parser.add_argument("--tasks-per-family", type=int, default=2)
    parser.add_argument("--cases-per-task", type=int, default=4)
    parser.add_argument(
        "--sandbox",
        action="store_true",
        help="Also run trusted references through SandboxFusion",
    )
    parser.add_argument("--sandbox-url", default="http://127.0.0.1:8080/run_code")
    parser.add_argument("--sandbox-timeout", type=int, default=10)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    train = build_variant_tasks(
        FAMILIES,
        split="train",
        tasks_per_family=args.tasks_per_family,
        cases_per_task=args.cases_per_task,
        seed=args.seed,
    )
    eval_tasks = build_variant_tasks(
        FAMILIES,
        split="eval",
        tasks_per_family=args.tasks_per_family,
        cases_per_task=args.cases_per_task,
        seed=args.seed,
    )
    assert_split_isolation(train, eval_tasks)
    local_receipts = [
        run_trusted_local(task, language)
        for task in (train[0], train[args.tasks_per_family])
        for language in LANGUAGES
    ]
    sandbox_receipts = (
        asyncio.run(
            run_trusted_sandbox(
                [train[0], train[args.tasks_per_family]],
                args.sandbox_url,
                args.sandbox_timeout,
            )
        )
        if args.sandbox
        else []
    )
    receipt = {
        "schema": "machine.project-euler.tool-rl-preflight.v1",
        "status": "passed",
        "seed": args.seed,
        "families": list(FAMILIES),
        "languages": list(LANGUAGES),
        "train_tasks": len(train),
        "eval_tasks": len(eval_tasks),
        "cases_per_task": args.cases_per_task,
        "train_digest": task_set_digest(train),
        "eval_digest": task_set_digest(eval_tasks),
        "split_overlap": 0,
        "local_runtime_receipts": local_receipts,
        "sandbox_requested": args.sandbox,
        "sandbox_receipts": sandbox_receipts,
        "tinker_called": False,
        "training_started": False,
    }
    print(json.dumps(receipt, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()