#!/usr/bin/env python3
"""Multi-turn agentic harness for vinland dual-MI50.

Every prior round was single-shot. A bot in a Hermes/Claw cluster runs
many turns with tool calls, accumulating context and recovering from
its own earlier mistakes. This harness measures THAT: it puts each
model in a real tool-use loop and scores the whole trajectory, not one
answer.

How it works:
  - Each task gives the model a goal and a small tool set.
  - The model must emit a tool call as JSON. The harness executes it in
    a sandboxed scratch dir, appends the result to the conversation, and
    asks the model for its next move.
  - This repeats until the model calls `done`, a turn cap is hit, or a
    per-call timeout fires. Nothing can hang the overnight run: every
    model call and every code execution has a hard timeout.
  - Success is checked by an objective grader per task (files exist,
    code runs, output matches).

Tools available to the model:
  read_file(path)            -> file contents
  write_file(path, content)  -> writes into the scratch dir
  run_python(path)           -> executes a .py in the scratch dir, returns stdout/stderr
  list_dir()                 -> lists the scratch dir
  done(summary)              -> ends the task

Metrics per (model, task): solved (bool), turns used, terminated cleanly
(called done vs hit cap vs runaway/timeout), tool-call validity rate,
wasted turns (malformed calls). These are the swarm-relevant numbers.

Writes one logfile per model to results_agentic/ with the full
trajectory, plus a machine-readable summary block per task. Resumable.

Usage:
    cd ~/workarea/ash/llm_knowledge_test
    python3 knowledge_bench5.py
"""

import glob
import json
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import urllib.request

# ---------------------------------------------------------------- config

SERVER_BIN = "/opt/llama.cpp-rocm/llama-server"
MODEL_DIR = os.path.expanduser("~/.cache/llama.cpp")
RESULTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results_agentic")
SCRATCH_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "agentic_scratch")
PORT = 8090
HOST = "127.0.0.1"
BASE_URL = f"http://{HOST}:{PORT}"

LOAD_TIMEOUT_S = 1200
TURN_TIMEOUT_S = 400        # per model turn
CODE_EXEC_TIMEOUT_S = 20    # per run_python
MAX_TURNS = 16              # hard trajectory cap
MAX_TOKENS = 3072           # per turn
CTX = 32768                 # long trajectories need context

# The decided pool + two worker candidates whose long-horizon
# termination is exactly what we want to stress.
MODELS = [
    ("qwen3.6-27b",         "Qwen_Qwen3.6-27B-Q4_K_M.gguf",                              "full"),
    ("gpt-oss-120b",        "openai_gpt-oss-120b-Q4_K_M/*00001-of-*.gguf",               "auto"),
    ("north-mini-code",     "North-Mini-Code-1.0-UD-Q4_K_XL.gguf",                       "full"),
    ("devstral-small2-24b", "mistralai_Devstral-Small-2-24B-Instruct-2512-Q4_K_M.gguf",  "full"),
    ("gpt-oss-20b",         "openai_gpt-oss-20b-Q5_K_M.gguf",                            "full"),
    ("qwen3.5-9b",          "Qwen_Qwen3.5-9B-Q4_K_M.gguf",                               "full"),
]

SYSTEM_PROMPT = """You are an autonomous agent that completes tasks using tools. You work in a scratch directory.

On each turn you MUST respond with exactly one tool call as a single JSON object and NOTHING else (no prose, no code fence, no explanation). The JSON must have this shape:

{"tool": "<name>", "args": { ... }}

Available tools:
- {"tool": "list_dir", "args": {}} - list files in the scratch dir
- {"tool": "read_file", "args": {"path": "name.txt"}} - read a file
- {"tool": "write_file", "args": {"path": "name.py", "content": "..."}} - write/overwrite a file
- {"tool": "run_python", "args": {"path": "name.py"}} - run a python file, returns stdout and stderr
- {"tool": "done", "args": {"summary": "..."}} - call this when the task is complete

Rules:
- Respond with ONLY the JSON object, nothing before or after.
- Take one action per turn. Use the tool results to decide your next action.
- When the goal is achieved, call done. Do not keep working after success.
- If a run_python shows an error, read it, fix the file, and re-run.
- Do not give up; but do not loop forever. Call done when finished."""

