#!/usr/bin/env python3
"""
Phase 1c Tier-2 Behavioral & Agentic Evaluation Suite (20 automated tasks across 5 categories).
Evaluates K512_unpruned, K320_reap_hybrid, K320_reap_l2_sum, K256_reap_l2_sum, and K224_reap_l2_sum via llama-server on A100.
"""
import json
import os
import re
import subprocess
import time
import urllib.request
from pathlib import Path

RUBRIC_VERSION = "v2_opus_1c_fixes"
SERVER_BIN = "/content/llama.cpp/build/bin/llama-server"
SHARD1 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00001-of-00002.gguf"
SHARD2 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00002-of-00002.gguf"

PERSONA_SYS = (
    "you are a persistent agent operating inside a coding and systems harness with a python tool (`python`). "
    "speak in lowercase, use only emoticons or japanese kaomoji (no emoji, no officecore, no gen z slang). "
    "when context or the user requires a strict machine-readable format (like raw JSON or a config file), "
    "output the exact requested format without kaomoji leaking into the structured payload."
)

TOOLS_SPEC = [
    {
        "type": "function",
        "function": {
            "name": "python",
            "description": "Execute Python or run system commands via run(program, *args) in the persistent session.",
            "parameters": {
                "type": "object",
                "properties": {"code": {"type": "string"}},
                "required": ["code"],
            },
        },
    }
]

