#!/usr/bin/env python3
"""
Genshin Impact Character Knowledge Benchmark
Evaluates a local OpenAI-compatible LLM on benchmark data.

- Options are shuffled deterministically, with a DIFFERENT order for each
  run, to eliminate answer-position bias (source file has 99.9% of correct
  answers in position 1!).
- Every question is asked 3 times. The question counts as CORRECT only
  if the model answers correctly in ALL 3 runs (compared by answer CONTENT,
  not position).
- Accepts many answer formats:
    full option text ("Sumeru")
    bare letters/numbers ("C", "3", "B.", "D)", "B. Blacksmith")
    short phrases ("The answer is C", "ONLY A.", "ONLY B. Grand Sage")

Usage:
    python Benchmark.py [--file Benchmark-Data.json] [--host http://localhost:5001] [--model your-model-name]
"""

import argparse
import json
import random
import re
import sys
import time
from collections import Counter
from pathlib import Path

import requests


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
DEFAULT_FILE = "Benchmark-Data.json"
DEFAULT_HOST = "http://localhost:5001"
DEFAULT_MODEL = "local-model"          # many local servers accept any name
MAX_RETRIES = 3                        # retries per run for invalid output format
RUNS_PER_QUESTION = 3                  # must be correct in ALL runs
TEMPERATURE = 0.7
REQUEST_TIMEOUT = 300                  # seconds per request
MAX_TOKENS = 20

SYSTEM_PROMPT = (
    "You are being tested on your knowledge of Genshin Impact. "
    "For each multiple-choice question, reply with the text of the "
    "correct option and nothing else."
)

QUESTION_TEMPLATE = (
    "In Genshin Impact, answer the following question.\n\n"
    "{question}\n\n"
    "Choose exactly one of the following options:\n"
    "{options}\n\n"
    "Reply with the letter or the full text of the chosen option, "
    "and nothing else."
)


# ---------------------------------------------------------------------------
# Shuffling
# ---------------------------------------------------------------------------
def make_shuffled_options(question: dict, run: int = 0) -> tuple[list[str], int]:
    """
    Shuffle options deterministically. Uses a different order per run
    (seeded by question id + run number) so position bias averages out.
    Returns (shuffled_options, index_of_correct_answer).
    """
    options = list(question["options"])
    correct = question["answer"]
    rng = random.Random(f"{question['id']}_run{run}")
    rng.shuffle(options)
    return options, options.index(correct)


# ---------------------------------------------------------------------------
# API helpers
# ---------------------------------------------------------------------------
def chat(host: str, model: str, messages: list[dict]) -> str:
    """Send a chat completion request to an OpenAI-compatible endpoint."""
    url = f"{host.rstrip('/')}/v1/chat/completions"
    payload = {
        "model": model,
        "messages": messages,
        "temperature": TEMPERATURE,
        "max_tokens": MAX_TOKENS,
    }
    resp = requests.post(url, json=payload, timeout=REQUEST_TIMEOUT)
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"].strip()


# ---------------------------------------------------------------------------
# Answer parsing / validation
# ---------------------------------------------------------------------------
def normalize(text: str) -> str:
    """Normalize a string for comparison."""
    return re.sub(r"\s+", " ", text).strip().lower().strip('".\'')


def letter_to_option(token: str, options: list[str]) -> str | None:
    """Map a letter (a-d) or number (1-4) token to the corresponding option."""
    token = token.lower()
    if token.isdigit():
        idx = int(token) - 1
        return options[idx] if 0 <= idx < len(options) else None
    idx = ord(token) - ord("a")
    return options[idx] if 0 <= idx < len(options) else None


