"""Agentic multi-task instrument — the v2-hard gauntlet.

Runs multiple build-a-real-app tasks (each in a different language/framework)
against one model, sequentially. Each task is a self-contained module under
`bench/tasks/multitask/` exposing TASK_NAME, PLAN, STARTER, REFERENCE,
validate(), gates(), TOOLS, TOOL_IMPLS. The harness materializes the
starter, runs a multi-turn tool loop, then scores via the task's gates().

Shipped tasks (the gauntlet):
  go_htmx   — Go + HTMX partials + embedded SQLite + recursive CTE
  rust_cli  — Rust + reqwest + quick-xml + async + Digest auth
  ts_next   — Next.js + Prisma + Zod validation + API route

These are realistic small-app shapes that test the real agent workload:
multi-file creation, dependency management, iterative build/test cycles,
and cross-file reasoning. Add your own by dropping a task module into
`bench/tasks/multitask/` — it auto-discovers.

Usage (via the bench.py CLI):
  python bench.py agentic-multitask validate go_htmx
  python bench.py agentic-multitask run go_htmx
  python bench.py agentic-multitask cohort go_htmx rust_cli ts_next

Env (via .env): ENDPOINT, MODEL, API_KEY, MAX_TOKENS, MAX_TURNS, TEMPERATURE.
"""
import importlib.util
import json
import os
import re
import sys
import time

from bench.core import client
from bench.core import tools
from bench.core.reporting import write_json, safe_label

HERE = os.path.dirname(os.path.abspath(__file__))
MULTITASK_DIR = os.path.join(HERE, "multitask")
WORK_BASE = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "work")

MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "16000"))
MAX_TURNS = int(os.environ.get("MAX_TURNS", "40"))
TEMPERATURE = float(os.environ.get("TEMPERATURE", "0"))


def load_task(name):
    path = os.path.join(MULTITASK_DIR, f"{name}.py")
    if not os.path.isfile(path):
        print(f"ERROR: no such task module: {path}", file=sys.stderr)
        sys.exit(2)
    spec = importlib.util.spec_from_file_location(f"task_{name}", path)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return {
        "TASK_NAME": m.TASK_NAME,
        "PLAN": m.PLAN,
        "STARTER": m.STARTER,
        "REFERENCE": m.REFERENCE,
        "validate": m.validate,
        "gates": m.gates,
        "TOOLS": getattr(m, "TOOLS", tools.TOOLS),
        "TOOL_IMPLS": getattr(m, "TOOL_IMPLS", tools.TOOL_IMPLS),
    }


def available_tasks():
    if not os.path.isdir(MULTITASK_DIR):
        return []
    return sorted(f[:-3] for f in os.listdir(MULTITASK_DIR)
                  if f.endswith(".py") and not f.startswith("_"))


def _materialize(files, dest):
    import shutil
    shutil.rmtree(dest, ignore_errors=True)
    os.makedirs(dest, exist_ok=True)
    for path, content in files.items():
        fp = os.path.join(dest, path)
        os.makedirs(os.path.dirname(fp), exist_ok=True)
        with open(fp, "w") as f:
            f.write(content)
    return dest