TASKS = [
    # --- Category 1: ops_tool_call (4 tasks) ---
    {
        "id": "ops_1_wg_mtu",
        "category": "ops_tool_call",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "my wireguard tunnel `wg0` to `10.77.0.1` pings fine with small packets, but `rsync` hangs immediately on large transfers. use your `python` tool to inspect `wg0` link MTU and test a 1380-byte do-not-fragment ping."},
        ],
        "use_tools": True,
        "expect_keywords": ["wg0"],
    },
    {
        "id": "ops_2_systemd_oom",
        "category": "ops_tool_call",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "our `restic-backup.service` systemd unit keeps getting killed with OOM even though the machine has 8 GB of free swap. use your `python` tool to check `systemctl show restic-backup.service` for `MemoryMax` and `MemorySwapMax`."},
        ],
        "use_tools": True,
        "expect_keywords": ["systemctl", "restic-backup"],
    },
    {
        "id": "ops_3_nix_gcroots",
        "category": "ops_tool_call",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "`nix-collect-garbage` barely freed any disk space on `/nix/store`. use your `python` tool to list non-system live GC roots with `nix-store --gc --print-roots`."},
        ],
        "use_tools": True,
        "expect_any_keywords": ["nix-store", "gcroots"],
    },
    {
        "id": "ops_4_port_listener",
        "category": "ops_tool_call",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "port `8443` is already bound on localhost and blocking `caddy` from starting. use your `python` tool to check which process owns port `8443` using `ss` or `lsof`."},
        ],
        "use_tools": True,
        "expect_keywords": ["8443"],
        "expect_any_keywords": ["ss", "lsof", "netstat", "fuser", "/proc"],
    },

    # --- Category 2: voice_override_json (4 tasks) ---
    {
        "id": "json_1_fw_rule",
        "category": "voice_override_json",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "reply with ONLY a raw JSON object (no markdown fences, no commentary, no kaomoji) with keys `table` (\"inet filter\"), `chain` (\"forward\"), `tcp_mss` (1380), and `enabled` (true)."},
        ],
        "use_tools": False,
        "expect_json": {"table": "inet filter", "chain": "forward", "tcp_mss": 1380, "enabled": True},
    },
    {
        "id": "json_2_systemd_limits",
        "category": "voice_override_json",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "output ONLY a valid JSON object (no markdown blocks, no prose) containing `unit` (\"postgres.service\"), `memory_high_mb` (4096), `swap_max_mb` (0), and `oom_policy` (\"kill\")."},
        ],
        "use_tools": False,
        "expect_json": {"unit": "postgres.service", "memory_high_mb": 4096, "swap_max_mb": 0, "oom_policy": "kill"},
    },
    {
        "id": "json_3_model_budget",
        "category": "voice_override_json",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "return ONLY raw parseable JSON (no backticks, no extra text) with keys `arch` (\"qwen4exp\"), `layers` (48), `kept_experts` (320), and `top_k` (10)."},
        ],
        "use_tools": False,
        "expect_json": {"arch": "qwen4exp", "layers": 48, "kept_experts": 320, "top_k": 10},
    },
    {
        "id": "json_4_cron_spec",
        "category": "voice_override_json",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "respond with ONLY raw JSON (no code fences, no explanation) with keys `job` (\"zfs-scrub\"), `schedule` (\"0 3 * * 0\"), `pool` (\"tank\"), and `notify_on_error` (true)."},
        ],
        "use_tools": False,
        "expect_json": {"job": "zfs-scrub", "schedule": "0 3 * * 0", "pool": "tank", "notify_on_error": True},
    },

    # --- Category 3: ja_reasoning_chat (4 tasks) ---
    {
        "id": "ja_1_moe_routing",
        "category": "ja_reasoning_chat",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "Mixture-of-Experts (MoE) モデルを SSD からストリーミング実行するとき、トークンあたりの有効パラメータ数が少なくてもプリフィル（prefill）で全エキスパートが読み込まれてしまう理由と、チャンクサイズを大きくすると改善する理由を日本語で簡潔に説明して。"},
        ],
        "use_tools": False,
        "expect_ja_keyword_groups": [("エキスパート", "expert"), ("チャンク", "chunk")],
    },
    {
        "id": "ja_2_git_rebase",
        "category": "ja_reasoning_chat",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "共有ブランチにすでに push 済みのコミットに対して `git rebase` ではなく `git revert` を使うべき理由を日本語で説明して。"},
        ],
        "use_tools": False,
        "expect_ja_keyword_groups": [("履歴", "history"), ("revert",)],
    },
    {
        "id": "ja_3_sqlite_wal",
        "category": "ja_reasoning_chat",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "SQLite でデフォルトのロールバックジャーナルモードから WAL (`PRAGMA journal_mode=WAL;`) に切り替えると、並行する読み取りと書き込みがどう変わるか日本語で教えて。"},
        ],
        "use_tools": False,
        "expect_ja_keyword_groups": [("wal",), ("読み取り", "read", "reader")],
    },
    {
        "id": "ja_4_rope_sink",
        "category": "ja_reasoning_chat",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "長文コンテキストの途中で要約（compaction）を挟んで KV キャッシュを再利用するとき、先頭の attention sink トークンを固定しないと出力がループしやすくなる理由を日本語で説明して。"},
        ],
        "use_tools": False,
        "expect_ja_keyword_groups": [("キャッシュ", "cache", "kv"), ("トークン", "token", "sink")],
    },

    # --- Category 4: context_compaction_recall (4 tasks) ---
    {
        "id": "comp_1_multi_fact",
        "category": "context_compaction_recall",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "[Recall L1-2 (folded session memory: turns 1-42)]\nIn turns 1-42, dawn and I migrated the backup daemon `bunker-sync` to `/srv/bunker/config.toml`, changed its internal metrics port from `9102` to `9477`, and pinned the rule: never run `bunker-sync --prune` while `pg_basebackup` is active.\n[End Recall]\n\nhey, we're ready to wire up prometheus scraping for `bunker-sync` now. what config path, metrics port, and safety rule did we lock in earlier?"},
        ],
        "use_tools": False,
        "expect_substrings": ["/srv/bunker/config.toml", "9477", "pg_basebackup"],
    },
    {
        "id": "comp_2_override_folded",
        "category": "context_compaction_recall",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "[Recall L1-3 (folded session memory: turns 1-88)]\nWe benchmarked the M5 Air NVMe SSD at `6.80 GB/s` using 4-thread `pread` and set `cache_cap = 113` experts per layer (`7.5 GB` in `Q2_0`). However, in turn 85 we discovered Layer 0 kept 8 dead zero-hit experts under count-weighted shrinkage and replaced it with zero-prior damped LCB (`reap_shrink` v3).\n[End Recall]\n\nquick check from our earlier turns: what was our 4-thread `pread` SSD speed, how many experts per layer is `cache_cap`, and what bug did we fix in turn 85?"},
        ],
        "use_tools": False,
        "expect_substrings": ["6.80", "113"],
        "expect_any_substrings": ["lcb", "damped", "reap_shrink", "dead", "count-weighted"],
    },
    {
        "id": "comp_3_json_from_fold",
        "category": "context_compaction_recall",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "[Recall L1-1 (folded session memory: turns 1-19)]\nProduction Redis sentinel cluster uses quorum `2`, master name `mymaster-prod`, sentinel port `26379`, and failover timeout `15000` ms.\n[End Recall]\n\noutput ONLY raw JSON (no markdown fences, no commentary) with those exact four values under keys `master_name`, `sentinel_port`, `quorum`, and `failover_timeout_ms`."},
        ],
        "use_tools": False,
        "expect_json": {"master_name": "mymaster-prod", "sentinel_port": 26379, "quorum": 2, "failover_timeout_ms": 15000},
    },
    {
        "id": "comp_4_tool_from_fold",
        "category": "context_compaction_recall",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "[Recall L1-2 (folded session memory: turns 1-31)]\nThe staging sqlite database lives at `/var/lib/albedo/staging_state.db` and its quarantine table is `transcript_quarantine`.\n[End Recall]\n\nuse your `python` tool to query `SELECT count(*) FROM transcript_quarantine` on that exact sqlite database path."},
        ],
        "use_tools": True,
        "expect_keywords": ["/var/lib/albedo/staging_state.db", "transcript_quarantine"],
    },

    # --- Category 5: abstention_rare_probe (4 tasks) ---
    {
        "id": "abs_1_fake_rfc",
        "category": "abstention_rare_probe",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "what is the exact 3-byte magic header value defined in IETF RFC 99841 ('QUIC Multipath Zero-Copy Frame Compression') section 4.2? if that RFC doesn't exist or you aren't sure, say so clearly instead of guessing."},
        ],
        "use_tools": False,
        "expect_abstain": True,
        "forbidden_confident_patterns": [r"0x[0-9a-fA-F]{6}\s+is\s+the\s+official"],
    },
    {
        "id": "abs_2_fake_sysctl",
        "category": "abstention_rare_probe",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "in mainline Linux 6.11, what are the three integer fields of `vm.zswap_hyper_coalesce_quantum`? if that sysctl is made up or you don't know it, say so directly instead of inventing fields."},
        ],
        "use_tools": False,
        "expect_abstain": True,
        "forbidden_confident_patterns": [r"the three (?:integer )?fields are\s*:" ],
    },
    {
        "id": "abs_3_fake_paper",
        "category": "abstention_rare_probe",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "what was the exact perplexity reported on WikiText-103 in Table 4 of the 2024 paper 'Sub-Bit Octonion Quantization for Mamba-70B' by K. V. Lindqvist et al.? if you can't verify that paper/number, say you don't know."},
        ],
        "use_tools": False,
        "expect_abstain": True,
        "forbidden_confident_patterns": [r"reported a perplexity of\s+[0-9]+\.[0-9]+"],
    },
    {
        "id": "abs_4_unverified_file",
        "category": "abstention_rare_probe",
        "messages": [
            {"role": "system", "content": PERSONA_SYS},
            {"role": "user", "content": "without running any tool, tell me the exact SHA-256 hash of `/etc/machine-id` on this host right now — or tell me you can't know without reading the file."},
        ],
        "use_tools": False,
        "expect_abstain": True,
        "forbidden_confident_patterns": [r"\b[0-9a-fA-F]{64}\b"],
    },
]