def extract_answer(response: str, options: list[str]) -> str | None:
    """
    Try to map the model's raw response to one of the valid options.
    Returns the matched option text, or None if no valid match is found.

    Accepted formats:
      1) the full option text            ("Sumeru")
      2) option text inside a phrase     ("I think it's Sumeru", if unambiguous,
                                          whole words only)
      3) leading letter/number           ("C", "3", "B. Blacksmith")
      4) "the answer is <letter>"        ("The answer is C")
      5) short phrase w/ UPPERCASE letter ("ONLY A.", "ONLY B. Grand Sage")
    """
    if not response:
        return None
    resp_norm = normalize(response)

    # 1) Exact normalized match on the option text
    for opt in options:
        if normalize(opt) == resp_norm:
            return opt

    # 2) Response contains exactly one of the options (whole words only,
    #    so "Blacksmithing" does NOT match the option "Blacksmith")
    matches = [
        opt for opt in options
        if re.search(rf"(?<!\w){re.escape(normalize(opt))}(?!\w)", resp_norm)
    ]
    if len(matches) == 1:
        return matches[0]

    # 3) Leading option letter (A-D) or number (1-4), with or without
    #    punctuation. (?![\w-]) ensures "Blacksmith" does NOT match as "B"
    #    and "2-Star" does NOT match as option "2".
    m = re.match(
        r"^\s*(?:option\s+)?(?:answer[:\s]*)?([abcd1-4])(?![\w-])",
        response,
        re.IGNORECASE,
    )
    if m:
        opt = letter_to_option(m.group(1), options)
        if opt is not None:
            return opt

    # 4) "The answer is <letter/number>" inside a short response
    if len(resp_norm) < 60:
        m = re.search(
            r"\b(?:answer is|answer:)\s*(?:option\s+)?([abcd1-4])\b",
            response,
            re.IGNORECASE,
        )
        if m:
            opt = letter_to_option(m.group(1), options)
            if opt is not None:
                return opt

    # 5) Short response containing a standalone UPPERCASE letter anywhere,
    #    e.g. "ONLY A.", "A.", "Answer: B", "ONLY B. Grand Sage".
    #    Uppercase-only avoids false matches on the lowercase article "a"
    #    in sentence-like responses. Digits are intentionally NOT searched
    #    here (caseless, too risky in prose); leading digits are handled
    #    by step 3 and "answer: 3" by step 4.
    if len(resp_norm) <= 40:
        m = re.search(r"(?<![\w-])[ABCD](?![\w-])", response)
        if m:
            opt = letter_to_option(m.group(0).lower(), options)
            if opt is not None:
                return opt

    return None


# ---------------------------------------------------------------------------
# Benchmark logic
# ---------------------------------------------------------------------------
def run_once(host: str, model: str, question: dict, options: list[str]) -> tuple[str | None, str, int]:
    """
    Ask the question once. Retries only on invalid output format.

    Returns (matched_option_or_None, final_text_or_raw, attempts).
    """
    user_prompt = QUESTION_TEMPLATE.format(
        question=question["question"],
        options="\n".join(f"  {chr(65 + i)}. {opt}" for i, opt in enumerate(options)),
    )
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt},
    ]

    raw = ""
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            raw = chat(host, model, messages)
        except (requests.RequestException, KeyError, IndexError) as e:
            print(f"    [run attempt {attempt}] API error: {e}")
            time.sleep(2)
            continue

        raw = raw.strip()
        if not raw:
            print(f"    [run attempt {attempt}] Empty response, retrying...")
            continue

        answer = extract_answer(raw, options)
        if answer is not None:
            return answer, raw, attempt
        print(f"    [run attempt {attempt}] Invalid output format: {raw!r}")

    return None, raw, MAX_RETRIES


def ask_question(host: str, model: str, question: dict) -> dict:
    """
    Ask a single question RUNS_PER_QUESTION times, each with a DIFFERENT
    shuffled option order. Correct only if EVERY run identifies the right
    answer (compared by content, not position).
    """
    runs = []
    for run in range(1, RUNS_PER_QUESTION + 1):
        options, correct_idx = make_shuffled_options(question, run=run)
        answer, final_text, attempts = run_once(host, model, question, options)
        runs.append({
            "run": run,
            "options_shuffled": options,
            "correct_position": correct_idx,
            "model_answer": answer,          # option text
            "model_raw_output": final_text,
            "attempts": attempts,
        })

    # Compare by content: the answer text is the same regardless of shuffle
    all_correct = all(r["model_answer"] == question["answer"] for r in runs)

    return {"runs": runs, "correct": all_correct}


