#!/usr/bin/env python3
"""Run a trusted local Python agent over a Gradia supplied-output dataset.

Python 3.10+, standard library only. Each case gets a fresh process and only its
input is passed to the function. This is NOT a sandbox: agent code has your OS
permissions, network access, environment and files, and can incur provider costs.
The runner itself calls no providers and never retries. Agent stdout/stderr are
discarded; only bounded JSON results are retained. POSIX process groups are killed
after each case. On Windows, only the direct worker is terminated; descendant
cleanup is not guaranteed. Do not use this runner for untrusted agent code.

Example:
  python gradia-agent-eval.py --cases cases.json --agent agent.py:run \
      --slot candidate --output candidate-results.json

Existing target-slot evidence and output files are never replaced. Checkpoints
retain completed/error cases, the other slot, and every still-missing case. An
interruption records the current case as an error and leaves later cases missing.
Results are descriptive supplied-output evidence, not execution attestations or
native task grades. Timeouts cannot roll back external actions already performed.
"""

from __future__ import annotations

import argparse
import asyncio
import copy
import importlib.util
import inspect
import json
import math
import os
import re
import signal
import subprocess
import sys
import tempfile
import time
from decimal import Decimal, InvalidOperation
from pathlib import Path

SCHEMA = "gradia.agent-eval.dataset.v1"
MAX_BYTES = 2_000_000
MAX_CASES = 200
MAX_DEPTH = 24
MAX_NODES = 50_000
MAX_NUMBER = 9_007_199_254_740_991
JS_WHITESPACE = "\t\n\v\f\r \u00a0\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000\ufeff"


class ValidationError(ValueError):
    """Invalid or unsupported dataset/result, without echoing its contents."""


def _length(text: str) -> int:
    return len(text.encode("utf-16-le", errors="surrogatepass")) // 2


def validate_json(value: object) -> None:
    nodes = 0

    def visit(item: object, depth: int) -> None:
        nonlocal nodes
        nodes += 1
        if depth > MAX_DEPTH or nodes > MAX_NODES:
            raise ValidationError("JSON exceeds the depth or value-count limit.")
        if item is None or isinstance(item, (str, bool)):
            return
        if isinstance(item, (int, float)):
            if not math.isfinite(item) or abs(item) > MAX_NUMBER:
                raise ValidationError("JSON numbers must be finite and within the safe integer range.")
            return
        if isinstance(item, list):
            for child in item:
                visit(child, depth + 1)
            return
        if isinstance(item, dict) and all(isinstance(key, str) for key in item):
            for child in item.values():
                visit(child, depth + 1)
            return
        raise ValidationError("Values must be JSON objects, arrays, strings, numbers, booleans or null.")

    visit(value, 0)


def json_bytes(value: object) -> bytes:
    validate_json(value)
    chunks: list[bytes] = []
    size = 0
    encoder = json.JSONEncoder(ensure_ascii=False, allow_nan=False, separators=(",", ":"))
    for chunk in encoder.iterencode(value):
        # Escape isolated surrogate code points like JSON.stringify, while keeping
        # ordinary Unicode compact and the file valid UTF-8.
        encoded = chunk.encode("utf-8", errors="backslashreplace")
        size += len(encoded)
        if size > MAX_BYTES:
            raise ValidationError("JSON exceeds the 2 MB limit.")
        chunks.append(encoded)
    return b"".join(chunks)


def strict_json(raw: bytes) -> object:
    if len(raw) > MAX_BYTES:
        raise ValidationError("JSON exceeds the 2 MB limit.")

    def pairs(items: list[tuple[str, object]]) -> dict[str, object]:
        result: dict[str, object] = {}
        for key, value in items:
            if key in result:
                raise ValidationError("Duplicate JSON fields are not allowed.")
            result[key] = value
        return result

    def constant(_: str) -> None:
        raise ValidationError("Non-finite numbers are not valid JSON.")

    def number(literal: str) -> int | float:
        if len(literal) > 128:
            raise ValidationError("JSON number literals must be at most 128 characters.")
        try:
            numeric = float(literal)
            if (not math.isfinite(numeric) or abs(numeric) > MAX_NUMBER
                    or Decimal(literal) != Decimal(str(numeric))):
                raise ValidationError("JSON numbers must round-trip without precision loss in the browser.")
            return int(literal) if not any(char in literal for char in ".eE") else numeric
        except (InvalidOperation, OverflowError) as error:
            raise ValidationError("Invalid or unsupported JSON number.") from error

    try:
        result = json.loads(raw.decode("utf-8"), object_pairs_hook=pairs, parse_constant=constant,
                            parse_int=number, parse_float=number)
        validate_json(result)
        return result
    except (UnicodeError, json.JSONDecodeError, RecursionError, OverflowError) as error:
        raise ValidationError("Invalid JSON or unsupported nesting.") from error


def _object(value: object, allowed: set[str], label: str) -> dict:
    if not isinstance(value, dict) or set(value) - allowed:
        raise ValidationError(f"{label} must be an object with supported fields only.")
    return value


