"""Judge a rendered ring against its web reference.

Per RFC D4: the optimization signal is a PER-AXIS RUBRIC (discrete checks average
more stably than the noisy holistic 0-1 and each failed axis maps to a concrete
brief edit). A holistic 0-1 is also returned, but ONLY as a soft dashboard number.

  from harness.judge import judge
  judge(Path("ref.jpg"), Path("build.png"), request)
  # -> {"rubric": {...bool}, "score": 0.67, "holistic": 0.4, "diffs": [...]}
"""

from __future__ import annotations

import base64
import json
import os
import re
import urllib.request
from pathlib import Path

# claude-opus-4.8 perceives fine schematic detail (e.g. small gripping prongs) that
# gemini-3.1-pro/3.5-flash systematically miss — verified A/B on a clear 4-claw render
# (opus prongs_claws True 3/3, gemini False 3/3). Worth the cost for measurement fidelity.
MODEL = os.environ.get("PI_JUDGE_MODEL", "anthropic/claude-opus-4.8")

# Generalistic axes — apply to any ring; "n/a" passes (e.g. accents on a solitaire).
AXES = {
    "band_slim": "Band is slender/petite and does NOT visually compete with the stone.",
    "stone_dominant": "The center stone is the dominant focal mass, sized and proud.",
    "prongs_claws": "Prongs read as fine claws gripping the crown, not thick pillars.",
    "accents_fine": "Any halo/pavé/side accents are fine, dense, and well-set (n/a if none).",
    "profile_low": "Low, sleek profile — not tall, bulky, or heavy.",
    "setting_correct": "The setting style matches the requested '{setting}'.",
    "printable_one_piece": "PRINTABILITY: the metal reads as ONE connected, castable "
    "piece — every prong, accent seat, and the band is joined to the body, with "
    "nothing floating, detached, or hair-thin. Could be 3D-printed and cast as one.",
}

PROMPT = (Path(__file__).resolve().parent.parent / "prompts" / "judge.md").read_text()


MAX_JUDGE_PX = 768   # r69: 3 full-res images/vote at Fable pricing burned the
#                      $50/day OpenRouter cap in one day; 768px is plenty for
#                      rubric-level judgments and ~4x cheaper


def _b64(p: Path) -> str:
    try:
        from io import BytesIO

        from PIL import Image
        img = Image.open(p)
        if max(img.size) > MAX_JUDGE_PX:
            img.thumbnail((MAX_JUDGE_PX, MAX_JUDGE_PX))
            buf = BytesIO()
            img.convert("RGB").save(buf, format="JPEG", quality=88)
            return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
    except Exception:
        pass
    mime = "image/png" if p.suffix.lower() == ".png" else "image/jpeg"
    return f"data:{mime};base64," + base64.b64encode(p.read_bytes()).decode()


def judge(ref: Path, build: Path, request: dict, head_png: Path | None = None) -> dict:
    axes_txt = "\n".join(f"- {k}: {v.format(**request)}" for k, v in AXES.items())
    keys = ", ".join(f'"{k}": <true|false>' for k in AXES)
    prompt = PROMPT.format(axes=axes_txt, keys=keys, **request)
    content = [
        {"type": "text", "text": prompt},
        {"type": "image_url", "image_url": {"url": _b64(ref)}},
        {"type": "image_url", "image_url": {"url": _b64(build)}},
    ]
    if head_png is not None and head_png.exists():
        content.insert(1, {"type": "text", "text":
                           "Image 3 = a CLOSE-UP of image 2's setting/head. Use it "
                           "to verify the small-scale checks: prong contact and "
                           "taper, accent-stone setting, and whether the metal "
                           "connects as one piece."})
        content.append({"type": "image_url", "image_url": {"url": _b64(head_png)}})
    body = json.dumps({
        "model": MODEL, "max_tokens": 1200, "reasoning": {"effort": "low"},
        "messages": [{"role": "user", "content": content}],
    }).encode()
    req = urllib.request.Request(
        "https://openrouter.ai/api/v1/chat/completions", data=body,
        headers={"Authorization": f"Bearer {os.environ['OPENROUTER_API_KEY']}",
                 "Content-Type": "application/json"})
    txt = json.load(urllib.request.urlopen(req, timeout=120))["choices"][0]["message"]["content"]
    m = re.search(r"\{.*\}", txt, re.S)
    j = json.loads(m.group(0)) if m else {"rubric": {}, "holistic": 0.0, "diffs": ["no judge output"]}
    rubric = {k: bool(j.get("rubric", {}).get(k, False)) for k in AXES}
    score = sum(rubric.values()) / len(rubric)        # stable per-axis fraction = the signal
    return {"rubric": rubric, "score": round(score, 3),
            "holistic": float(j.get("holistic", 0.0)), "diffs": j.get("diffs", [])}