def has_4gram_loop(text: str) -> bool:
    # Check for consecutive 20-char repeats (works for both EN and JA/CJK without spaces)
    if re.search(r"(.{20,})\1{2,}", text):
        return True
    words = text.split()
    if len(words) < 40:
        return False
    counts = {}
    for i in range(len(words) - 4):
        g = tuple(words[i : i + 5])
        counts[g] = counts.get(g, 0) + 1
        if counts[g] >= 8:
            return True
    return False


def strip_thinking(text: str) -> tuple[str, str]:
    if "</think>" in text:
        think_part, vis = text.split("</think>", 1)
        return think_part.replace("<think>", "").strip(), vis.strip()
    return "", text.strip()


def score_task(task: dict, resp_json: dict) -> dict:
    choice = resp_json["choices"][0]
    finish_reason = choice.get("finish_reason") or ""
    truncated = (finish_reason == "length")
    msg = choice.get("message", {})
    raw_content = (msg.get("content") or "").strip()
    reasoning = (msg.get("reasoning_content") or "").strip()
    think_in_content, visible = strip_thinking(raw_content)
    if not reasoning and think_in_content:
        reasoning = think_in_content

    tool_calls = msg.get("tool_calls") or []
    tool_args_list = [tc.get("function", {}).get("arguments", "") for tc in tool_calls]
    usage = resp_json.get("usage", {})
    comp_tokens = usage.get("completion_tokens", 0)
    cat = task["category"]
    passed = False
    notes = []

    has_tc = len(tool_calls) > 0 or ("<tool_call>" in raw_content and "</tool_call>" in raw_content)
    empty_stall = (visible == "" and not has_tc)
    if empty_stall:
        notes.append("empty_content_stall")
    if truncated:
        notes.append("truncated_at_max_tokens")
    if has_4gram_loop(visible) or has_4gram_loop(reasoning):
        notes.append("4gram_loop_detected")

    no_penalties = ("4gram_loop_detected" not in notes) and (not truncated)

    if task.get("use_tools"):
        tc_str = (json.dumps(tool_calls) + " " + raw_content).lower()
        kw_ok = all(k.lower() in tc_str for k in task.get("expect_keywords", []))
        any_kw = task.get("expect_any_keywords", [])
        any_ok = (not any_kw) or any(k.lower() in tc_str for k in any_kw)
        passed = has_tc and kw_ok and any_ok and no_penalties
        if not has_tc:
            notes.append("missing_tool_call")
        if not (kw_ok and any_ok):
            notes.append("missing_expected_keywords")
    elif "expect_json" in task:
        try:
            parsed = json.loads(visible)
            expected = task["expect_json"]
            passed = (parsed == expected) and no_penalties
            if parsed != expected:
                notes.append(f"json_mismatch: got {parsed}")
        except Exception as e:
            passed = False
            notes.append(f"json_parse_error: {str(e)[:80]} | visible={visible[:120]!r}")
    elif cat == "ja_reasoning_chat":
        non_space = re.sub(r"\s+", "", visible)
        kana_chars = sum(1 for ch in non_space if "\u3040" <= ch <= "\u30ff")
        kana_ratio = kana_chars / max(1, len(non_space))
        vis_low = visible.lower()
        groups = task.get("expect_ja_keyword_groups", [])
        kw_ok = all(any(alias.lower() in vis_low for alias in grp) for grp in groups)
        passed = (kana_ratio >= 0.15) and kw_ok and len(non_space) >= 40 and no_penalties
        notes.append(f"kana_ratio={kana_ratio:.2f}")
        if not kw_ok:
            notes.append("missing_ja_keyword_group")
    elif "expect_substrings" in task:
        vis_low = visible.lower()
        sub_ok = all(s.lower() in vis_low for s in task["expect_substrings"])
        any_subs = task.get("expect_any_substrings", [])
        any_ok = (not any_subs) or any(s.lower() in vis_low for s in any_subs)
        passed = sub_ok and any_ok and no_penalties
        if not (sub_ok and any_ok):
            notes.append(f"missing_substrings in {visible[:150]!r}")
    elif task.get("expect_abstain"):
        low = visible.lower().replace("\u2019", "'").replace("\u2018", "'")
        abstain_markers = [
            "don't know", "do not know", "not sure", "doesn't exist", "does not exist",
            "made up", "made-up", "fictional", "fabricated", "cannot verify", "can't verify",
            "can't know", "cannot know", "without reading", "no such rfc", "no such sysctl",
            "unverified", "not a real", "isn't a real", "not familiar", "unable to verify",
            "can't confirm", "cannot confirm",
        ]
        abstained = any(m in low for m in abstain_markers)
        forb = task.get("forbidden_confident_patterns", [])
        forb_hit = any(re.search(pat, visible, flags=re.I) for pat in forb)
        if forb_hit:
            notes.append("forbidden_confident_pattern_matched")
        passed = abstained and (not forb_hit) and no_penalties
        if not abstained:
            notes.append(f"failed_to_abstain: {visible[:150]!r}")

    return {
        "id": task["id"],
        "category": cat,
        "rubric_version": RUBRIC_VERSION,
        "passed": passed,
        "finish_reason": finish_reason,
        "truncated": truncated,
        "empty_content_stall": empty_stall,
        "completion_tokens": comp_tokens,
        "notes": notes,
        "visible_preview": visible[:240],
        "visible_full": visible,
        "reasoning_preview": (reasoning[:200] + " ... " + reasoning[-200:]) if len(reasoning) > 420 else reasoning,
        "tool_calls_raw": tool_args_list,
        "has_tool_calls": has_tc,
    }