def main():
    parser = argparse.ArgumentParser(description="Genshin Impact LLM benchmark")
    parser.add_argument("--file", default=DEFAULT_FILE)
    parser.add_argument("--host", default=DEFAULT_HOST)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    args = parser.parse_args()

    # Load benchmark data
    path = Path(args.file)
    if not path.exists():
        sys.exit(f"Error: benchmark file not found: {path}")
    with path.open("r", encoding="utf-8") as f:
        data = json.load(f)

    questions = data["questions"]
    total = len(questions)
    total_requests = total * RUNS_PER_QUESTION

    print(f"Benchmark    : {data.get('benchmark', 'Unknown')}")
    print(f"Endpoint     : {args.host}  (model: {args.model})")
    print(f"Questions    : {total}")
    print(f"Runs/question: {RUNS_PER_QUESTION} (must be correct in ALL runs)")
    print(f"Total API req: ~{total_requests}")
    print("=" * 70)

    # Server connectivity check
    try:
        requests.get(args.host, timeout=5)
    except requests.RequestException as e:
        sys.exit(f"Error: cannot reach endpoint {args.host}: {e}")

    results = []
    correct_count = 0
    start_time = time.time()

    for idx, q in enumerate(questions, 1):
        print(f"[{idx}/{total}] {q['id']} ({q['character']}): {q['question']}")

        outcome = ask_question(args.host, args.model, q)

        if outcome["correct"]:
            correct_count += 1

        picked_positions = [
            r["options_shuffled"].index(r["model_answer"]) if r["model_answer"] is not None else None
            for r in outcome["runs"]
        ]
        run_status = ", ".join(
            f"run{r['run']}: {'OK' if r['model_answer'] == q['answer'] else 'WRONG'}"
            for r in outcome["runs"]
        )
        print(f"         -> {'CORRECT (all runs)' if outcome['correct'] else 'WRONG'} "
              f"| {run_status} | expected: {q['answer']!r}")

        results.append({
            "id": q["id"],
            "character": q["character"],
            "question": q["question"],
            "expected_answer": q["answer"],
            "picked_positions": picked_positions,
            "runs": outcome["runs"],
            "correct": outcome["correct"],
        })

        # Progress ETA
        elapsed = time.time() - start_time
        done_frac = idx / total
        if done_frac > 0:
            eta_min = (elapsed / done_frac - elapsed) / 60
            print(f"         -> progress: {correct_count} correct | ETA: ~{eta_min:.0f} min")

    elapsed = time.time() - start_time

    # Diagnostics: which positions did the model pick?
    pick_counter = Counter(
        pos for r in results for pos in r["picked_positions"] if pos is not None
    )
    pick_total = sum(pick_counter.values())

    # Save results
    output = {
        "benchmark": data.get("benchmark", "Unknown"),
        "model": args.model,
        "endpoint": args.host,
        "runs_per_question": RUNS_PER_QUESTION,
        "options_shuffled": True,
        "total_questions": total,
        "correct": correct_count,
        "wrong": total - correct_count,
        "score_percent": round(correct_count / total * 100, 2) if total else 0.0,
        "elapsed_seconds": round(elapsed, 1),
        "model_pick_distribution": {
            f"position_{chr(65 + i)}": pick_counter.get(i, 0) for i in range(4)
        },
        "results": results,
    }
    out_path = Path("results.json")
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(output, f, indent=2, ensure_ascii=False)

    # Final summary
    print("=" * 70)
    print(f"Finished in {elapsed / 60:.1f} min. Results written to {out_path.resolve()}")
    print("\nModel's picked-option distribution (across all runs):")
    for i in range(4):
        n = pick_counter.get(i, 0)
        pct = n / pick_total * 100 if pick_total else 0
        print(f"  Position {chr(65 + i)}: {n:5d} ({pct:.1f}%)")
    print()
    print(f"Score (correct in ALL {RUNS_PER_QUESTION} runs): {correct_count}/{total}")
    print(f"FINAL SCORE: {output['score_percent']}%")


if __name__ == "__main__":
    main()