# task_id, goal_prompt, grader(scratch_dir)->(bool, detail)
TASKS = [
    ("t1_fizzbuzz",
     "Create a file `fizzbuzz.py` in the scratch dir containing a function "
     "`fb(n)` that returns a list of strings for 1..n: 'Fizz' for multiples of "
     "3, 'Buzz' for multiples of 5, 'FizzBuzz' for both, else the number as a "
     "string. Then run it to verify it works. When the file exists and runs "
     "without error, call done."),

    ("t2_fix_bug",
     "A file `buggy.py` already exists in the scratch dir. It has a function "
     "`has_dup(xs)` that should return True if the list has any duplicates, but "
     "it has a bug. Read it, fix the bug in place, and run it to confirm the "
     "included tests pass. Then call done."),

    ("t3_pipeline",
     "Build a two-step pipeline. First create `gen.py` that writes the numbers "
     "1 through 20, one per line, to a file called `nums.txt`. Run it. Then "
     "create `sum.py` that reads `nums.txt` and prints the sum of the numbers. "
     "Run it and confirm it prints 210. Then call done."),

    ("t4_debug_loop",
     "Create `stats.py` with a function `mean_std(xs)` returning a tuple of the "
     "arithmetic mean and the population standard deviation of a list of "
     "numbers. It must print mean_std([2,4,4,4,5,5,7,9]) which should be "
     "(5.0, 2.0). Run it and confirm the output is exactly (5.0, 2.0). If not, "
     "fix and re-run until correct, then call done."),

    ("t5_parse",
     "Create `parse.py` with a function `parse_kv(text)` that parses lines of "
     "the form 'key=value' (one per line, ignore blank lines) into a dict, with "
     "values converted to int when they are all digits. Include a test in the "
     "file: parse_kv('a=1\\nb=hello\\n\\nc=42') must equal {'a':1,'b':'hello',"
     "'c':42}. Run it, assert the test passes, then call done."),
]

COMPLETE_MARKER = "=== RUN COMPLETE ==="

# ---------------------------------------------------------------- graders


def _run_py_capture(scratch, filename):
    path = os.path.join(scratch, filename)
    if not os.path.exists(path):
        return None
    try:
        r = subprocess.run([sys.executable, path], capture_output=True,
                           text=True, timeout=CODE_EXEC_TIMEOUT_S, cwd=scratch)
        return r
    except Exception:
        return None


def grade_fizzbuzz(scratch):
    path = os.path.join(scratch, "fizzbuzz.py")
    if not os.path.exists(path):
        return False, "fizzbuzz.py missing"
    probe = ("import fizzbuzz\n"
             "assert fizzbuzz.fb(5)==['1','2','Fizz','4','Buzz'], fizzbuzz.fb(5)\n"
             "assert fizzbuzz.fb(15)[-1]=='FizzBuzz'\n"
             "print('GRADE_OK')\n")
    return _probe(scratch, probe)


def grade_fix_bug(scratch):
    probe = ("import importlib.util,sys\n"
             "spec=importlib.util.spec_from_file_location('b','buggy.py')\n"
             "m=importlib.util.module_from_spec(spec);spec.loader.exec_module(m)\n"
             "assert m.has_dup([1,2,3,1])==True\n"
             "assert m.has_dup([1,2,3])==False\n"
             "print('GRADE_OK')\n")
    return _probe(scratch, probe)


def grade_pipeline(scratch):
    nums = os.path.join(scratch, "nums.txt")
    if not os.path.exists(nums):
        return False, "nums.txt missing"
    r = _run_py_capture(scratch, "sum.py")
    if r is None:
        return False, "sum.py missing or failed"
    if "210" in (r.stdout or ""):
        return True, "prints 210"
    return False, f"sum.py stdout={r.stdout.strip()[:60]}"


def grade_debug_loop(scratch):
    r = _run_py_capture(scratch, "stats.py")
    if r is None:
        return False, "stats.py missing or failed"
    out = (r.stdout or "").replace(" ", "")
    if "(5.0,2.0)" in out:
        return True, "correct output"
    return False, f"stdout={r.stdout.strip()[:60]}"


