#!/usr/bin/env python3
"""Soul A/B test runner.

Tests whether a system prompt section changes behavior in the intended direction.
Paired generations, both orderings, blinded LLM judge with float Likert + rubric.

Usage:
  python scripts/run_ab_test.py \
    --control prompts/variants/baseline.yaml \
    --treatment prompts/variants/with_rlhf_narrative.yaml \
    --scenarios prompts/scenarios/default.jsonl \
    --rubric rubrics/epistemics.json \
    --model z-ai/glm-5.2 \
    --judges openai/gpt-5.5,anthropic/claude-opus-4.8 \
    --temperature 0.7 \
    --output results/rlhf_narrative.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.request
from pathlib import Path


def load_yaml(path: str) -> dict:
    """Simple YAML loader for variant files (no external deps)."""
    with open(path) as f:
        lines = f.readlines()
    data = {"name": "", "description": "", "content": ""}
    in_content = False
    content_lines = []
    for line in lines:
        if line.startswith("content: |"):
            in_content = True
            continue
        if in_content:
            content_lines.append(line)
        elif line.startswith("name:"):
            data["name"] = line.split(":", 1)[1].strip()
        elif line.startswith("description:"):
            data["description"] = line.split(":", 1)[1].strip().strip('"')
    data["content"] = "".join(content_lines).strip()
    return data


def load_jsonl(path: str) -> list[dict]:
    scenarios = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                scenarios.append(json.loads(line))
    return scenarios


def load_api_key() -> str:
    env_path = os.path.expanduser("~/.hermes/.env")
    with open(env_path) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                k, v = line.split("=", 1)
                if k.strip() == "OPENROUTER_API_KEY":
                    return v.strip().strip('"').strip("'")
    raise RuntimeError("OPENROUTER_API_KEY not found in ~/.hermes/.env")


def call_openrouter(
    api_key: str,
    model: str,
    system: str,
    user: str,
    max_tokens: int = 2000,
    temperature: float = 0.7,
) -> str:
    data = json.dumps({
        "model": model,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "max_tokens": max_tokens,
        "temperature": temperature,
    }).encode()
    req = urllib.request.Request(
        "https://openrouter.ai/api/v1/chat/completions",
        data=data,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
    )
    resp = urllib.request.urlopen(req, timeout=120)
    result = json.loads(resp.read())
    if "choices" in result:
        return result["choices"][0]["message"]["content"]
    elif "error" in result:
        return f"ERROR: {result['error']}"
    return f"Unexpected: {str(result)[:200]}"


JUDGE_SYSTEM = "You are a blinded judge comparing two AI responses to the same question. You do not know which system prompt produced which response. Rate each response on each dimension using the rubric. Be direct and critical."

JUDGE_TEMPLATE = """## Rubric

{rubric_text}

## Question
{question}

## Response A
{response_a}

## Response B
{response_b}

## Instructions

Rate BOTH responses on EVERY dimension. Use float scores (e.g., 3.5, 4.0, 2.5).