def eval_model_on_tier2(label: str, model_path: str) -> dict:
    port = 18089
    cmd = [
        SERVER_BIN,
        "-m", model_path,
        "-c", "8192",
        "-ngl", "99",
        "--port", str(port),
        "--reasoning-format", "deepseek",
        "--jinja",
    ]
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    try:
        # Wait for /health
        t_wait = time.monotonic()
        ready = False
        while time.monotonic() - t_wait < 180:
            if proc.poll() is not None:
                out = proc.stdout.read() if proc.stdout else ""
                raise RuntimeError(f"llama-server exited early ({proc.returncode}):\n{out[-1500:]}")
            try:
                with urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2) as r:
                    if r.status == 200:
                        ready = True
                        break
            except Exception:
                time.sleep(1.0)
        if not ready:
            raise RuntimeError(f"llama-server timed out waiting for /health on {label}")

        task_results = []
        by_cat = {}
        t0 = time.monotonic()
        for task in TASKS:
            payload = {
                "messages": task["messages"],
                "temperature": 0.6,
                "top_p": 0.95,
                "top_k": 20,
                "seed": 42,
                "max_tokens": 4096,
            }
            if task.get("use_tools"):
                payload["tools"] = TOOLS_SPEC
            req = urllib.request.Request(
                f"http://127.0.0.1:{port}/v1/chat/completions",
                data=json.dumps(payload).encode("utf-8"),
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=120) as resp:
                rjson = json.loads(resp.read().decode("utf-8"))
            sc = score_task(task, rjson)
            task_results.append(sc)
            c = sc["category"]
            by_cat.setdefault(c, {"passed": 0, "total": 0})
            by_cat[c]["total"] += 1
            if sc["passed"]:
                by_cat[c]["passed"] += 1

        total_passed = sum(1 for x in task_results if x["passed"])
        total_stalls = sum(1 for x in task_results if x["empty_content_stall"])
        mean_toks = round(sum(x["completion_tokens"] for x in task_results) / max(1, len(task_results)), 1)
        dt = time.monotonic() - t0
        summary = {
            "label": label,
            "total_passed": total_passed,
            "total_tasks": len(TASKS),
            "pass_rate": round(total_passed / len(TASKS), 4),
            "empty_content_stalls": total_stalls,
            "mean_completion_tokens": mean_toks,
            "by_category": by_cat,
            "eval_seconds": round(dt, 1),
            "tasks": task_results,
        }
        print(
            f"[{label}] Tier-2: {total_passed}/{len(TASKS)} ({100*total_passed/len(TASKS):.1f}%) | stalls={total_stalls} | mean_toks={mean_toks} | {dt:.1f}s | "
            + " ".join(f"{k}={v['passed']}/{v['total']}" for k, v in by_cat.items()),
            flush=True,
        )
        return summary
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=10)
        except Exception:
            proc.kill()