def _attempt(value: object) -> dict:
    attempt = _object(value, {"status", "output", "error", "duration_ms"}, "Attempt")
    if "duration_ms" in attempt:
        duration = attempt["duration_ms"]
        if isinstance(duration, bool) or not isinstance(duration, (int, float)) or duration < 0:
            raise ValidationError("Attempt duration_ms must be a non-negative number.")
    if attempt.get("status") == "completed" and "output" in attempt and "error" not in attempt:
        return attempt
    error = attempt.get("error")
    if (attempt.get("status") == "error" and isinstance(error, str) and error.strip(JS_WHITESPACE)
            and _length(error) <= 500 and "output" not in attempt):
        return attempt
    raise ValidationError("Attempt needs completed with output, or error with a short description.")


def parse_dataset(raw: bytes) -> dict:
    data = _object(strict_json(raw), {"schema", "name", "cases"}, "Dataset")
    if data.get("schema") != SCHEMA:
        raise ValidationError("Use the gradia.agent-eval.dataset.v1 schema.")
    name = data.get("name")
    if not isinstance(name, str) or not name.strip(JS_WHITESPACE) or _length(name) > 160:
        raise ValidationError("Dataset name must contain 1–160 characters.")
    cases = data.get("cases")
    if not isinstance(cases, list) or not 1 <= len(cases) <= MAX_CASES:
        raise ValidationError("Include 1–200 cases.")
    seen = set()
    for item in cases:
        case = _object(item, {"id", "input", "expected", "baseline", "candidate"}, "Case")
        identifier = case.get("id")
        if (not isinstance(identifier, str) or not identifier.strip(JS_WHITESPACE) or _length(identifier) > 120
                or identifier in seen):
            raise ValidationError("Case IDs must be unique non-empty strings of at most 120 characters.")
        seen.add(identifier)
        if "input" not in case or "expected" not in case:
            raise ValidationError("Each case requires input and expected; explicit null is allowed.")
        for slot in ("baseline", "candidate"):
            if slot in case:
                _attempt(case[slot])
    data["name"] = name.strip(JS_WHITESPACE)
    return data


def _error(reason: str, duration: float = 0) -> dict:
    return {"status": "error", "error": reason, "duration_ms": round(max(0, duration), 3)}


def _fits_with_remaining(data: dict, slot: str) -> None:
    reserved = copy.deepcopy(data)
    for case in reserved["cases"]:
        if slot not in case:
            case[slot] = _error("output_limit_exceeded", MAX_NUMBER)
    json_bytes(reserved)


def _checkpoint(output: Path, data: dict) -> None:
    payload = json_bytes(data)
    fd, temporary = tempfile.mkstemp(prefix=f".{output.name}.", dir=output.parent)
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(payload)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, output)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def _stop(process: subprocess.Popen) -> None:
    try:
        if os.name == "posix":
            os.killpg(process.pid, signal.SIGKILL)
        elif process.poll() is None:
            process.kill()
    except ProcessLookupError:
        pass
    process.wait()


def _worker(agent: str, function: str, response: str) -> int:
    # stdout and stderr point at DEVNULL in the parent. JSON uses a separate file
    # so print(), logging and native-library writes cannot corrupt the protocol.
    try:
        value = strict_json(sys.stdin.buffer.read(MAX_BYTES + 1))
        sys.path.insert(0, str(Path(agent).parent))
        spec = importlib.util.spec_from_file_location("_gradia_evaluation_agent", agent)
        if spec is None or spec.loader is None:
            raise ImportError("No Python loader")
        module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = module
        spec.loader.exec_module(module)
        callable_agent = getattr(module, function)
        if not callable(callable_agent):
            raise TypeError("Agent is not callable")
    except BaseException:
        result = _error("agent_load_error")
    else:
        try:
            output = callable_agent(value)
            if inspect.isawaitable(output):
                async def await_output():
                    return await output
                output = asyncio.run(await_output())
            try:
                json_bytes(output)
                result = {"status": "completed", "output": output}
            except (ValueError, TypeError, RecursionError, OverflowError):
                result = _error("invalid_agent_output")
        except BaseException as error:
            # Never retain exception messages: they may contain inputs or secrets.
            kind = re.sub(r"[^A-Za-z0-9_]", "", type(error).__name__)[:60]
            result = _error(f"agent_exception:{kind or 'Error'}")
    try:
        payload = json_bytes(result)
    except (ValueError, TypeError, RecursionError, OverflowError):
        payload = json_bytes(_error("output_limit_exceeded"))
    with open(response, "wb") as handle:
        handle.write(payload)
    return 0


