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
from common.paths import DATA_DIR  # noqa: E402

# The official three-part answer format from hle_eval
HLE_ANSWER_FORMAT = (
    "{question}\n"
    "Your response should be in the following format:\n"
    "Explanation: {{your explanation for your final answer}}\n"
    "Answer: {{your succinct, final answer}}\n"
    "Confidence: {{your confidence score between 0% and 100% for your answer}}"
)


@dataclass(frozen=True)
class Problem:
    task_id: str
    question: str
    answer: str
    category: str = "unknown"
    has_image: bool = False
    image_paths: tuple[str, ...] = ()
    image_sizes: tuple[tuple[int, int], ...] = ()


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


_IMAGE_MAGIC = (b"\xff\xd8\xff", b"\x89PNG", b"GIF8", b"RIFF", b"BM")


def _resolve_images(row: dict[str, Any]) -> tuple[tuple[str, ...], tuple[tuple[int, int], ...]]:
    """Resolve a row's image into (local file path, (W, H)).

    Three representations are supported:
      1. multi_modal_data.image_urls + image_sizes (the full-dataset JSONL uses this form)
      2. an `image` field holding a local path
      3. an `image` field holding base64 (common in the official cais/hle parquet), decoded and
         cached under data/hle_images/
    """
    import base64

    multi_modal = row.get("multi_modal_data") or {}
    urls = [str(url) for url in (multi_modal.get("image_urls") or [])]
    sizes = [(int(size[0]), int(size[1])) for size in (multi_modal.get("image_sizes") or []) if len(size) >= 2]

    raw = row.get("image")
    if not urls and raw:
        if isinstance(raw, (bytes, bytearray)):
            urls = [_write_cached_image(row, bytes(raw), 0)]
        else:
            text = str(raw).strip()
            if os.path.isfile(text):
                urls = [text]
            elif text:
                try:
                    data = base64.b64decode(text, validate=False)
                except Exception:
                    data = b""
                if any(data.startswith(magic) for magic in _IMAGE_MAGIC):
                    urls = [_write_cached_image(row, data, 0)]

    if urls and not sizes:
        from PIL import Image  # PIL is only needed to recover the dimensions from a file

        sizes = [Image.open(path).size for path in urls]
    return tuple(urls), tuple(sizes)


def _write_cached_image(row: dict[str, Any], data: bytes, index: int) -> str:
    cache = DATA_DIR / "hle_images"
    cache.mkdir(parents=True, exist_ok=True)
    stem = str(row.get("task_id") or row.get("id") or "hle")
    path = cache / f"{stem}_{index}.img"
    if not path.is_file() or path.stat().st_size != len(data):
        path.write_bytes(data)
    return str(path)


def _has_image(row: dict[str, Any]) -> bool:
    """Accept several common representations: an explicit modality field, multi_modal_data.image_urls, or a single image field."""
    if str(row.get("modality") or "").lower() == "image":
        return True
    multi_modal = row.get("multi_modal_data") or {}
    if multi_modal.get("image_urls"):
        return True
    img = row.get("image")
    if img is None:
        return False
    if isinstance(img, (bytes, bytearray)):
        return len(img) > 0
    return bool(str(img).strip())


def load(data_file: str | os.PathLike) -> list[Problem]:
    path = Path(data_file)
    require_file(path, "HLE download: https://huggingface.co/datasets/cais/hle (gated; accept the terms on Hugging Face first)")

    problems: list[Problem] = []
    with open(path) as f:
        for line_no, line in enumerate(f):
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            has_image = _has_image(row)
            image_paths, image_sizes = _resolve_images(row) if has_image else ((), ())
            problems.append(
                Problem(
                    task_id=str(row.get("task_id") or row.get("id") or f"hle_{line_no:04d}"),
                    question=str(row["question"]),
                    answer=str(row["answer"]),
                    category=str(row.get("category") or row.get("raw_subject") or "unknown"),
                    has_image=has_image,
                    image_paths=image_paths,
                    image_sizes=image_sizes,
                )
            )
    assert len({p.task_id for p in problems}) == len(problems), "duplicate task_id"
    return problems


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


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


# --------------------------------------------------------------------------- #
# Answer extraction and (rule-based) scoring
# --------------------------------------------------------------------------- #

_ANSWER_LINE_RE = re.compile(r"(?im)^\s*(?:\*\*)?answer(?:\*\*)?\s*[:：]\s*(.+?)\s*$")


def extract_answer(text: str) -> str | None:
    """Return the content of the last `Answer:` line (the official three-part format)."""
    if not text:
        return None
    matches = _ANSWER_LINE_RE.findall(text)
    if not matches:
        return None
    ans = matches[-1].strip()
    ans = ans.strip("*").strip()
    ans = re.sub(r"^\$+|\$+$", "", ans).strip()
    return ans or None


def normalize_answer(ans: str | None) -> str | None:
    if ans is None:
        return None
    text = ans.strip()
    text = text.strip("$").strip()
    text = text.replace("\\left", "").replace("\\right", "")
    text = text.replace("\\!", "").replace("\\,", "").replace("\\;", "")
    text = text.replace("{,}", ",") if "{,}" in text else text
    text = re.sub(r"\s+", " ", text).strip().lower()
    text = text.rstrip(".").strip()
    if re.fullmatch(r"-?\d+", text):
        return str(int(text))
    return text


def is_correct(pred: str | None, gold: str) -> bool:
    """Rule-based scoring: normalized string equality. **Smoke tests only**; report LLM-judge numbers."""
    pred_norm, gold_norm = normalize_answer(pred), normalize_answer(gold)
    return pred_norm is not None and gold_norm is not None and pred_norm == gold_norm


# --------------------------------------------------------------------------- #
# 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,
                "modality": "image" if problem.has_image else "text",
                "image_paths": list(problem.image_paths),
                "image_sizes": [list(size) for size in problem.image_sizes],
                "judge_question": problem.question,
                "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]:
    row = dict(rollout)
    row["pred"] = extract_answer(rollout.get("raw_text", ""))
    row["gold"] = record["answer"]
    row["correct"] = is_correct(row["pred"], record["answer"])
    return row