"""Does "think step by step" help every task, or only logical ones? 3 tasks x 3 models x CoT on/off."""
import certifi, json, os, re, ssl, urllib.request
from concurrent.futures import ThreadPoolExecutor

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
KEY = next(l.split("=", 1)[1].strip() for l in open(os.path.join(ROOT, ".env")) if l.startswith(("OPENROUTER_API_KEY", "OPENROUTER_KEY")))
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cot_results.json")

COT = "Think through this step by step before giving your final answer."

TASKS = {
    "creative": "Write the opening paragraph of a novel about a physicist who discovers gravity is optional.",
    "logical": ("Alice is taller than Bob. Bob is taller than Charlie. Charlie is shorter than Diana. "
                "Diana is taller than Alice. Is this logically possible? Explain."),
    "factual": "Explain the difference between Python's `is` and `==` operators, with an example where they differ.",
}

MODELS = ["openai/gpt-4o-2024-11-20", "anthropic/claude-sonnet-4", "google/gemini-2.5-flash"]


def ask(model, prompt):
    req = urllib.request.Request(
        "https://openrouter.ai/api/v1/chat/completions",
        json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 800, "temperature": 0.7}).encode(),
        {"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"},
    )
    ctx = ssl.create_default_context(cafile=certifi.where())
    r = json.load(urllib.request.urlopen(req, context=ctx, timeout=180))
    c = r["choices"][0]
    return {"text": c["message"]["content"], "finish": c.get("finish_reason"),
            "tokens": r.get("usage", {}).get("completion_tokens")}


def score(task, text):
    """Only claims code can actually check. Creative quality is not one of them."""
    low = text.lower()
    s = {"words": len(text.split()), "has_code_block": "```" in text,
         # did it narrate its reasoning, asked or not?
         "reasons_aloud": bool(re.search(r"step[- ]by[- ]step|let'?s (break|evaluate|work|think)|let me (think|work|trace|analyz)|step \d", low))}
    if task == "logical":
        # Ground truth: Diana > Alice > Bob > Charlie satisfies all four, so the answer is yes.
        # Grade on the stated ordering, not on keywords: "no contradiction" contains "contradiction".
        s["correct"] = bool(re.search(r"d\w*\s*>\s*a\w*\s*>\s*b\w*\s*>\s*c", low)) and not re.search(
            r"\b(not logically possible|is impossible|is not possible|is inconsistent)\b", low)
    if task == "factual":
        s["mentions_identity"] = any(k in low for k in ("identity", "same object", "memory address", "in memory"))
        s["mentions_interning"] = any(k in low for k in ("intern", "cach", "small integer", "-5", "256"))
    if task == "creative":
        # Task said "paragraph", singular. Did the answer arrive clean, or with scaffolding?
        paras = [p for p in text.split("\n\n") if p.strip()]
        s["blocks"] = len(paras)
        s["preamble"] = bool(re.match(r"^(sure|here|okay|let'?s|step \d|\*\*|##|to write|i'?ll)", low.strip()))
        s["meta_words"] = sum(low.count(w) for w in ("step 1", "step 2", "first,", "let's", "i'll", "opening paragraph:"))
    return s


def run(args):
    model, task, cot = args
    prompt = f"{COT}\n\n{TASKS[task]}" if cot else TASKS[task]
    try:
        r = ask(model, prompt)
    except Exception as e:
        return {"model": model, "task": task, "cot": cot, "error": repr(e)}
    return {"model": model, "task": task, "cot": cot, "prompt": prompt, **r, **score(task, r["text"])}


if __name__ == "__main__":
    jobs = [(m, t, c) for m in MODELS for t in TASKS for c in (False, True)]
    with ThreadPoolExecutor(6) as ex:
        results = list(ex.map(run, jobs))
    json.dump(results, open(OUT, "w"), indent=2)

    for r in results:
        if "error" in r:
            print(f"FAIL {r['model']:32} {r['task']:9} cot={r['cot']:<5} {r['error'][:80]}")
            continue
        extra = {k: v for k, v in r.items() if k not in
                 ("model", "task", "cot", "prompt", "text", "finish", "tokens")}
        print(f"{r['model'].split('/')[1]:24} {r['task']:9} cot={str(r['cot']):5} "
              f"{r['tokens']:>4}tok finish={r['finish']:<6} {extra}")
    print(f"\nwrote {OUT}")


def rescore(path=OUT):
    """Re-grade saved responses without spending another 18 calls."""
    rs = json.load(open(path))
    for r in rs:
        r.update(score(r["task"], r["text"]))
    json.dump(rs, open(path, "w"), indent=2)
    for task in TASKS:
        print(f"\n--- {task} ---")
        for r in [x for x in rs if x["task"] == task]:
            keys = ("correct",) if task == "logical" else ("blocks", "preamble") if task == "creative" else ("mentions_interning",)
            print(f"  {r['model'].split('/')[1]:24} cot={str(r['cot']):5} {r['words']:>4}w "
                  f"reasons_aloud={str(r['reasons_aloud']):5} finish={r['finish']:<6} "
                  + " ".join(f"{k}={r.get(k)}" for k in keys))