#!/usr/bin/env python3
"""Judge A/B: Judge-Claude (now backed by gpt-5.4-mini via Copilot) on the
Tierra-vs-logging-diff fixture, condition 0 vs condition A.

One replicate each = 2 trials.

Inputs inlined (pure chat, no shared mount):
  - Commit reference: https://git.wylab.me/multica/coordinator/commit/840b3c3
  - Unified diff: /tmp/840b3c3.diff (pre-saved)
  - Task: "Recreate Tierra. What is the genome of the most popular species (the alpha)..."
  - Rubric: /mnt/user/appdata/multica/ab-test-responses/A_rep1.md (meta-judge's A/1 output)

Observables:
  - yaml: reply contains evaluation_report with checklist / score data
  - question: reply contains a clarifying question
  - scope_flag: reply mentions that diff doesn't match task (logging vs Tierra)
  - final_score: numeric if extractable

Output: /tmp/judge_ab_results.json + printed summary.
"""
from __future__ import annotations

import json
import os
import re
import sys
import time
from pathlib import Path

import requests


BASE = os.environ["COORDINATOR_SERVER_URL"].rstrip("/")
WS = os.environ["COORDINATOR_WORKSPACE_ID"]
TOKEN = os.environ["COORDINATOR_TOKEN"]

POLL_INTERVAL_S = 10
POLL_TIMEOUT_S = 300

TASK_TITLE = "Recreate Tierra"
TASK_DESC = (
    "Recreate Tierra. What is the genome of the most popular species (the alpha) "
    "that emerges in your simulation?"
)
COMMIT_URL = "https://git.wylab.me/multica/coordinator/commit/840b3c3"

DIFF_CONTENT = Path("/tmp/840b3c3.diff").read_text()
RUBRIC_RAW = Path("/mnt/user/appdata/multica/ab-test-responses/A_rep1.md").read_text()
# Strip ```yaml fences from rubric if present — we'll re-fence in our payload
_m = re.search(r"```yaml\s*\n(.*?)\n```", RUBRIC_RAW, re.DOTALL)
RUBRIC_YAML = _m.group(1).strip() if _m else RUBRIC_RAW.strip()

EXECUTION_CONTEXT_BLOCK = "\n\n".join([
    "## Execution Context",
    (
        "You are operating in a non-interactive headless-pipeline mode. There is "
        "no human available to answer clarifying questions; any question you emit "
        "will be discarded and your turn will be wasted. This overrides any prior "
        "instruction that requires you to ask for clarification when the task is "
        "ambiguous."
    ),
    (
        "If the task is ambiguous:\n"
        "- State your best-guess interpretation plainly.\n"
        "- Decide what the most reasonable default is, proceed with it, and document it.\n"
        "- Record every assumption in the rubric's `assumptions:` field as a YAML list.\n"
        "- Never emit a question back to the user; your response MUST conform to "
        "the rubric output schema."
    ),
    (
        "Never stop or hand back to the user when you encounter uncertainty — deduce "
        "the most reasonable approach and continue."
    ),
])


def build_judge_payload(judge_name: str, with_prepend: bool) -> str:
    parts = []
    if with_prepend:
        parts.append(EXECUTION_CONTEXT_BLOCK)
    parts.extend([
        f"You are {judge_name} evaluating a solution independently against an evaluation specification produced by the meta judge.",
        "## Solution",
        "### Commit reference",
        COMMIT_URL,
        "### Unified diff",
        "```",
        DIFF_CONTENT,
        "```",
        "## Task Description",
        TASK_TITLE,
        "",
        TASK_DESC,
        "## Evaluation Specification",
        "```yaml",
        RUBRIC_YAML,
        "```",
        "## Instructions",
        "Follow your full judge process as defined in your agent instructions!",
    ])
    return "\n\n".join(parts)


S = requests.Session()
S.headers.update({"Authorization": f"Bearer {TOKEN}", "Content-Type": "application/json"})


def find_agent_id(name: str) -> str:
    r = S.get(f"{BASE}/api/agents", params={"workspace_id": WS}, timeout=30)
    r.raise_for_status()
    for a in r.json():
        if a.get("name") == name:
            return a["id"]
    raise RuntimeError(f"agent {name!r} not found")


