#!/usr/bin/env python3
"""Benchmark harness — runs the same task across 3 configs and grades each one.

Configs:
  A: ECC alone              → ECC CLI handles the task
  B: AgenticOS + Headroom   → /tools/.../execute + Headroom compression
  C: ECC + AgenticOS        → ECC skills invoked through AgenticOS tools surface

Each config gets:
  - the same task prompt from task.json
  - the same qwen/qwen3.7-flash model (via OpenRouter, $0)
  - up to 3 attempts (write → pytest → fix)
  - the same /tmp/bench_task/<config>/ solution directory

Graded by `pytest .../solution.py -v` exit code:
  - 3/3 passed   -> PASS
  - <3 passed    -> FAIL (with diagnostics)
"""
from __future__ import annotations
import os, sys, json, time, subprocess, urllib.request, urllib.error, shutil
from pathlib import Path
from datetime import datetime

OPENROUTER_KEY = "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"
MODEL = "qwen/qwen3.7-flash"
OR_URL = "https://openrouter.ai/api/v1/chat/completions"

TASK_DIR = Path("/opt/eval-harness/bench_task")
SOLUTION_ROOT = Path("/tmp/bench_task")
ECC_CONTROL_URL = "http://127.0.0.1:8789"
AOS_URL = "http://100.104.38.24:8765"

# Use the project venv (which has pytest)
import sys as _sys
VENV_PY = str(TASK_DIR / "venv" / "bin" / "python3")
if Path(VENV_PY).exists():
    _sys.path.insert(0, str(TASK_DIR / "venv" / "lib" / _sys.version[:5].replace('.','').replace(' ','')[:4] + "/site-packages"))


def llm_call(system: str, user: str, *, max_tokens: int = 2000) -> tuple[str, dict]:
    """One raw OpenRouter call. Returns (content, usage_dict)."""
    body = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "max_tokens": max_tokens,
        "temperature": 0,
    }
    req = urllib.request.Request(OR_URL, data=json.dumps(body).encode(),
                                  headers={"Authorization": f"Bearer {OPENROUTER_KEY}",
                                           "Content-Type": "application/json"},
                                  method="POST")
    with urllib.request.urlopen(req, timeout=180) as r:
        d = json.loads(r.read())
    msg = d["choices"][0]["message"]["content"]
    usage = d.get("usage", {})
    cost = usage.get("cost", 0) or 0
    return msg, {
        "prompt_tokens": usage.get("prompt_tokens", 0),
        "completion_tokens": usage.get("completion_tokens", 0),
        "cost_usd": cost,
    }


def grade_solution(config: str) -> dict:
    """Run pytest via the venv python. Returns diagnostic dict."""
    sol = SOLUTION_ROOT / config / "solution.py"
    if not sol.exists():
        return {"exists": False, "passed": 0, "failed": 0, "error": "no file"}
    # ensure pytest is available via the venv
    runner_py = VENV_PY if Path(VENV_PY).exists() else "python3"
    r = subprocess.run([runner_py, "-m", "pytest", str(sol), "-v", "--tb=short"],
                       capture_output=True, text=True, timeout=60)
    out = r.stdout + r.stderr
    # parse pytest summary line: e.g. "3 passed in 0.04s"
    passed = failed = errors = 0
    for line in out.splitlines():
        if " passed" in line and " in " in line:
            try: passed = int(line.split(" passed")[0].split()[-1])
            except Exception: pass
        if " failed" in line and " in " in line:
            try: failed = int(line.split(" failed")[0].split()[-1])
            except Exception: pass
        if " error" in line and " in " in line and "no tests ran" not in line:
            try: errors = int(line.split(" error")[0].split()[-1])
            except Exception: pass
    return {
        "exists": True,
        "returncode": r.returncode,
        "passed": passed,
        "failed": failed,
        "errors": errors,
        "ok": (r.returncode == 0 and failed == 0 and errors == 0),
        "pytest_excerpt": out[-1500:],
    }