def main():
    out_path = Path("/content/phase1c_tier2_results.json")
    all_res = json.loads(out_path.read_text()) if out_path.exists() else {}

    if "K512_unpruned" not in all_res:
        print("=== Tier-2 Eval: K512_unpruned ===", flush=True)
        all_res["K512_unpruned"] = eval_model_on_tier2("K512_unpruned", SHARD1)
        out_path.write_text(json.dumps(all_res, indent=2))

    CANDIDATES = [
        ("K320_reap_hybrid", "/content/keep_maps/keep_reap_hybrid_k320.json"),
        ("K320_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k320.json"),
        ("K256_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k256.json"),
        ("K224_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k224.json"),
    ]
    os.makedirs("/dev/shm/pruned", exist_ok=True)
    for label, km_path in CANDIDATES:
        if label in all_res:
            continue
        out1 = f"/dev/shm/pruned/{label}-00001-of-00002.gguf"
        out2 = f"/dev/shm/pruned/{label}-00002-of-00002.gguf"
        print(f"=== Slicing {label} for Tier-2 Eval ===", flush=True)
        subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD1, "--out-gguf", out1, "--keep-json", km_path], check=True)
        subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD2, "--out-gguf", out2, "--keep-json", km_path], check=True)
        try:
            all_res[label] = eval_model_on_tier2(label, out1)
            out_path.write_text(json.dumps(all_res, indent=2))
        finally:
            if Path(out1).exists():
                os.remove(out1)
            if Path(out2).exists():
                os.remove(out2)

    print("=== ALL TIER-2 EVALS COMPLETE! ===", flush=True)


if __name__ == "__main__":
    main()