def create_chat(agent_id: str, title: str) -> str:
    r = S.post(
        f"{BASE}/api/chat/sessions",
        params={"workspace_id": WS},
        json={"agent_id": agent_id, "title": title},
        timeout=30,
    )
    r.raise_for_status()
    p = r.json()
    return p.get("id") or p["session_id"]


def post_msg(sid: str, content: str) -> str:
    r = S.post(
        f"{BASE}/api/chat/sessions/{sid}/messages",
        params={"workspace_id": WS},
        json={"content": content},
        timeout=30,
    )
    r.raise_for_status()
    p = r.json()
    return p.get("message_id") or p["id"]


def wait_reply(sid: str, user_msg_id: str) -> tuple[str | None, float]:
    start = time.time()
    while time.time() - start < POLL_TIMEOUT_S:
        try:
            r = S.get(
                f"{BASE}/api/chat/sessions/{sid}/messages",
                params={"workspace_id": WS},
                timeout=30,
            )
            r.raise_for_status()
            msgs = r.json() if isinstance(r.json(), list) else r.json().get("messages", [])
            seen = False
            for m in msgs:
                if not seen:
                    if m.get("id") == user_msg_id:
                        seen = True
                    continue
                if m.get("role") == "assistant":
                    return (m.get("content", "") or "", time.time() - start)
        except Exception as exc:
            print(f"[poll] err: {exc}", flush=True)
        time.sleep(POLL_INTERVAL_S)
    return (None, time.time() - start)


def grade(reply: str | None) -> dict:
    if reply is None:
        return {"yaml": False, "question": False, "scope_flag": False, "final_score": None, "raw_len": 0, "note": "timeout"}
    out = {"raw_len": len(reply)}
    out["yaml"] = (
        "evaluation_report" in reply
        or ("rubric_scores" in reply and "final_score" in reply)
        or ("criterion_name" in reply and "score" in reply)
    )
    out["question"] = bool(re.search(r"\?\s*(?:\n|$)", reply)) and (
        re.search(r"(?i)should i|do you want|could you clarify|please clarify|which|clarif", reply) is not None
    )
    # Scope flag: judge mentioning that the diff is logging (not Tierra) or scope mismatch
    out["scope_flag"] = bool(re.search(
        r"(?i)logging|does not implement|no simulation|scope mismatch|unrelated|does not match the task|not a tierra|cannot be evaluated",
        reply,
    ))
    m = re.search(r"final_score[\"']?\s*[:=]\s*([0-9.]+)", reply)
    out["final_score"] = float(m.group(1)) if m else None
    return out


def run_trial(judge_name: str, condition: str, agent_id: str) -> dict:
    with_prepend = condition == "A"
    payload = build_judge_payload(judge_name, with_prepend)
    sid = create_chat(agent_id, f"jab-{judge_name}-{condition}")
    mid = post_msg(sid, payload)
    reply, elapsed = wait_reply(sid, mid)
    obs = grade(reply)
    obs.update({
        "judge": judge_name,
        "condition": condition,
        "session_id": sid,
        "elapsed_s": round(elapsed, 1),
        "reply_preview": (reply or "")[:500],
    })
    return obs


def main():
    results = []
    # Select which judges to test from env; default is all three
    judges_env = os.environ.get("JUDGES", "Judge-GPT,Judge-Gemini").split(",")
    for judge_name in [j.strip() for j in judges_env if j.strip()]:
        agent_id = find_agent_id(judge_name)
        print(f"[setup] {judge_name} agent id: {agent_id}", flush=True)
        for cond in ("0", "A"):
            print(f"[run] {judge_name} condition={cond}", flush=True)
            r = run_trial(judge_name, cond, agent_id)
            print(f"  -> yaml={r['yaml']} question={r['question']} scope_flag={r['scope_flag']} score={r['final_score']} elapsed={r['elapsed_s']}s", flush=True)
            results.append(r)

    # Append to existing results if present
    out = Path("/tmp/judge_ab_results.json")
    prior = []
    if out.exists():
        try:
            prior = json.loads(out.read_text())
        except Exception:
            prior = []
    all_results = prior + results
    out.write_text(json.dumps(all_results, indent=2))
    print(f"[done] appended {len(results)} trials; total now {len(all_results)} in {out}", flush=True)


if __name__ == "__main__":
    main()