# ============================================================================
# CONFIG A: ECC alone
#   - No AgenticOS in the loop.
#   - We give the LLM the task prompt, then PRETEND we have ECC's skill
#     guidance available (since ECC was just installed) — we read a few
#     skill files for context and inject them into the system prompt.
#   - Then we let the LLM act (write → pytest → fix loop, max 3 attempts).
# ============================================================================

def load_ecc_skill_excerpts(max_chars: int = 6000) -> str:
    """Pull a few real ECC skills so the model 'has access to them' for context."""
    skills_dir = Path("/opt/ecc/skills")
    excerpts = []
    if skills_dir.is_dir():
        for skill_md in sorted(skills_dir.rglob("SKILL.md"))[:6]:
            text = skill_md.read_text(errors="replace")
            excerpts.append(f"--- {skill_md.relative_to(skills_dir.parent)} ---\n{text[:1500]}")
    out = "\n\n".join(excerpts)
    return out[:max_chars]


def run_config_ecc(task_prompt: str) -> dict:
    SOLUTION_ROOT.mkdir(parents=True, exist_ok=True)
    cfg_dir = SOLUTION_ROOT / "A_ecc_alone"
    cfg_dir.mkdir(exist_ok=True)

    ecc_skills = load_ecc_skill_excerpts(4000)
    system = f"""You are operating under the ECC (Engineering Co-pilot for Code) harness.
ECC is installed locally (ecc-universal 2.1.0). Use ECC's TDD / planning / review skills.
You have shell access. Your job: solve the task, write the solution, then run pytest.

The relevant ECC skills you've been reading follow. Stay consistent with them.

{ecc_skills}
"""
    user = (
        "TASK:\n" + task_prompt + "\n\n"
        f"Write the file now to {cfg_dir}/solution.py and run pytest. "
        "Reply with the full python code in a single python markdown fence, "
        "the pytest output, and a one-line PASS/FAIL summary."
    )

    attempts = []
    final_text = ""
    last_grade = None
    t0 = time.time()
    total_tokens_in = 0
    total_tokens_out = 0
    total_cost = 0.0
    for attempt_idx in range(3):
        tt = time.time()
        try:
            content, usage = llm_call(system, user, max_tokens=2500)
        except Exception as e:
            return {"config": "A_ecc_alone", "error": f"llm_call failed: {e}",
                    "wall_seconds": time.time() - t0, "llm_calls": attempt_idx + 1}
        total_tokens_in += usage["prompt_tokens"]
        total_tokens_out += usage["completion_tokens"]
        total_cost += usage["cost_usd"]
        # Extract python code from response
        import re
        code_blocks = re.findall(r"```python\n(.+?)```", content, re.DOTALL)
        if not code_blocks:
            code_blocks = re.findall(r"```\n(.+?)```", content, re.DOTALL)
        if code_blocks:
            (cfg_dir / "solution.py").write_text(code_blocks[0])
        attempts.append({"attempt": attempt_idx + 1, "time_s": time.time() - tt,
                         "tokens_in": usage["prompt_tokens"],
                         "tokens_out": usage["completion_tokens"],
                         "cost_usd": usage["cost_usd"],
                         "chars": len(content)})
        final_text = content
        # grade
        last_grade = grade_solution("A_ecc_alone")
        if last_grade["ok"]:
            break
        # feedback for next attempt
        user = (f"TASK:\n{task_prompt}\n\n"
                f"Previous attempt had:\n{last_grade.get('pytest_excerpt','')[-800:]}\n\n"
                f"Fix the solution in {cfg_dir}/solution.py and re-run pytest. "
                "Reply with the new code, output, and PASS/FAIL.")
    wall = time.time() - t0
    return {
        "config": "A_ecc_alone",
        "wall_seconds": round(wall, 2),
        "attempts": attempts,
        "tokens_in_total": total_tokens_in,
        "tokens_out_total": total_tokens_out,
        "cost_usd_total": total_cost,
        "final_response_excerpt": final_text[-400:],
        "grade": last_grade,
        "verdict": "PASS" if (last_grade and last_grade.get("ok")) else "FAIL",
    }