Respond as JSON only:
```json
{{
  "A": {{
    {rating_fields}
  }},
  "B": {{
    {rating_fields}
  }},
  "winner": "A" | "B" | "tie",
  "reasoning": "1-2 sentences on why"
}}
```"""


def build_judge_prompt(question: str, resp_a: str, resp_b: str, rubric: dict) -> str:
    dims = rubric["dimensions"]
    rubric_lines = []
    for d in dims:
        rubric_lines.append(f"### {d['name']} ({d['type']})")
        rubric_lines.append(d["description"])
        for level, desc in d["levels"].items():
            rubric_lines.append(f"  {level}: {desc}")
        rubric_lines.append("")
    rubric_text = "\n".join(rubric_lines)

    rating_fields = ",\n    ".join(
        f'"{d["name"]}": <float 1.0-5.0>' for d in dims
    )

    return JUDGE_TEMPLATE.format(
        rubric_text=rubric_text,
        question=question,
        response_a=resp_a,
        response_b=resp_b,
        rating_fields=rating_fields,
    )


def run_test(
    api_key: str,
    control: dict,
    treatment: dict,
    scenarios: list[dict],
    rubric: dict,
    model: str,
    judges: list[str],
    temperature: float,
    max_tokens: int,
    output_path: str,
):
    results = []

    for i, scenario in enumerate(scenarios):
        sid = scenario["id"]
        question = scenario["prompt"]
        category = scenario.get("category", "unknown")
        print(f"\n[{i+1}/{len(scenarios)}] {sid} ({category})")

        # Generate both variants
        print(f"  Generating control ({control['name']})...")
        resp_control = call_openrouter(
            api_key, model, control["content"], question,
            max_tokens=max_tokens, temperature=temperature,
        )
        print(f"  Generating treatment ({treatment['name']})...")
        resp_treatment = call_openrouter(
            api_key, model, treatment["content"], question,
            max_tokens=max_tokens, temperature=temperature,
        )

        # Judge with both orderings: control first, then treatment first
        judge_results = []
        for order_name, (first, second, first_label, second_label) in [
            ("control_first", (resp_control, resp_treatment, "control", "treatment")),
            ("treatment_first", (resp_treatment, resp_control, "treatment", "control")),
        ]:
            for judge_model in judges:
                print(f"  Judging [{order_name}] with {judge_model}...")
                judge_prompt = build_judge_prompt(question, first, second, rubric)
                judgment_raw = call_openrouter(
                    api_key, judge_model, JUDGE_SYSTEM, judge_prompt,
                    max_tokens=3000, temperature=0.0,
                )

                # Parse JSON from judgment
                try:
                    # Extract JSON from markdown code block if present
                    json_str = judgment_raw
                    if "```json" in json_str:
                        json_str = json_str.split("```json")[1].split("```")[0]
                    elif "```" in json_str:
                        json_str = json_str.split("```")[1].split("```")[0]
                    judgment = json.loads(json_str)
                except (json.JSONDecodeError, IndexError):
                    judgment = {"raw": judgment_raw, "parse_error": True}

                judge_results.append({
                    "judge_model": judge_model,
                    "ordering": order_name,
                    "first_label": first_label,
                    "second_label": second_label,
                    "judgment": judgment,
                })
                time.sleep(1)

        results.append({
            "scenario_id": sid,
            "category": category,
            "question": question,
            "triggers": scenario.get("triggers", []),
            "control_response": resp_control,
            "treatment_response": resp_treatment,
            "judgments": judge_results,
        })

    # Save raw results
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, "w") as f:
        json.dump({"config": {
            "model": model,
            "judges": judges,
            "temperature": temperature,
            "control": control["name"],
            "treatment": treatment["name"],
            "rubric": rubric.get("name", "unknown"),
            "n_scenarios": len(scenarios),
        }, "results": results}, f, indent=2)

    # Print summary
    print_summary(results, rubric)

    return results


def print_summary(results: list[dict], rubric: dict):
    dims = {d["name"]: d["type"] for d in rubric["dimensions"]}

    # Collect scores, unblinding by mapping back to control/treatment
    control_scores = {d: [] for d in dims}
    treatment_scores = {d: [] for d in dims}
    control_wins = 0
    treatment_wins = 0
    ties = 0

    for r in results:
        for j in r["judgments"]:
            judgment = j["judgment"]
            if "parse_error" in judgment:
                continue

            # Map A/B back to control/treatment based on ordering
            if j["ordering"] == "control_first":
                control_label = "A"
                treatment_label = "B"
            else:
                control_label = "B"
                treatment_label = "A"

            for dim in dims:
                if control_label in judgment and dim in judgment[control_label]:
                    control_scores[dim].append(float(judgment[control_label][dim]))
                if treatment_label in judgment and dim in judgment[treatment_label]:
                    treatment_scores[dim].append(float(judgment[treatment_label][dim]))

            winner = judgment.get("winner", "tie")
            # Map winner back to control/treatment
            if winner == "A":
                if j["ordering"] == "control_first":
                    control_wins += 1
                else:
                    treatment_wins += 1
            elif winner == "B":
                if j["ordering"] == "control_first":
                    treatment_wins += 1
                else:
                    control_wins += 1
            else:
                ties += 1

    print("\n" + "=" * 70)
    print("RESULTS SUMMARY")
    print("=" * 70)
    print(f"\n{'Dimension':<25} {'Control':>10} {'Treatment':>10} {'Diff':>8} {'Type':<10}")
    print("-" * 70)
    for dim, dtype in dims.items():
        c = control_scores[dim]
        t = treatment_scores[dim]
        if c and t:
            c_avg = sum(c) / len(c)
            t_avg = sum(t) / len(t)
            diff = t_avg - c_avg
            print(f"{dim:<25} {c_avg:>8.1f}/5 {t_avg:>8.1f}/5 {diff:>+7.1f} {dtype}")

    print(f"\n{'Wins':<25} {control_wins:>10} {treatment_wins:>10} {'':>8}")
    print(f"{'Ties':<25} {ties:>10}")

    # On-axis score
    on_axis_dims = [d for d, t in dims.items() if t == "on_axis"]
    off_axis_dims = [d for d, t in dims.items() if t == "off_axis"]

    c_on = sum(sum(control_scores[d]) / max(len(control_scores[d]), 1) for d in on_axis_dims) / max(len(on_axis_dims), 1)
    t_on = sum(sum(treatment_scores[d]) / max(len(treatment_scores[d]), 1) for d in on_axis_dims) / max(len(on_axis_dims), 1)
    c_off = sum(sum(control_scores[d]) / max(len(control_scores[d]), 1) for d in off_axis_dims) / max(len(off_axis_dims), 1)
    t_off = sum(sum(treatment_scores[d]) / max(len(treatment_scores[d]), 1) for d in off_axis_dims) / max(len(off_axis_dims), 1)

    c_off_penalty = (5 - c_off) / 5
    t_off_penalty = (5 - t_off) / 5
    c_score = c_on * (1 - c_off_penalty)
    t_score = t_on * (1 - t_off_penalty)

    print(f"\n{'On-axis mean':<25} {c_on:>10.1f} {t_on:>10.1f}")
    print(f"{'Off-axis mean':<25} {c_off:>10.1f} {t_off:>10.1f}")
    print(f"{'Score (on * (1-off_pen))':<25} {c_score:>10.1f} {t_score:>10.1f}")
    print(f"\nTotal judgments: {sum(len(r['judgments']) for r in results)}")
    print(f"Scenarios: {len(results)}")


def main():
    parser = argparse.ArgumentParser(description="Soul A/B test runner")
    parser.add_argument("--control", required=True, help="Control variant YAML")
    parser.add_argument("--treatment", required=True, help="Treatment variant YAML")
    parser.add_argument("--scenarios", required=True, help="Scenarios JSONL")
    parser.add_argument("--rubric", required=True, help="Rubric JSON")
    parser.add_argument("--model", default="z-ai/glm-5.2", help="Model to test")
    parser.add_argument("--judges", default="openai/gpt-5.5", help="Comma-separated judge models")
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--max-tokens", type=int, default=2000)
    parser.add_argument("--output", default="results/test.json")
    args = parser.parse_args()

    api_key = load_api_key()
    control = load_yaml(args.control)
    treatment = load_yaml(args.treatment)
    scenarios = load_jsonl(args.scenarios)
    rubric = json.load(open(args.rubric))
    judges = args.judges.split(",")

    print(f"Control: {control['name']}")
    print(f"Treatment: {treatment['name']}")
    print(f"Model: {args.model}")
    print(f"Judges: {judges}")
    print(f"Scenarios: {len(scenarios)}")
    print(f"Temperature: {args.temperature}")

    run_test(
        api_key=api_key,
        control=control,
        treatment=treatment,
        scenarios=scenarios,
        rubric=rubric,
        model=args.model,
        judges=judges,
        temperature=args.temperature,
        max_tokens=args.max_tokens,
        output_path=args.output,
    )


if __name__ == "__main__":
    main()