def _invoke(agent: Path, function: str, value: object, timeout: float) -> tuple[dict, bool]:
    started = time.monotonic()
    interrupted = False
    with tempfile.TemporaryDirectory(prefix="gradia-eval-") as directory:
        response = Path(directory) / "response.json"
        with tempfile.TemporaryFile() as request:
            request.write(json_bytes(value))
            request.seek(0)
            process = subprocess.Popen(
                [sys.executable, str(Path(__file__).resolve()), "--_worker", str(agent), function, str(response)],
                stdin=request, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                start_new_session=os.name == "posix",
            )
            try:
                process.wait(timeout=timeout)
                if process.returncode != 0:
                    attempt = _error("agent_process_failed")
                else:
                    try:
                        with response.open("rb") as handle:
                            attempt = _attempt(strict_json(handle.read(MAX_BYTES + 1)))
                    except (OSError, ValueError, TypeError):
                        attempt = _error("invalid_worker_response")
            except subprocess.TimeoutExpired:
                attempt = _error("agent_timeout")
            except KeyboardInterrupt:
                interrupted = True
                attempt = _error("agent_interrupted")
            finally:
                _stop(process)
    attempt["duration_ms"] = round((time.monotonic() - started) * 1000, 3)
    return attempt, interrupted


def run(cases: Path, agent_spec: str, slot: str, output: Path, timeout: float = 30) -> int:
    if slot not in {"baseline", "candidate"}:
        raise ValidationError("Choose baseline or candidate.")
    if not math.isfinite(timeout) or not 0 < timeout <= 300:
        raise ValidationError("Timeout must be greater than zero and at most 300 seconds.")
    agent_name, separator, function = agent_spec.rpartition(":")
    if not separator or not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", function):
        raise ValidationError("Use --agent path/to/agent.py:function.")
    agent = Path(agent_name).expanduser().resolve()
    if not agent.is_file() or agent.suffix.lower() != ".py":
        raise ValidationError("Agent must be an existing Python .py file.")
    with cases.open("rb") as handle:
        data = parse_dataset(handle.read(MAX_BYTES + 1))
    if any(slot in case for case in data["cases"]):
        raise ValidationError("The selected slot already contains evidence. Use a fresh dataset or the other slot.")
    _fits_with_remaining(data, slot)
    output = output.expanduser().absolute()
    # Reserve before importing or invoking any agent. O_EXCL also rejects existing
    # symlinks; checkpoint replacement is only for the file this run reserved.
    descriptor = os.open(output, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    os.close(descriptor)
    _checkpoint(output, data)
    print("Running trusted local code with your permissions; this is not a sandbox.", file=sys.stderr)
    if os.name != "posix":
        print("Windows: timeout cleanup terminates the worker only, not all descendants.", file=sys.stderr)
    interrupted = False
    active_case = None
    try:
        for index, case in enumerate(data["cases"]):
            active_case = case
            attempt, interrupted = _invoke(agent, function, case["input"], timeout)
            case[slot] = attempt
            try:
                _fits_with_remaining(data, slot)
            except ValidationError:
                case[slot] = _error("output_limit_exceeded", attempt["duration_ms"])
            _checkpoint(output, data)
            print(f"Case {index + 1}/{len(data['cases'])}: {case[slot]['status']}", file=sys.stderr)
            if interrupted:
                break
    except KeyboardInterrupt:
        # An interrupt between cases does not pretend the next case was attempted.
        # If a result was already received, persist it before reporting its counts;
        # a signal can otherwise land while validating or replacing a checkpoint.
        interrupted = True
        if active_case is not None and slot in active_case:
            try:
                _fits_with_remaining(data, slot)
            except ValidationError:
                active_case[slot] = _error("output_limit_exceeded", active_case[slot]["duration_ms"])
        _checkpoint(output, data)
    completed = sum(case.get(slot, {}).get("status") == "completed" for case in data["cases"])
    errors = sum(case.get(slot, {}).get("status") == "error" for case in data["cases"])
    missing = sum(slot not in case for case in data["cases"])
    print(f"Saved {completed} completed, {errors} error, {missing} missing cases to {output}")
    return 130 if interrupted else (1 if errors else 0)


def main(argv: list[str] | None = None) -> int:
    argv = sys.argv[1:] if argv is None else argv
    if argv and argv[0] == "--_worker":
        if len(argv) != 4:
            return 2
        return _worker(*argv[1:])
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--cases", required=True, type=Path)
    parser.add_argument("--agent", required=True, help="Trusted Python file.py:function; receives case input only")
    parser.add_argument("--slot", required=True, choices=("baseline", "candidate"))
    parser.add_argument("--output", required=True, type=Path, help="New output file; existing files are never replaced")
    parser.add_argument("--timeout", type=float, default=30, help="Per-case seconds, greater than 0 up to 300 (default 30)")
    args = parser.parse_args(argv)
    try:
        return run(args.cases, args.agent, args.slot, args.output, args.timeout)
    except KeyboardInterrupt:
        print("Evaluation interrupted; any saved checkpoint is retained.", file=sys.stderr)
        return 130
    except (OSError, ValidationError) as error:
        # OSError details may contain local paths, so keep the message bounded.
        message = str(error) if isinstance(error, ValidationError) else type(error).__name__
        print(f"Cannot run evaluation: {message}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())