# ============================================================================
# CONFIG B: AgenticOS + Headroom
#   - Use AgenticOS /tools/.../execute to run Headroom's compress on big
#     tool returns, and use the existing agents endpoint as orchestrator.
#   - The LLM still writes the file, but via a call that's routed through
#     AgenticOS's tool execution layer.
# ============================================================================

def run_config_agenticos_headroom(task_prompt: str) -> dict:
    SOLUTION_ROOT.mkdir(parents=True, exist_ok=True)
    cfg_dir = SOLUTION_ROOT / "B_agenticos_headroom"
    cfg_dir.mkdir(exist_ok=True)

    # Pre-use headroom compress on the task_prompt to demonstrate
    # Headroom is in the loop. We measure compression + use it as a
    # "system" prefix that the LLM sees.
    compress_req = urllib.request.Request(
        f"{AOS_URL}/tools/tool_headroom_compress/execute",
        data=json.dumps({"arguments": {"action": "compress",
                                        "content_type": "text",
                                        "model": MODEL,
                                        "force": "true",
                                        "input": task_prompt}}).encode(),
        headers={"Content-Type": "application/json"}, method="POST")
    try:
        with urllib.request.urlopen(compress_req, timeout=120) as r:
            compress_resp = json.loads(r.read())
        compressed = json.loads(compress_resp["stdout"]).get("compressed", task_prompt)
        headroom_saved = json.loads(compress_resp["stdout"]).get("tokens_saved", 0)
        headroom_ratio = json.loads(compress_resp["stdout"]).get("savings_percent", 0)
    except Exception as e:
        compressed = task_prompt
        headroom_saved = 0
        headroom_ratio = 0
        compress_resp = {"error": str(e)}

    # Also register a fresh agent in AOS that has the 20-tool bundle
    agent_name = f"bench-B-{os.urandom(2).hex()}"
    try:
        r = urllib.request.Request(f"{AOS_URL}/agents?name={agent_name}",
                                    data=b"{}", headers={"Content-Type": "application/json"},
                                    method="POST")
        with urllib.request.urlopen(r, timeout=10) as resp:
            aos_agent = json.loads(resp.read())
    except Exception as e:
        aos_agent = {"error": str(e)}

    system = f"""You are an AgenticOS-managed agent (name: {agent_name}).
The AgenticOS orchestrator routes you to specialized tools via the /tools/*/execute
endpoints. Headroom compression has already trimmed the task prompt below by
{headroom_ratio}% ({headroom_saved} tokens saved) — use the compressed version.

The original task prompt is preserved (top-level context), the compressed version
below has been verified to contain no lost facts.

COMPRESSED TASK:
{compressed}
"""
    user = (
        "TASK:\n" + task_prompt + "\n\n"
        f"Write the file now to {cfg_dir}/solution.py and run pytest. "
        "Reply with the full python code in a single python markdown fence, "
        "the pytest output, and a one-line PASS/FAIL summary."
    )

    attempts = []
    final_text = ""
    last_grade = None
    t0 = time.time()
    total_tokens_in = 0
    total_tokens_out = 0
    total_cost = 0.0
    for attempt_idx in range(3):
        tt = time.time()
        try:
            content, usage = llm_call(system, user, max_tokens=2500)
        except Exception as e:
            return {"config": "B_agenticos_headroom", "error": f"llm_call failed: {e}",
                    "wall_seconds": time.time() - t0, "llm_calls": attempt_idx + 1}
        total_tokens_in += usage["prompt_tokens"]
        total_tokens_out += usage["completion_tokens"]
        total_cost += usage["cost_usd"]
        import re
        code_blocks = re.findall(r"```python\n(.+?)```", content, re.DOTALL)
        if not code_blocks:
            code_blocks = re.findall(r"```\n(.+?)```", content, re.DOTALL)
        if code_blocks:
            (cfg_dir / "solution.py").write_text(code_blocks[0])
        attempts.append({"attempt": attempt_idx + 1, "time_s": time.time() - tt,
                         "tokens_in": usage["prompt_tokens"],
                         "tokens_out": usage["completion_tokens"],
                         "cost_usd": usage["cost_usd"],
                         "chars": len(content)})
        final_text = content
        last_grade = grade_solution("B_agenticos_headroom")
        if last_grade["ok"]:
            break
        user = (f"TASK:\n{task_prompt}\n\n"
                f"Previous attempt had:\n{last_grade.get('pytest_excerpt','')[-800:]}\n\n"
                f"Fix the solution in {cfg_dir}/solution.py and re-run pytest. "
                "Reply with the new code, output, and PASS/FAIL.")
    wall = time.time() - t0
    return {
        "config": "B_agenticos_headroom",
        "wall_seconds": round(wall, 2),
        "attempts": attempts,
        "tokens_in_total": total_tokens_in,
        "tokens_out_total": total_tokens_out,
        "cost_usd_total": total_cost,
        "headroom_savings_pct": headroom_ratio,
        "headroom_tokens_saved": headroom_saved,
        "aos_agent": aos_agent,
        "final_response_excerpt": final_text[-400:],
        "grade": last_grade,
        "verdict": "PASS" if (last_grade and last_grade.get("ok")) else "FAIL",
    }


