#!/usr/bin/env python3
"""LLM bench runner — OpenAI-compatible chat.completions + tools + multi-turn.

Runs each task JSON under <tasks-dir>/<category>/*.json against an OpenAI-compatible
endpoint (mlx_lm.server, vLLM, ollama, etc). Saves per-task response JSON and
appends compact row to run_log.jsonl. Stdlib-only (urllib + json).

Canonical: dock/scripts/llm-bench-run.py
Task JSON schema: experiments/llm-benchmarks/.../tasks/<cat>/<id>_<slug>.json
Spec: 03-benchmarks/specs/qwen3.6-35b-a3b-probe.yaml

Example (Mac Row 4):
  python3 scripts/llm-bench-run.py \\
    --endpoint http://127.0.0.1:8000/v1 \\
    --model ~/ai/weights/qwen3.6-35b-a3b-4bit-dwq/ \\
    --model-id qwen3.6_4bit_dwq \\
    --tasks-dir ~/ai/experiments/llm-benchmarks/.../tasks \\
    --output-dir ~/ai/experiments/llm-benchmarks/.../qwen3.6_4bit_dwq \\
    --run-log ~/ai/experiments/llm-benchmarks/.../run_log.jsonl
"""
from __future__ import annotations
import argparse
import json
import re
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path


def http_post_json(url, body, headers, timeout):
    data = json.dumps(body).encode()
    req = urllib.request.Request(
        url, data=data, headers={**headers, "Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode()), r.getcode()


def run_task(task, endpoint, model, api_key, timeout, max_tokens_override=None):
    url = endpoint.rstrip("/") + "/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}

    messages = list(task.get("messages") or [])
    system = task.get("system")
    if system:
        messages = [{"role": "system", "content": system}] + messages

    gen_cfg = dict(task.get("generation_config") or {})
    if max_tokens_override is not None:
        gen_cfg["max_tokens"] = max_tokens_override

    body = {
        "model": model,
        "messages": messages,
        **gen_cfg,
    }
    tools = task.get("tools") or []
    if tools:
        body["tools"] = tools
        body["tool_choice"] = "auto"

    turns = []
    t0 = time.perf_counter()
    try:
        r1, _ = http_post_json(url, body, headers, timeout)
    except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError) as e:
        return {
            "ok": False,
            "error": f"{type(e).__name__}: {e}",
            "wall_s": round(time.perf_counter() - t0, 3),
            "turns": [],
        }
    turns.append(r1)

    grading = task.get("grading") or {}
    multi_turn = bool(grading.get("multi_turn"))
    mock = grading.get("tool_response_mock") or {}

    first_msg = ((r1.get("choices") or [{}])[0]).get("message") or {}
    tool_calls = first_msg.get("tool_calls") or []

    if multi_turn and tool_calls and mock:
        messages.append(
            {"role": "assistant", "content": first_msg.get("content"), "tool_calls": tool_calls}
        )
        for tc in tool_calls:
            fn_name = ((tc.get("function") or {}).get("name")) or ""
            mock_result = mock.get(fn_name, "(no mock response)")
            if not isinstance(mock_result, str):
                mock_result = json.dumps(mock_result, ensure_ascii=False)
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tc.get("id", ""),
                    "content": mock_result,
                }
            )
        body2 = {**body, "messages": messages}
        try:
            r2, _ = http_post_json(url, body2, headers, timeout)
            turns.append(r2)
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError) as e:
            return {
                "ok": False,
                "error": f"second_turn: {type(e).__name__}: {e}",
                "wall_s": round(time.perf_counter() - t0, 3),
                "turns": turns,
            }

    wall_s = round(time.perf_counter() - t0, 3)
    prompt_tokens = sum((t.get("usage") or {}).get("prompt_tokens", 0) for t in turns)
    completion_tokens = sum(
        (t.get("usage") or {}).get("completion_tokens", 0) for t in turns
    )
    tok_s = round(completion_tokens / wall_s, 2) if wall_s > 0 and completion_tokens else None

    final_choice = (turns[-1].get("choices") or [{}])[0]
    final_msg = final_choice.get("message") or {}
    final_finish = final_choice.get("finish_reason")
    # mlx_lm.server / vLLM reasoning model: thinking 은 별 필드. 둘 다 캡쳐.
    reasoning_text = final_msg.get("reasoning") or final_msg.get("reasoning_content") or ""
    content_text = final_msg.get("content") or ""
    return {
        "ok": True,
        "wall_s": wall_s,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "tok_s_decode": tok_s,
        "finish_reason": final_finish,
        "turns": turns,
        "final_text": content_text,
        "final_reasoning": reasoning_text,
        "final_tool_calls": final_msg.get("tool_calls") or [],
        **auto_grade(task, turns),
    }


