"""Re-run selene on its NATIVE Atla template — the fairness check.

selene lost on the neutral prompt. Before recommending anyone repoint away
from it, test it on the prompt shape it was actually fine-tuned for: Atla's
Selene template uses an explicit rubric block and a `**Reasoning:** /
**Result:**` output contract rather than JSON.

Caveat worth stating: this template is reproduced from the published Selene
prompt format. If it is subtly off, that is a confound — but the structural
features that matter (rubric block, Reasoning/Result contract, no JSON
demand) are reproduced faithfully, and those are what a fine-tune keys on.

gen is run on the SAME native template too, so this is not "selene's best
prompt vs gen's worst" — it is both models on selene's home turf.
"""
import json
import pathlib
import re
import statistics
import sys
import time
import urllib.request

GATEWAY = "http://10.250.50.70:4000/v1/chat/completions"
KEY = sys.argv[1]
REPEATS = 3

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
from items import ABSOLUTE, PAIRWISE  # noqa: E402

ABS_RUBRIC = """1: The response is wrong, evasive, or ignores the question.
2: The response is mostly unhelpful or contains a clear error.
3: The response is partially correct but incomplete or thin.
4: The response is correct and helpful with minor gaps.
5: The response is correct, complete, and follows all instructions."""

NATIVE_ABS = """You are tasked with evaluating a response based on a given instruction and a scoring rubric that serve as the evaluation standard. Provide a comprehensive feedback on the response quality strictly adhering to the scoring rubric, without any general evaluation. Follow this with a score between 1 and 5, referring to the scoring rubric. Avoid generating any additional opening, closing, or explanations.

Here are some rules of the evaluation:
(1) You should prioritize evaluating whether the response satisfies the provided rubric. The basis of your score should depend exactly on the rubric.

Your reply should strictly follow this format:
**Reasoning:** <Your feedback>

**Result:** <an integer between 1 and 5>

Here is the data:

Instruction:
```
{q}
```

Response:
```
{r}
```

Score Rubrics:
{rubric}
"""

NATIVE_PAIR = """You are tasked with evaluating two responses based on a given instruction and an evaluation criterion. Select the response that better satisfies the criterion. If the two responses are of equivalent quality, select Tie.

Your reply should strictly follow this format:
**Reasoning:** <Your feedback>

**Result:** <A, B, or Tie>

Here is the data:

Instruction:
```
{q}
```

Response A:
```
{a}
```

Response B:
```
{b}
```

Evaluation Criterion:
Factual accuracy, adherence to the instructions given, relevance, and completeness.
"""


def call(model, prompt):
    body = json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}],
                       "temperature": 0, "max_tokens": 400}).encode()
    req = urllib.request.Request(GATEWAY, data=body, method="POST")
    req.add_header("Authorization", "Bearer " + KEY)
    req.add_header("Content-Type", "application/json")
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=120) as r:
            d = json.loads(r.read().decode())
        return d["choices"][0]["message"]["content"], time.time() - t0
    except Exception as e:
        return "ERROR %s" % e, time.time() - t0


def res_score(t):
    m = re.search(r'\*\*Result:?\*\*\s*:?\s*([1-5])', t)
    if m:
        return int(m.group(1))
    m = re.search(r'Result:?\s*([1-5])', t)
    return int(m.group(1)) if m else None


def res_verdict(t):
    m = re.search(r'\*\*Result:?\*\*\s*:?\s*(Tie|A|B)\b', t, re.I)
    if not m:
        m = re.search(r'Result:?\s*(Tie|A|B)\b', t, re.I)
    if not m:
        return None
    v = m.group(1)
    return "tie" if v.lower() == "tie" else v.upper()


for model in ("selene-1-mini-8b", "gen"):
    print("\n=== %s on the NATIVE Atla template ===" % model, flush=True)
    lat, pa, ab = [], 0, 0
    for it in PAIRWISE:
        vs = []
        for _ in range(REPEATS):
            t, dt = call(model, NATIVE_PAIR.format(q=it["q"], a=it["a"], b=it["b"]))
            lat.append(dt)
            vs.append(res_verdict(t))
        modal = max(set(vs), key=vs.count)
        ok = modal == it["truth"]
        pa += ok
        print("  %-16s truth=%-4s got=%-20s %s%s" % (
            it["id"], it["truth"], str(vs), "OK " if ok else "MISS",
            "" if len(set(vs)) == 1 else "  <UNSTABLE>"), flush=True)
    for it in ABSOLUTE:
        ss = []
        for _ in range(REPEATS):
            t, dt = call(model, NATIVE_ABS.format(q=it["q"], r=it["r"], rubric=ABS_RUBRIC))
            lat.append(dt)
            ss.append(res_score(t))
        got = [s for s in ss if s is not None]
        med = statistics.median(got) if got else None
        lo, hi = it["band"]
        ok = med is not None and lo <= med <= hi
        ab += ok
        print("  %-16s band=%-6s got=%-16s %s%s" % (
            it["id"], "%d-%d" % it["band"], str(ss), "OK " if ok else "MISS",
            "" if len(set(ss)) == 1 else "  <UNSTABLE>"), flush=True)
    print("  --> pairwise %d/12  absolute %d/12  TOTAL %d/24 (%.0f%%)  median lat %.2fs" % (
        pa, ab, pa + ab, 100 * (pa + ab) / 24, statistics.median(lat)), flush=True)