from __future__ import annotations

import json
import os
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

EVAL_DIR = Path(__file__).resolve().parent.parent  # eval/
if str(EVAL_DIR) not in sys.path:
    sys.path.insert(0, str(EVAL_DIR))

from common.io import require_file  # noqa: E402

GPQA_PROMPT = (
    "Answer the following multiple choice question. The last line of your response should be "
    "of the following format: 'ANSWER: $LETTER' (without quotes) where LETTER is one of ABCD. "
    "Think step by step before answering.\n\n{question}"
)


@dataclass(frozen=True)
class Problem:
    task_id: str
    question: str
    answer: str  # a single letter A-D
    category: str = "unknown"


# --------------------------------------------------------------------------- #
# Data
# --------------------------------------------------------------------------- #


# The public mirror appends an answer-format instruction to each question, which conflicts with our
# prompt template, so it is stripped uniformly.
_TRAILING_INSTR_RE = re.compile(r"\n*\s*Please write your final answer in the form of.*$", re.I | re.S)


def _strip_trailing_instruction(question: str) -> str:
    return _TRAILING_INSTR_RE.sub("", question).strip()


def _normalize_row(row: dict[str, Any], index: int) -> Problem | None:
    """Accept both the parquet-derived format and the public-mirror format."""
    question = row.get("question") or row.get("problem")
    raw_answer = row.get("answer") or row.get("solution")
    if question is None or raw_answer is None:
        return None
    # \boxed{D} / ANSWER: D / a bare D are all normalized to a single letter.
    # Do not search for the first A-D: the "B" in "BOXED" inside \boxed{D} would match.
    raw = str(raw_answer).strip()
    match = (
        re.search(r"\\boxed\{\s*\(?([A-D])\)?\s*\}", raw, re.I)
        or re.search(r"(?i)answer\s*[:：]\s*\(?([A-D])\b", raw)
        or re.fullmatch(r"\s*\(?([A-D])\)?\s*", raw)
    )
    if not match:
        return None
    task_id = str(row.get("task_id") or row.get("id") or f"gpqa_diamond_{index}")
    category = str(row.get("category") or row.get("domain") or "unknown")
    return Problem(
        task_id=task_id,
        question=_strip_trailing_instruction(str(question)),
        answer=match.group(1),
        category=category,
    )


def load(data_file: str | os.PathLike) -> list[Problem]:
    path = Path(data_file)
    require_file(
        path,
        "GPQA Diamond download: https://huggingface.co/datasets/Idavidrein/gpqa (gated)"
        " or the public mirror https://huggingface.co/datasets/hendrydong/gpqa_diamond_mc",
    )

    problems: list[Problem] = []
    skipped = 0
    with open(path) as f:
        for line_no, line in enumerate(f):
            line = line.strip()
            if not line:
                continue
            problem = _normalize_row(json.loads(line), line_no)
            if problem is None:
                skipped += 1
                continue
            problems.append(problem)
    assert problems, f"no problem parsed from {path}"
    assert len({p.task_id for p in problems}) == len(problems), "duplicate task_id"
    if skipped:
        print(f"[gpqa] skipped {skipped} lines (missing fields, or an answer that is not A-D)", flush=True)
    return problems


def build_question(problem: Problem) -> str:
    return GPQA_PROMPT.format(question=problem.question.strip())


def build_messages(problem: Problem) -> list[dict[str, Any]]:
    return [{"role": "user", "content": build_question(problem)}]


# --------------------------------------------------------------------------- #
# Answer extraction and scoring
# --------------------------------------------------------------------------- #

_ANSWER_LINE_RE = re.compile(r"(?im)^\s*ANSWER\s*[:：]\s*\(?([A-D])\b")
_BOXED_RE = re.compile(r"\\boxed\{\s*\(?([A-D])\)?\s*\}")
_STANDALONE_RE = re.compile(r"(?<![A-Za-z])([A-D])(?![A-Za-z])")


def extract_answer_with_source(text: str) -> tuple[str | None, str]:
    """Extract the answer letter by priority and report which path matched."""
    if not text:
        return None, "none"
    matches = _ANSWER_LINE_RE.findall(text)
    if matches:
        return matches[-1].upper(), "answer_line"
    matches = _BOXED_RE.findall(text)
    if matches:
        return matches[-1].upper(), "boxed"
    tail = text.rsplit("</think>", 1)[-1] if "</think>" in text else text[-400:]
    matches = _STANDALONE_RE.findall(tail)
    if matches:
        return matches[-1].upper(), "tail_letter"
    return None, "none"


def is_correct(pred: str | None, gold: str) -> bool:
    return pred is not None and pred.upper() == str(gold).strip().upper()


# --------------------------------------------------------------------------- #
# Entry points used by run_eval.py / score.py
# --------------------------------------------------------------------------- #


def build_records(problems: list[Problem], *, limit: int | None = None) -> list[dict[str, Any]]:
    records = []
    for problem in problems:
        records.append(
            {
                "task_id": problem.task_id,
                "question": problem.question,
                "answer": problem.answer,
                "category": problem.category,
                "judge_question": build_question(problem),
                "messages": build_messages(problem),
            }
        )
        if limit is not None and len(records) >= limit:
            break
    return records


def score_rollout(rollout: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
    pred, source = extract_answer_with_source(rollout.get("raw_text", ""))
    row = dict(rollout)
    row["pred"] = pred
    row["extract_source"] = source
    row["gold"] = record["answer"]
    row["correct"] = is_correct(pred, record["answer"])
    return row