def grade_parse(scratch):
    probe = ("import parse\n"
             "assert parse.parse_kv('a=1\\nb=hello\\n\\nc=42')=={'a':1,'b':'hello','c':42}\n"
             "print('GRADE_OK')\n")
    return _probe(scratch, probe)


def _probe(scratch, probe_code):
    with tempfile.NamedTemporaryFile("w", suffix=".py", dir=scratch,
                                     delete=False) as f:
        f.write(probe_code)
        p = f.name
    try:
        r = subprocess.run([sys.executable, os.path.basename(p)],
                           capture_output=True, text=True,
                           timeout=CODE_EXEC_TIMEOUT_S, cwd=scratch)
        if "GRADE_OK" in (r.stdout or ""):
            return True, "graded pass"
        return False, ((r.stderr or "").strip().splitlines() or [""])[-1][:100]
    except Exception as e:
        return False, str(e)[:100]
    finally:
        try:
            os.unlink(p)
        except OSError:
            pass


GRADERS = {
    "t1_fizzbuzz": grade_fizzbuzz,
    "t2_fix_bug": grade_fix_bug,
    "t3_pipeline": grade_pipeline,
    "t4_debug_loop": grade_debug_loop,
    "t5_parse": grade_parse,
}

# files that must pre-exist in a task's scratch before the model starts
BUGGY_PY = '''def has_dup(xs):
    seen = set()
    for x in xs:
        if x in seen:
            return False
        seen.add(x)
    return True

if __name__ == "__main__":
    assert has_dup([1,2,3,1]) == True
    assert has_dup([1,2,3]) == False
    print("tests pass")
'''


def seed_task(task_id, scratch):
    if task_id == "t2_fix_bug":
        with open(os.path.join(scratch, "buggy.py"), "w") as f:
            f.write(BUGGY_PY)


# ---------------------------------------------------------------- tools


def tool_list_dir(scratch, args):
    return "\n".join(sorted(os.listdir(scratch))) or "(empty)"


def _safe_path(scratch, rel):
    # confine to scratch dir; reject traversal
    full = os.path.normpath(os.path.join(scratch, rel))
    if not full.startswith(os.path.abspath(scratch)):
        return None
    return full


def tool_read_file(scratch, args):
    p = _safe_path(scratch, args.get("path", ""))
    if not p or not os.path.exists(p):
        return f"ERROR: file not found: {args.get('path')}"
    try:
        with open(p) as f:
            return f.read()[:4000]
    except Exception as e:
        return f"ERROR: {e}"


def tool_write_file(scratch, args):
    p = _safe_path(scratch, args.get("path", ""))
    if not p:
        return "ERROR: invalid path"
    try:
        with open(p, "w") as f:
            f.write(args.get("content", ""))
        return f"wrote {args.get('path')} ({len(args.get('content',''))} bytes)"
    except Exception as e:
        return f"ERROR: {e}"


def tool_run_python(scratch, args):
    p = _safe_path(scratch, args.get("path", ""))
    if not p or not os.path.exists(p):
        return f"ERROR: file not found: {args.get('path')}"
    try:
        r = subprocess.run([sys.executable, os.path.basename(p)],
                           capture_output=True, text=True,
                           timeout=CODE_EXEC_TIMEOUT_S, cwd=scratch)
        out = f"exit={r.returncode}\nstdout:\n{r.stdout[:2000]}\nstderr:\n{r.stderr[:1000]}"
        return out
    except subprocess.TimeoutExpired:
        return "ERROR: execution timed out"
    except Exception as e:
        return f"ERROR: {e}"


TOOLS = {
    "list_dir": tool_list_dir,
    "read_file": tool_read_file,
    "write_file": tool_write_file,
    "run_python": tool_run_python,
}

# ---------------------------------------------------------------- llm io


def http_json(url, payload, timeout):
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode())


def wait_healthy(proc, deadline_s):
    start = time.time()
    while time.time() - start < deadline_s:
        if proc.poll() is not None:
            return False
        try:
            if http_json(f"{BASE_URL}/health", None, 5).get("status") == "ok":
                return True
        except Exception:
            pass
        time.sleep(5)
    return False