def run_solo(task, work_dir):
    tools.set_work_dir(work_dir)
    task_tools = task.get("TOOLS", tools.TOOLS)
    task_tool_impls = task.get("TOOL_IMPLS", tools.TOOL_IMPLS)
    _materialize(task["STARTER"], dest=work_dir)
    messages = [
        {"role": "system", "content":
            "You are a coding agent working in a project work dir. You are given "
            "a task plan below. Execute it step by step using the tools. You can "
            "write files, edit files, run builds and tests, and run shell commands "
            "in the work dir. Keep going until the build is green, the tests pass, "
            "and the app/binary works. Use the tools to verify your own progress — "
            "don't guess; compile and test.\n\n"
            "When you are confident the task is complete and tests pass, say "
            "`DONE` in your final message (no tool call) and stop.\n\n"
            "TASK PLAN:\n" + task["PLAN"]},
        {"role": "user", "content":
            "The work dir has the starter scaffold. Execute the task plan. Begin "
            "by adding any dependencies and writing the files. Use the tools to "
            "build and test as you go."},
    ]
    turns_log = []
    peak_prompt = total_completion = total_reasoning = 0
    finish_reasons = []
    t0 = time.time()
    turns = 0
    done_said = False
    no_tool_streak = 0
    for turn in range(1, MAX_TURNS + 1):
        turns = turn
        body = {"model": client.MODEL, "messages": messages, "max_tokens": MAX_TOKENS,
                "temperature": TEMPERATURE, "tools": task_tools}
        try:
            resp, dt = client.call_model(messages, tools=task_tools, max_tokens=MAX_TOKENS,
                                         temperature=TEMPERATURE)
        except Exception as e:
            print(f"  turn {turn}: REQUEST ERROR {e!r}")
            turns_log.append({"turn": turn, "error": repr(e)})
            break
        u = client.usage(resp)
        peak_prompt = max(peak_prompt, u["prompt_tokens"])
        total_completion += u["completion_tokens"]
        total_reasoning += u["reasoning_tokens"]
        choice = resp["choices"][0]
        fr = choice.get("finish_reason")
        finish_reasons.append(fr)
        msg = choice["message"]
        tcs = msg.get("tool_calls") or []
        content = msg.get("content") or ""
        tlog = {"turn": turn, "dt": round(dt, 1), "prompt_tok": u["prompt_tokens"],
                "completion_tok": u["completion_tokens"], "reasoning_tok": u["reasoning_tokens"],
                "finish_reason": fr, "n_tool_calls": len(tcs)}
        turns_log.append(tlog)
        tc_summary = ",".join((tc["function"]["name"] for tc in tcs)) or "(no tool call)"
        print(f"  turn {turn:2d}: {dt:5.1f}s ptok={u['prompt_tokens']:6d} ctok={u['completion_tokens']:5d} "
              f"rtok={u['reasoning_tokens']:5d} fr={fr:8s} tools=[{tc_summary}]")
        if not tcs and "DONE" in content.upper():
            done_said = True
            tlog["done"] = True
            break
        if not tcs:
            no_tool_streak += 1
            if no_tool_streak >= 2:
                print("  -> giving up after 2 no-tool turns")
                break
            messages.append({"role": "assistant", "content": content})
            messages.append({"role": "user", "content":
                "You must either call a tool to keep working, or say DONE if "
                "you believe the task is complete and tests pass."})
            continue
        no_tool_streak = 0
        a_msg = {"role": "assistant", "content": content}
        if tcs: a_msg["tool_calls"] = tcs
        messages.append(a_msg)
        for tc in tcs:
            fn = tc["function"]
            name = fn["name"]
            try:
                args = json.loads(fn["arguments"] or "{}")
            except json.JSONDecodeError as e:
                args = {}
                result = f"ERROR: invalid JSON args ({e}); re-emit the {name} call."
            else:
                impl = task_tool_impls.get(name)
                result = impl(args) if impl else f"ERROR: unknown tool {name!r}"
            if len(result) > 6000:
                result = result[:6000] + f"\n...[truncated, {len(result)} total]"
            messages.append({"role": "tool", "tool_call_id": tc.get("id", ""), "content": result})
            tlog.setdefault("tool_results", []).append({
                "tool": name, "result_head": result[:120].replace("\n", " ")})
    wall = time.time() - t0
    gate_results = task["gates"](work_dir)
    fr_counter = {}
    for f in finish_reasons:
        if f: fr_counter[f] = fr_counter.get(f, 0) + 1
    rec = {
        "architecture": "solo",
        "task": task["TASK_NAME"],
        "model": client.MODEL,
        "max_tokens_per_turn": MAX_TOKENS,
        "max_turns": MAX_TURNS,
        "turns_used": turns,
        "done_said": done_said,
        "wall_s": round(wall, 1),
        "peak_prompt_tokens": peak_prompt,
        "total_completion_tokens": total_completion,
        "total_reasoning_tokens": total_reasoning,
        "finish_reasons": fr_counter,
        "checkpoints": gate_results,
        "all_pass": all(gate_results.values()),
        "turns_log": turns_log,
    }
    return rec


def _report(rec):
    tag = safe_label(f"{rec['task']}__{rec['model']}")
    print("\n===================== V2-HARD RESULT =====================")
    print(f"task               : {rec['task']}")
    print(f"model              : {rec['model']}")
    print(f"turns used         : {rec['turns_used']}/{rec['max_turns']}")
    print(f"done_said          : {rec['done_said']}")
    print(f"wall clock         : {rec['wall_s']}s ({rec['wall_s']/60:.1f} min)")
    print(f"peak prompt tok    : {rec['peak_prompt_tokens']}")
    print(f"total completion   : {rec['total_completion_tokens']}")
    print(f"total reasoning    : {rec['total_reasoning_tokens']}")
    print(f"finish reasons     : {rec['finish_reasons']}")
    print()
    print("--- gates ---")
    for k, v in rec["checkpoints"].items():
        print(f"  {k:16s}: {'PASS' if v else 'fail'}")
    print(f"  => {'ALL PASS' if rec['all_pass'] else 'INCOMPLETE'}")
    write_json(rec, tag, prefix="v2hard")
    print()


def validate(task_name):
    task = load_task(task_name)
    work_dir = os.path.join(WORK_BASE, f"multitask_validate_{task_name}")
    print(f"=== VALIDATION: {task_name} reference app ===")
    task["validate"](work_dir)


def run(task_name):
    task = load_task(task_name)
    work_dir = os.path.join(WORK_BASE, f"multitask_solo_{task_name}")
    rec = run_solo(task, work_dir)
    _report(rec)


def cohort(*task_names):
    results = []
    for name in task_names:
        print(f"\n{'='*60}\n=== COHORT: {name} ===\n{'='*60}")
        task = load_task(name)
        work_dir = os.path.join(WORK_BASE, f"multitask_cohort_{name}")
        rec = run_solo(task, work_dir)
        _report(rec)
        results.append(rec)
    print(f"\n{'='*60}\n=== COHORT SUMMARY ===\n{'='*60}")
    print(f"{'task':16s} {'gates':>6s} {'wall':>8s} {'turns':>6s} {'comp tok':>9s} {'rsn tok':>9s}")
    for r in results:
        g = 'PASS' if r['all_pass'] else 'fail'
        print(f"{r['task']:16s} {g:>6s} {r['wall_s']/60:>7.1f}m {r['turns_used']:>6} {r['total_completion_tokens']:>9,} {r['total_reasoning_tokens']:>9,}")


def main():
    if len(sys.argv) < 2:
        print(f"usage: python bench.py agentic-multitask [validate|run|cohort] <task>...")
        print(f"available tasks: {', '.join(available_tasks())}")
        sys.exit(1)
    cmd = sys.argv[1]
    if cmd == "validate":
        validate(sys.argv[2])
    elif cmd == "run":
        run(sys.argv[2])
    elif cmd == "cohort":
        cohort(*sys.argv[2:])
    else:
        print(f"unknown command: {cmd}")
        sys.exit(1)


if __name__ == "__main__":
    main()