#!/usr/bin/env python3
"""Run a task suite against any agent entry point, log the trajectories and report the cost.

Purpose: the measurement half of Part 26. It loads an agent entry point by file path,
    runs every task in a suite against it (optionally several times), scores each answer
    with deterministic checks, writes one JSON-lines trajectory file per run, and prints
    success rate, mean steps, mean tokens and, when you supply a cost per million tokens
    from Part 23's cost-model.py, a cost per task. It knows nothing about frameworks:
    anything that satisfies the four-function contract below can be measured with it.
Platform: all (pure Python over HTTP; the models may be served on any track or machine)
Minimum memory: 16 GB on the machine serving the models; this script needs almost none
Assumes: Python 3.9 or later and no third-party packages for the harness itself. The agent
    entry point may need whatever its framework needs. A task file in the shape of
    agent-tasks.json. Nothing here starts a server: the endpoints in --base-url and
    --option must already be answering.

    The entry-point contract, which is the whole interface:

        def run_task(task: str, options: dict) -> dict
            Required. Runs one task and returns at least
              {"answer": str, "steps": int, "tokens": int}
            and optionally "prompt_tokens", "completion_tokens", "seconds", "stopped"
            and "trajectory": a list of JSON-serialisable step records.

        def build(options: dict) -> None      Optional. Called once, before task one.
        def close() -> None                   Optional. Called once, at the end.
        SCAFFOLD_NAME: str                    Optional. Used in reports and file names.

Usage: python3 agent-eval.py --agent multi-agent-system.py --tasks agent-tasks.json \\
           --model local/chat --workspace ./agent-workspace --out results.json
       python3 agent-eval.py --agent scaffold-minimal.py --tasks agent-tasks.json \\
           --model local/router --repeats 3 --cost-per-million 0.42 --labbook labbook.md
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import platform
import re
import statistics
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional


# --------------------------------------------------------------------------------------
# Loading an agent entry point
# --------------------------------------------------------------------------------------

def load_agent(path: str):
    """Import a Python file as a module and check it satisfies the contract."""
    file = Path(path)
    if not file.exists():
        sys.exit(f"agent entry point {path} does not exist")
    spec = importlib.util.spec_from_file_location(file.stem.replace("-", "_"), file)
    if spec is None or spec.loader is None:
        sys.exit(f"cannot load {path} as a Python module")
    module = importlib.util.module_from_spec(spec)
    # The entry point may import files that sit beside it, so its directory goes on the
    # path before it is executed.
    sys.path.insert(0, str(file.resolve().parent))
    spec.loader.exec_module(module)
    if not hasattr(module, "run_task"):
        sys.exit(f"{path} defines no run_task(task, options); see the contract in this file")
    return module


# --------------------------------------------------------------------------------------
# Scoring
# --------------------------------------------------------------------------------------

def check_answer(task: dict, record: dict) -> dict:
    """Deterministic checks only. No model grades anything here, so a score is repeatable."""
    answer = (record.get("answer") or "")
    lowered = answer.lower()
    missing = [s for s in task.get("expect_in_answer", []) if s.lower() not in lowered]
    forbidden = [s for s in task.get("expect_not_in_answer", []) if s.lower() in lowered]

    pattern = task.get("expect_regex")
    regex_ok = True if not pattern else bool(re.search(pattern, answer, re.IGNORECASE))

    used = [step.get("tool") for step in record.get("trajectory", []) if step.get("tool")]
    wanted_tool = task.get("expect_tool")
    tool_ok = True if not wanted_tool else wanted_tool in used

    limit = task.get("max_steps")
    steps = int(record.get("steps") or 0)
    within_steps = True if not limit else steps <= limit

    return {
        "missing": missing,
        "forbidden": forbidden,
        "regex_ok": regex_ok,
        "tool_ok": tool_ok,
        "within_steps": within_steps,
        "tools_used": used,
        "passed": (not missing and not forbidden and regex_ok and tool_ok and within_steps),
    }


# --------------------------------------------------------------------------------------
# Reporting
# --------------------------------------------------------------------------------------

def summarise(rows: List[dict], cost_per_million: Optional[float]) -> dict:
    """Aggregate the per-run records. Every figure here is arithmetic over the runs."""
    if not rows:
        return {"runs": 0}
    passed = [r for r in rows if r["checks"]["passed"]]
    steps = [r["steps"] for r in rows]
    tokens = [r["tokens"] for r in rows]
    seconds = [r["seconds"] for r in rows]
    summary = {
        "runs": len(rows),
        "passed": len(passed),
        "success_rate": round(len(passed) / len(rows), 3),
        "mean_steps": round(statistics.fmean(steps), 2),
        "mean_tokens": round(statistics.fmean(tokens), 1),
        "mean_seconds": round(statistics.fmean(seconds), 1),
        "total_tokens": sum(tokens),
    }
    if len(rows) > 1:
        summary["stdev_steps"] = round(statistics.pstdev(steps), 2)
        summary["stdev_tokens"] = round(statistics.pstdev(tokens), 1)
    if cost_per_million is not None:
        summary["cost_per_million_tokens"] = cost_per_million
        summary["cost_total"] = round(sum(tokens) / 1e6 * cost_per_million, 6)
        summary["cost_per_task"] = round(summary["cost_total"] / len(rows), 6)
    return summary


def per_task_view(rows: List[dict]) -> List[dict]:
    """One line per task id, so a repeated suite shows where the variance is."""
    by_id: Dict[str, List[dict]] = {}
    for row in rows:
        by_id.setdefault(row["id"], []).append(row)
    view = []
    for task_id, runs in by_id.items():
        successes = sum(1 for r in runs if r["checks"]["passed"])
        steps = [r["steps"] for r in runs]
        tokens = [r["tokens"] for r in runs]
        view.append({
            "id": task_id,
            "category": runs[0]["category"],
            "attempts": len(runs),
            "successes": successes,
            "success_rate": round(successes / len(runs), 3),
            "steps_min": min(steps), "steps_max": max(steps),
            "tokens_min": min(tokens), "tokens_max": max(tokens),
            "first_failure": next((r["checks"] for r in runs if not r["checks"]["passed"]), None),
        })
    return view


# --------------------------------------------------------------------------------------
# The run
# --------------------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--agent", required=True,
                        help="Python file defining run_task(task, options); see the header")
    parser.add_argument("--tasks", default="agent-tasks.json")
    parser.add_argument("--only", default=None, help="run one category only")
    parser.add_argument("--task-id", default=None, help="run one task by id")
    parser.add_argument("--repeats", type=int, default=1,
                        help="run every task this many times; 3 or more to see the variance")
    parser.add_argument("--out", default="agent-results.json")
    parser.add_argument("--trajectory-dir", default="trajectories")

    # Everything below is handed to the agent unchanged, in the options dictionary.
    parser.add_argument("--base-url", default="http://127.0.0.1:4000/v1",
                        help="OpenAI-compatible endpoint, normally the Part 9 gateway")
    parser.add_argument("--model", required=True, help="alias the answering model is served as")
    parser.add_argument("--router-model", default=None, help="alias of the small routing model")
    parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"))
    parser.add_argument("--workspace", default=None, help="the only directory tools may read")
    parser.add_argument("--index", default=None, help="a Part 10 document index, optional")
    parser.add_argument("--option", action="append", default=[], metavar="KEY=VALUE",
                        help="any further option the entry point understands; repeatable")

    # Recording.
    parser.add_argument("--cost-per-million", type=float, default=None,
                        help="your own cost per million tokens from Part 23's cost-model.py. "
                             "There is no default: this script invents no price.")
    parser.add_argument("--quant", default="unknown",
                        help="the quantisation actually loaded; not discoverable over the API")
    parser.add_argument("--engine", default="unknown", help="engine name, for the record")
    parser.add_argument("--engine-version", default="unknown", help="engine version, for the record")
    parser.add_argument("--notes", default="", help="one line about what this run is testing")
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()

    options: Dict[str, Any] = {
        "base_url": args.base_url,
        "model": args.model,
        "router_model": args.router_model or args.model,
        "api_key": args.api_key,
        "workspace": args.workspace,
        "index": args.index,
    }
    for item in args.option:
        if "=" not in item:
            sys.exit(f"--option needs KEY=VALUE, got {item!r}")
        key, value = item.split("=", 1)
        options[key.strip()] = value.strip()

    spec = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
    tasks = spec["tasks"]
    if args.only:
        tasks = [t for t in tasks if t.get("category") == args.only]
    if args.task_id:
        tasks = [t for t in tasks if t["id"] == args.task_id]
    if not tasks:
        sys.exit("no tasks selected")

    agent = load_agent(args.agent)
    scaffold = getattr(agent, "SCAFFOLD_NAME", Path(args.agent).stem)
    if hasattr(agent, "build"):
        agent.build(options)

    out_dir = Path(args.trajectory_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    stamp = time.strftime("%Y%m%d-%H%M%S")
    safe_model = args.model.replace("/", "_")

    rows: List[dict] = []
    started = time.time()
    print(f"{scaffold} x {args.model}: {len(tasks)} task(s), {args.repeats} attempt(s) each\n")

    try:
        for attempt in range(1, args.repeats + 1):
            for task in tasks:
                began = time.time()
                try:
                    record = agent.run_task(task["task"], options)
                except Exception as exc:  # an agent that crashes is a failed task, not a crash
                    record = {"answer": "", "steps": 0, "tokens": 0,
                              "stopped": f"entry point raised {type(exc).__name__}: {exc}",
                              "trajectory": []}
                elapsed = record.get("seconds", round(time.time() - began, 1))

                trajectory = record.get("trajectory", [])
                path = out_dir / f"{stamp}-{scaffold}-{safe_model}-{task['id']}-r{attempt}.jsonl"
                with path.open("w", encoding="utf-8") as handle:
                    handle.write(json.dumps({"meta": {
                        "scaffold": scaffold, "model": args.model, "task_id": task["id"],
                        "attempt": attempt, "task": task["task"]}}) + "\n")
                    for step in trajectory:
                        handle.write(json.dumps(step) + "\n")

                checks = check_answer(task, record)
                rows.append({
                    "id": task["id"], "category": task.get("category", "uncategorised"),
                    "attempt": attempt, "answer": record.get("answer", ""),
                    "steps": int(record.get("steps") or 0),
                    "tokens": int(record.get("tokens") or 0),
                    "prompt_tokens": record.get("prompt_tokens"),
                    "completion_tokens": record.get("completion_tokens"),
                    "seconds": float(elapsed),
                    "stopped": record.get("stopped", ""),
                    "checks": checks, "trajectory_file": str(path),
                })
                mark = "ok " if checks["passed"] else "BAD"
                print(f"  {mark} {task['id']:<16} r{attempt} "
                      f"{record.get('steps', 0):>2} step(s) {record.get('tokens', 0):>7} tok "
                      f"{elapsed:>6.1f}s")
                if args.verbose and not checks["passed"]:
                    print(f"      answer: {(record.get('answer') or '')[:200]}")
                    if checks["missing"]:
                        print(f"      missing: {checks['missing']}")
    finally:
        if hasattr(agent, "close"):
            agent.close()

    summary = summarise(rows, args.cost_per_million)
    run = {
        "lab": "part-26/evaluating-agents",
        "run_id": stamp,
        "scaffold": scaffold,
        "agent_entry_point": args.agent,
        "task_set": spec.get("name", args.tasks),
        "task_set_version": spec.get("version"),
        "model": args.model,
        "router_model": options["router_model"],
        "quant": args.quant,
        "engine": args.engine,
        "engine_version": args.engine_version,
        "base_url": args.base_url,
        "repeats": args.repeats,
        "host": platform.platform(),
        "date": time.strftime("%Y-%m-%d"),
        "wall_seconds": round(time.time() - started, 1),
        "notes": args.notes,
        "summary": summary,
    }

    Path(args.out).write_text(
        json.dumps({"run": run, "per_task": per_task_view(rows), "results": rows}, indent=2),
        encoding="utf-8")

    print(f"\nsuccess rate  {summary['success_rate']:.2f}  "
          f"({summary['passed']} of {summary['runs']})")
    print(f"mean steps    {summary['mean_steps']}")
    print(f"mean tokens   {summary['mean_tokens']}")
    print(f"mean seconds  {summary['mean_seconds']}")
    if "cost_per_task" in summary:
        print(f"cost per task {summary['cost_per_task']} "
              f"(at {summary['cost_per_million_tokens']} per million tokens, your figure)")
    else:
        print("cost          not computed: pass --cost-per-million from Part 23's cost-model.py")
    print(f"\nresults written to {args.out}; trajectories in {out_dir}/")

    if args.labbook:
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(run) + "\n")
        print(f"recorded in {args.labbook}")

    sys.exit(0 if summary["passed"] == summary["runs"] else 2)


if __name__ == "__main__":
    main()