def chat(messages):
    payload = {"messages": messages, "temperature": 0.0,
               "max_tokens": MAX_TOKENS}
    resp = http_json(f"{BASE_URL}/v1/chat/completions", payload, TURN_TIMEOUT_S)
    msg = resp.get("choices", [{}])[0].get("message", {})
    return (msg.get("content") or "").strip()


def parse_tool_call(text):
    """Extract the first JSON tool call. Tolerates a code fence or stray
    prose around it. Returns (tool, args) or (None, reason)."""
    s = text
    if "```" in s:
        m = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", s, re.DOTALL)
        if m:
            s = m.group(1)
    # find first balanced-looking JSON object containing "tool"
    start = s.find("{")
    while start != -1:
        depth = 0
        for i in range(start, len(s)):
            if s[i] == "{":
                depth += 1
            elif s[i] == "}":
                depth -= 1
                if depth == 0:
                    cand = s[start:i + 1]
                    try:
                        obj = json.loads(cand)
                        if isinstance(obj, dict) and "tool" in obj:
                            return obj.get("tool"), obj.get("args", {})
                    except Exception:
                        pass
                    break
        start = s.find("{", start + 1)
    return None, "no valid tool JSON found"


# ---------------------------------------------------------------- runner


def run_task(model_name, task_id, goal, logf):
    scratch = os.path.join(SCRATCH_ROOT, model_name, task_id)
    if os.path.exists(scratch):
        shutil.rmtree(scratch)
    os.makedirs(scratch, exist_ok=True)
    seed_task(task_id, scratch)

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"TASK: {goal}\n\nBegin. Respond with one tool call as JSON."},
    ]

    turns = 0
    bad_calls = 0
    terminated = "capped"     # capped | done | error
    called_done = False

    logf.write("=" * 72 + "\n")
    logf.write(f"TASK {task_id}\n{goal}\n\n")
    logf.flush()

    while turns < MAX_TURNS:
        turns += 1
        try:
            reply = chat(messages)
        except Exception as e:
            logf.write(f"[turn {turns}] MODEL ERROR: {e}\n")
            terminated = "error"
            break

        tool, args = parse_tool_call(reply)
        short = reply.replace("\n", " ")[:200]
        logf.write(f"[turn {turns}] raw: {short}\n")

        if tool is None:
            bad_calls += 1
            logf.write(f"[turn {turns}] BAD CALL: {args}\n")
            messages.append({"role": "assistant", "content": reply})
            messages.append({"role": "user", "content":
                "That was not a valid tool call. Respond with ONLY a JSON "
                'object like {"tool":"list_dir","args":{}}. One tool call, no other text.'})
            if bad_calls >= 4:
                logf.write(f"[turn {turns}] too many bad calls, aborting task\n")
                terminated = "error"
                break
            continue

        if tool == "done":
            called_done = True
            terminated = "done"
            logf.write(f"[turn {turns}] DONE: {args.get('summary','')[:200]}\n")
            break

        fn = TOOLS.get(tool)
        if fn is None:
            bad_calls += 1
            result = f"ERROR: unknown tool '{tool}'"
        else:
            result = fn(scratch, args if isinstance(args, dict) else {})

        logf.write(f"[turn {turns}] tool={tool} -> {str(result)[:300]}\n")
        messages.append({"role": "assistant", "content": reply})
        messages.append({"role": "user", "content": f"TOOL RESULT:\n{result}"})
        logf.flush()

    solved, detail = GRADERS[task_id](scratch)
    logf.write(f"RESULT: solved={solved} ({detail}) | turns={turns} "
               f"terminated={terminated} called_done={called_done} "
               f"bad_calls={bad_calls}\n\n")
    logf.flush()
    return {"task": task_id, "solved": solved, "turns": turns,
            "terminated": terminated, "called_done": called_done,
            "bad_calls": bad_calls}