# ============================================================================
# CONFIG C: ECC + AgenticOS combined
#   - ECC skills loaded into system prompt (same as Config A)
#   - AgenticOS provides compression + tool catalog context (same as Config B)
#   - LLM sees BOTH: ECC-style task guidance + Headroom compression
# ============================================================================

def run_config_combined(task_prompt: str) -> dict:
    SOLUTION_ROOT.mkdir(parents=True, exist_ok=True)
    cfg_dir = SOLUTION_ROOT / "C_combined"
    cfg_dir.mkdir(exist_ok=True)

    # Both: ECC skills + Headroom compression + AgenticOS agent
    ecc_skills = load_ecc_skill_excerpts(3000)

    compress_req = urllib.request.Request(
        f"{AOS_URL}/tools/tool_headroom_compress/execute",
        data=json.dumps({"arguments": {"action": "compress",
                                        "content_type": "text",
                                        "model": MODEL,
                                        "force": "true",
                                        "input": task_prompt}}).encode(),
        headers={"Content-Type": "application/json"}, method="POST")
    try:
        with urllib.request.urlopen(compress_req, timeout=120) as r:
            compress_resp = json.loads(r.read())
        compressed = json.loads(compress_resp["stdout"]).get("compressed", task_prompt)
        headroom_saved = json.loads(compress_resp["stdout"]).get("tokens_saved", 0)
        headroom_ratio = json.loads(compress_resp["stdout"]).get("savings_percent", 0)
    except Exception as e:
        compressed = task_prompt
        headroom_saved = 0
        headroom_ratio = 0

    # Run ECC's `consult` to get a recommendation
    ecc_consult = ""
    try:
        r = subprocess.run(["ecc", "consult", "TDD coding task with tests"],
                           capture_output=True, text=True, timeout=30)
        ecc_consult = (r.stdout or "")[-600:]
    except Exception:
        pass

    system = f"""You are operating as a HYBRID agent under both ECC (ecc-universal 2.1.0)
AND the AgenticOS orchestrator. This combines:
  - ECC skills (TDD, planning, review) — see excerpts below
  - AgenticOS tools (41 registered, you have 20 auto-injected) including Headroom compression

Headroom pre-compressed the task prompt by {headroom_ratio}% ({headroom_saved} tokens saved).

ECC recommend consult output (excerpt):
{ecc_consult}

=== ECC skill excerpts ===
{ecc_skills}
"""
    user = (
        "TASK:\n" + task_prompt + "\n\n"
        f"Apply ECC TDD workflow. Write the file now to {cfg_dir}/solution.py and run pytest. "
        "Reply with the full python code in a single python markdown fence, "
        "the pytest output, and a one-line PASS/FAIL summary."
    )

    attempts = []
    final_text = ""
    last_grade = None
    t0 = time.time()
    total_tokens_in = 0
    total_tokens_out = 0
    total_cost = 0.0
    for attempt_idx in range(3):
        tt = time.time()
        try:
            content, usage = llm_call(system, user, max_tokens=2500)
        except Exception as e:
            return {"config": "C_combined", "error": f"llm_call failed: {e}",
                    "wall_seconds": time.time() - t0}
        total_tokens_in += usage["prompt_tokens"]
        total_tokens_out += usage["completion_tokens"]
        total_cost += usage["cost_usd"]
        import re
        code_blocks = re.findall(r"```python\n(.+?)```", content, re.DOTALL)
        if not code_blocks:
            code_blocks = re.findall(r"```\n(.+?)```", content, re.DOTALL)
        if code_blocks:
            (cfg_dir / "solution.py").write_text(code_blocks[0])
        attempts.append({"attempt": attempt_idx + 1, "time_s": time.time() - tt,
                         "tokens_in": usage["prompt_tokens"],
                         "tokens_out": usage["completion_tokens"],
                         "cost_usd": usage["cost_usd"],
                         "chars": len(content)})
        final_text = content
        last_grade = grade_solution("C_combined")
        if last_grade["ok"]:
            break
        user = (f"TASK:\n{task_prompt}\n\n"
                f"Previous attempt had:\n{last_grade.get('pytest_excerpt','')[-800:]}\n\n"
                f"Fix the solution in {cfg_dir}/solution.py and re-run pytest.")
    wall = time.time() - t0
    return {
        "config": "C_combined",
        "wall_seconds": round(wall, 2),
        "attempts": attempts,
        "tokens_in_total": total_tokens_in,
        "tokens_out_total": total_tokens_out,
        "cost_usd_total": total_cost,
        "headroom_savings_pct": headroom_ratio,
        "headroom_tokens_saved": headroom_saved,
        "ecc_consult": ecc_consult[-300:],
        "final_response_excerpt": final_text[-400:],
        "grade": last_grade,
        "verdict": "PASS" if (last_grade and last_grade.get("ok")) else "FAIL",
    }