def auto_grade(task, turns):
    grading = task.get("grading") or {}
    category = task.get("category")
    out = {"auto_pass": None, "auto_notes": ""}
    if not turns:
        return out

    first_msg = ((turns[0].get("choices") or [{}])[0]).get("message") or {}
    first_tool_calls = first_msg.get("tool_calls") or []
    first_called = [(tc.get("function") or {}).get("name") for tc in first_tool_calls]

    if grading.get("no_tool_expected"):
        if first_tool_calls:
            out["auto_pass"] = False
            out["auto_notes"] = f"unexpected tool call: {first_called}"
        else:
            out["auto_pass"] = True
    else:
        expected = grading.get("expected_tool_call")
        if expected:
            accepted = {expected, *(grading.get("alternative_tool_calls") or [])}
            if any(c in accepted for c in first_called):
                must = grading.get("expected_args_must_contain") or []
                if must and first_tool_calls:
                    args_raw = first_tool_calls[0].get("function", {}).get("arguments", "")
                    if not isinstance(args_raw, str):
                        args_raw = json.dumps(args_raw, ensure_ascii=False)
                    missing = [k for k in must if k not in args_raw]
                    if missing:
                        out["auto_pass"] = False
                        out["auto_notes"] = f"missing args keys: {missing}"
                    else:
                        out["auto_pass"] = True
                else:
                    out["auto_pass"] = True
            else:
                out["auto_pass"] = False
                out["auto_notes"] = f"expected {expected}, got {first_called}"

    ref = task.get("reference_answer")
    if category == "reasoning" and ref:
        final_msg = ((turns[-1].get("choices") or [{}])[0]).get("message", {})
        # reasoning 모델: content 또는 reasoning 필드에 답이 있을 수 있음
        merged = (final_msg.get("content") or "") + "\n" + (
            final_msg.get("reasoning") or final_msg.get("reasoning_content") or ""
        )
        key_tokens = re.findall(r"[\d.]+", ref)
        if key_tokens:
            hit = sum(1 for k in key_tokens if k in merged)
            out["auto_reasoning_hit"] = f"{hit}/{len(key_tokens)}"

    return out


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--endpoint", required=True, help="OpenAI-compatible base URL")
    p.add_argument("--model", required=True, help="model name or path sent to server")
    p.add_argument("--model-id", required=True, help="short id for output subdir")
    p.add_argument("--tasks-dir", required=True, type=Path)
    p.add_argument("--output-dir", required=True, type=Path)
    p.add_argument("--run-log", required=True, type=Path)
    p.add_argument(
        "--categories",
        default="tool_call,reasoning,translation,combined,safety",
    )
    p.add_argument("--api-key", default="")
    p.add_argument("--timeout", type=float, default=300.0)
    p.add_argument(
        "--max-tokens",
        type=int,
        default=None,
        help="override generation_config.max_tokens for all tasks (reasoning 모델은 4096-8192 권장)",
    )
    p.add_argument("--overwrite", action="store_true")
    p.add_argument("--dry-run", action="store_true")
    p.add_argument(
        "--no-think",
        action="store_true",
        help="Qwen3/3.5/3.6 계열: 마지막 user message 에 '/no_think' soft switch 추가 (reasoning 모드 OFF)",
    )
    args = p.parse_args()

    cats = [c.strip() for c in args.categories.split(",") if c.strip()]
    tasks = []
    for cat in cats:
        cat_dir = args.tasks_dir / cat
        if not cat_dir.is_dir():
            print(f"[warn] missing category dir: {cat_dir}", file=sys.stderr)
            continue
        for f in sorted(cat_dir.glob("*.json")):
            with open(f) as fh:
                t = json.load(fh)
            t["_path"] = str(f)
            tasks.append(t)

    if not tasks:
        print("no tasks found", file=sys.stderr)
        return 1

    print(f"[info] {len(tasks)} tasks loaded")
    for t in tasks:
        print(f"       - {t.get('category')}/{t.get('id')}_{t.get('slug')}")
    if args.dry_run:
        return 0

    args.output_dir.mkdir(parents=True, exist_ok=True)
    args.run_log.parent.mkdir(parents=True, exist_ok=True)

    for task in tasks:
        cat = task["category"]
        tid = task["id"]
        out_cat_dir = args.output_dir / cat
        out_cat_dir.mkdir(parents=True, exist_ok=True)
        out_file = out_cat_dir / f"{tid}_response.json"
        if out_file.exists() and not args.overwrite:
            print(f"[skip] {cat}/{tid} (exists — pass --overwrite to redo)")
            continue

        print(f"[run]  {cat}/{tid} ({task.get('slug')})", flush=True)
        if args.no_think:
            # /no_think soft switch — Qwen3 계열 reasoning mode OFF
            msgs = task.get("messages") or []
            for m in reversed(msgs):
                if m.get("role") == "user":
                    m["content"] = (m.get("content") or "") + "\n/no_think"
                    break
        result = run_task(
            task,
            endpoint=args.endpoint,
            model=args.model,
            api_key=args.api_key,
            timeout=args.timeout,
            max_tokens_override=args.max_tokens,
        )

        payload = {
            "model_id": args.model_id,
            "model": args.model,
            "task_id": tid,
            "category": cat,
            "slug": task.get("slug"),
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            **result,
        }
        with open(out_file, "w") as fh:
            json.dump(payload, fh, ensure_ascii=False, indent=2)

        log_row = {
            "ts": payload["timestamp"],
            "model_id": args.model_id,
            "category": cat,
            "task_id": tid,
            "ok": result.get("ok"),
            "error": result.get("error"),
            "wall_s": result.get("wall_s"),
            "prompt_tokens": result.get("prompt_tokens"),
            "completion_tokens": result.get("completion_tokens"),
            "tok_s_decode": result.get("tok_s_decode"),
            "auto_pass": result.get("auto_pass"),
            "auto_notes": result.get("auto_notes"),
            "auto_reasoning_hit": result.get("auto_reasoning_hit"),
            "finish_reason": result.get("finish_reason"),
        }
        with open(args.run_log, "a") as fh:
            fh.write(json.dumps(log_row, ensure_ascii=False) + "\n")

        print(
            f"       wall={result.get('wall_s')}s tok={result.get('completion_tokens')} "
            f"tok/s={result.get('tok_s_decode')} auto_pass={result.get('auto_pass')}"
        )

    print("[done]")
    return 0


if __name__ == "__main__":
    sys.exit(main())