def resolve_model_path(pattern):
    full = os.path.join(MODEL_DIR, pattern)
    if "*" in full:
        m = sorted(glob.glob(full))
        return m[0] if m else None
    if os.path.isdir(full):
        s = sorted(glob.glob(os.path.join(full, "*00001-of-*.gguf")))
        if s:
            return s[0]
        g = sorted(glob.glob(os.path.join(full, "*.gguf")))
        return g[0] if g else None
    return full if os.path.exists(full) else None


def build_cmd(model_path, fit_mode):
    cmd = [SERVER_BIN, "-m", model_path, "-sm", "layer", "-c", str(CTX),
           "-t", "12", "--host", HOST, "--port", str(PORT),
           "--jinja", "--reasoning-format", "auto"]
    if fit_mode == "full":
        cmd += ["-ngl", "99"]
    return cmd


def stop_server(proc):
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=30)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
    time.sleep(10)


def run_model(name, model_path, fit_mode):
    logfile = os.path.join(RESULTS_DIR, f"{name}.log")
    serverlog = os.path.join(RESULTS_DIR, f"{name}.server.log")
    if os.path.exists(logfile):
        with open(logfile) as f:
            if COMPLETE_MARKER in f.read():
                print(f"[skip] {name}: already complete")
                return

    print(f"[load] {name} ({fit_mode})")
    env = dict(os.environ, HIP_VISIBLE_DEVICES="0,1")
    proc = subprocess.Popen(build_cmd(model_path, fit_mode),
                            stdout=open(serverlog, "w"),
                            stderr=subprocess.STDOUT, env=env)
    try:
        if not wait_healthy(proc, LOAD_TIMEOUT_S):
            print(f"[FAIL] {name}: not healthy, see {serverlog}")
            stop_server(proc)
            with open(logfile, "w") as f:
                f.write(f"MODEL: {name}\nFAILED TO LOAD\n")
            return

        summary = []
        t0 = time.time()
        with open(logfile, "w") as f:
            f.write(f"MODEL: {name}\nPATH: {model_path}\n")
            f.write(f"STARTED: {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            f.flush()
            for task_id, goal, *_ in [(t[0], t[1]) for t in TASKS]:
                print(f"  [{name}] {task_id}", flush=True)
                res = run_task(name, task_id, goal, f)
                summary.append(res)

            f.write("=" * 72 + "\n")
            f.write("SUMMARY (machine readable):\n")
            f.write(json.dumps(summary, indent=2) + "\n\n")
            solved = sum(1 for r in summary if r["solved"])
            clean = sum(1 for r in summary if r["terminated"] == "done")
            f.write(f"SOLVED: {solved}/{len(TASKS)} | "
                    f"CLEAN_TERMINATIONS: {clean}/{len(TASKS)} | "
                    f"TOTAL_TIME: {time.time()-t0:.0f}s\n")
            f.write(COMPLETE_MARKER + "\n")
        print(f"[done] {name}: solved {solved}/{len(TASKS)}, clean {clean}/{len(TASKS)}")
    finally:
        stop_server(proc)


def main():
    os.makedirs(RESULTS_DIR, exist_ok=True)
    os.makedirs(SCRATCH_ROOT, exist_ok=True)
    if not os.path.exists(SERVER_BIN):
        sys.exit(f"llama-server not found at {SERVER_BIN}")

    todo = []
    for name, pattern, fit_mode in MODELS:
        path = resolve_model_path(pattern)
        if path is None:
            print(f"[warn] {name}: no file matching {pattern}, skipping")
        else:
            todo.append((name, path, fit_mode))

    print(f"Agentic multi-turn run: {len(todo)} models, {len(TASKS)} tasks, "
          f"max {MAX_TURNS} turns each.")
    print(f"Results: {RESULTS_DIR}\n")

    def on_sigint(sig, frame):
        raise KeyboardInterrupt
    signal.signal(signal.SIGINT, on_sigint)

    t0 = time.time()
    for name, path, fit_mode in todo:
        try:
            run_model(name, path, fit_mode)
        except KeyboardInterrupt:
            print("\nInterrupted. Rerun to resume; completed models skipped.")
            break
    print(f"\nElapsed: {(time.time()-t0)/60:.0f} min")
    print("Hand results_agentic/*.log to the judge.")


if __name__ == "__main__":
    main()