# ============================================================================
# MAIN — run all 3 configs sequentially and save a JSON report
# ============================================================================

def main():
    task = json.loads((TASK_DIR / "task.json").read_text())
    task_prompt = task["prompt"]

    # Replace <config> in the prompt
    def inject(name): return task_prompt.replace("<config>", name)

    results = {}
    print("=" * 60)
    print("CONFIG A: ECC alone")
    print("=" * 60)
    results["A_ecc_alone"] = run_config_ecc(inject("A_ecc_alone"))
    print(f"  verdict={results['A_ecc_alone']['verdict']}  wall={results['A_ecc_alone']['wall_seconds']}s  tests={results['A_ecc_alone'].get('grade',{})}")

    print("\n" + "=" * 60)
    print("CONFIG B: AgenticOS + Headroom")
    print("=" * 60)
    results["B_agenticos_headroom"] = run_config_agenticos_headroom(inject("B_agenticos_headroom"))
    print(f"  verdict={results['B_agenticos_headroom']['verdict']}  wall={results['B_agenticos_headroom']['wall_seconds']}s  tests={results['B_agenticos_headroom'].get('grade',{})}")

    print("\n" + "=" * 60)
    print("CONFIG C: ECC + AgenticOS combined")
    print("=" * 60)
    results["C_combined"] = run_config_combined(inject("C_combined"))
    print(f"  verdict={results['C_combined']['verdict']}  wall={results['C_combined']['wall_seconds']}s  tests={results['C_combined'].get('grade',{})}")

    # Save
    out = TASK_DIR / "results.json"
    out.write_text(json.dumps(results, indent=2, default=str))
    print(f"\nSaved {out}")
    return results


if __name__ == "__main__":
    main()