#!/usr/bin/env python3
"""IOL-AI 2026 submission script.

Reads /tmp/data/test.csv and writes submission.csv.
Model weights must be shipped in the repo; default local path is ./model.
"""
from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Any


def _ensure_deps() -> None:
    """Fail clearly instead of installing anything during evaluation."""
    try:
        import pandas  # noqa: F401
    except ImportError as exc:
        raise RuntimeError(
            "Missing required dependency: pandas. Install dependencies before evaluation; "
            "script.py will not download packages or call the internet at runtime."
        ) from exc


_ensure_deps()

import pandas as pd  # noqa: E402


INPUT_CSV = Path("/tmp/data/test.csv")
OUTPUT_CSV = Path("submission.csv")
MODEL_DIR = os.environ.get("MODEL_DIR", "./model")
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "768"))
DUMMY_MODE = os.environ.get("IOL_DUMMY", "0") == "1"

# Force local/offline loading for HF libraries.
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

ANSWER_KEYS = ("answers", "answer", "pred", "prediction", "predictions")
FENCED_BLOCK_RE = re.compile(r"```(?:json)?\s*(.*?)```", flags=re.S | re.I)
LIST_ITEM_RE = re.compile(
    r"^\s*(?:(?:[-*+]\s*)?(?:\(?\d{1,3}\)?[\).:]|[A-Za-z][\).:])|[-*+\u2013\u2014\u2022\u2023\u2043\u2219\u25e6])\s+(.*?)\s*$"
)
HEADER_LINE_RE = re.compile(r"^\s*(?:answers?|predictions?|output|final answers?)\s*:?\s*$", flags=re.I)
INTRO_LINE_RE = re.compile(r"^\s*(?:here are|the answers are|my answers are)\b", flags=re.I)
JSON_SCAFFOLD_LINE_RE = re.compile(
    r"^\s*(?:[\{\}\[\],]+|[\"']?(?:answers?|predictions?|pred|prediction)[\"']?\s*:\s*\[?)\s*$", flags=re.I
)


def expected_item_count(query: str) -> int | None:
    """Estimate number of numbered items in the query.

    Returns None if no reliable numbering is visible.
    """
    # Common IOL format: each item starts with `17.`, `17)`, `(17)` etc.
    matches = re.findall(r"(?m)^\s*(?:\(?\d{1,3}\)?[\).:]|[A-Z][\).:])\s+", query)
    if matches:
        return len(matches)
    return None


def clean_answer_text(text: str) -> str:
    """Remove only formatting wrappers around an answer."""
    text = text.strip()
    quote_pairs = {
        '"': '"',
        "'": "'",
        "\u201c": "\u201d",
        "\u2018": "\u2019",
    }
    if len(text) >= 2 and quote_pairs.get(text[0]) == text[-1]:
        text = text[1:-1].strip()
    return text


def fit_answer_count(answers: list[str], n_expected: int | None) -> list[str]:
    if n_expected is None:
        return answers
    if len(answers) > n_expected:
        return answers[:n_expected]
    if len(answers) < n_expected:
        return answers + [""] * (n_expected - len(answers))
    return answers


def answer_value_to_text(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return clean_answer_text(value)
    if isinstance(value, (dict, list)):
        return json.dumps(value, ensure_ascii=False)
    return clean_answer_text(str(value))


def answers_from_json_value(value: Any) -> list[str] | None:
    if isinstance(value, dict):
        for key in ANSWER_KEYS:
            if key in value:
                return answers_from_json_value(value[key])
        return None
    if isinstance(value, list):
        answers: list[str] = []
        for item in value:
            if isinstance(item, dict):
                item_answers = answers_from_json_value(item)
                answers.append(item_answers[0] if item_answers else answer_value_to_text(item))
            else:
                answers.append(answer_value_to_text(item))
        return answers
    if isinstance(value, (str, int, float, bool)) or value is None:
        return [answer_value_to_text(value)]
    return None


def iter_json_values(source: str, scan_embedded: bool) -> list[Any]:
    decoder = json.JSONDecoder()
    source = source.strip()
    if not source:
        return []

    starts: list[int] = []
    if source[:1] in "[{":
        starts.append(0)
    if scan_embedded:
        starts.extend(i for i, char in enumerate(source) if char in "[{" and i not in starts)

    values: list[Any] = []
    seen: set[str] = set()
    for start in starts:
        try:
            value, _ = decoder.raw_decode(source[start:])
        except json.JSONDecodeError:
            continue
        signature = json.dumps(value, ensure_ascii=False, sort_keys=True)
        if signature not in seen:
            values.append(value)
            seen.add(signature)
    return values


def select_answer_candidate(candidates: list[list[str]], n_expected: int | None) -> list[str] | None:
    if not candidates:
        return None
    if n_expected is not None:
        for answers in candidates:
            if len(answers) == n_expected:
                return answers
    return candidates[0]


def has_expected_count(answers: list[str], n_expected: int | None) -> bool:
    return n_expected is None or len(answers) == n_expected


def extract_json_object(text: str) -> dict[str, Any] | None:
    """Try to recover a JSON object from a model response."""
    sources = [text, *FENCED_BLOCK_RE.findall(text)]
    for source in sources:
        for value in iter_json_values(source, scan_embedded=True):
            if isinstance(value, dict):
                return value
    return None


def parse_json_answers(raw_text: str, n_expected: int | None, scan_embedded: bool) -> list[str] | None:
    sources = [raw_text, *FENCED_BLOCK_RE.findall(raw_text)]
    candidates: list[list[str]] = []
    for source in sources:
        for value in iter_json_values(source, scan_embedded=scan_embedded):
            answers = answers_from_json_value(value)
            if answers is not None:
                candidates.append(answers)
    return select_answer_candidate(candidates, n_expected)


def parse_list_item_answers(raw_text: str) -> list[str]:
    answers: list[str] = []
    for line in raw_text.splitlines():
        match = LIST_ITEM_RE.match(line)
        if match:
            answer = clean_answer_text(match.group(1))
            if answer:
                answers.append(answer)
    return answers


def parse_plain_line_answers(raw_text: str, n_expected: int | None) -> list[str]:
    lines: list[str] = []
    for line in raw_text.splitlines():
        answer = clean_answer_text(line)
        if not answer or answer.startswith("```") or HEADER_LINE_RE.match(answer) or JSON_SCAFFOLD_LINE_RE.match(answer):
            continue
        lines.append(answer)

    if n_expected is not None and len(lines) > n_expected:
        filtered = [line for line in lines if not INTRO_LINE_RE.match(line)]
        if len(filtered) >= n_expected:
            return filtered[:n_expected]
    return lines


def normalize_answers(raw_text: str, n_expected: int | None) -> list[str]:
    """Convert model text into a list of answer strings."""
    json_answers = parse_json_answers(raw_text, n_expected, scan_embedded=False)
    if json_answers is not None and has_expected_count(json_answers, n_expected):
        return fit_answer_count(json_answers, n_expected)

    list_answers = parse_list_item_answers(raw_text)
    if list_answers and (json_answers is None or has_expected_count(list_answers, n_expected)):
        return fit_answer_count(list_answers, n_expected)
    if json_answers is not None:
        return fit_answer_count(json_answers, n_expected)

    embedded_json_answers = parse_json_answers(raw_text, n_expected, scan_embedded=True)
    if embedded_json_answers is not None:
        return fit_answer_count(embedded_json_answers, n_expected)

    plain_answers = parse_plain_line_answers(raw_text, n_expected)
    return fit_answer_count(plain_answers, n_expected)


def build_prompt(row: pd.Series, n_expected: int | None) -> list[dict[str, str]]:
    task_type = row.get("task_type", "")
    eval_type = row.get("eval_type", "")
    count_instruction = (
        f"Return exactly {n_expected} answers." if n_expected is not None else "Return one answer per numbered item."
    )
    system = (
        "You solve International Linguistics Olympiad problems using only the data in the problem. "
        "Infer the pattern from the examples. Think silently. Return valid JSON only: "
        "{\"answers\": [\"...\"]}"
    )
    user = f"""Task type: {task_type}
Evaluation type: {eval_type}
{count_instruction}

CONTEXT:
{str(row.get('context', '')).strip()}

QUERY:
{str(row.get('query', '')).strip()}"""
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def validate_model_dir(model_dir: str) -> Path:
    path = Path(model_dir)
    if not path.exists():
        raise FileNotFoundError(
            f"Missing local model folder: {path}. Ship model weights with the repo, "
            "or set MODEL_DIR to an existing local directory. Runtime downloads are disabled."
        )
    if not path.is_dir():
        raise NotADirectoryError(f"MODEL_DIR must be a local directory, got: {path}")
    if not (path / "config.json").exists():
        raise FileNotFoundError(f"Local model folder is missing config.json: {path}")
    return path


def load_model():
    if DUMMY_MODE:
        return None, None
    model_path = validate_model_dir(MODEL_DIR)

    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError as exc:
        raise RuntimeError(
            "Missing required model dependency: torch/transformers. Install dependencies before evaluation; "
            "script.py will not download packages, model weights, or call the internet at runtime."
        ) from exc

    tok = AutoTokenizer.from_pretrained(model_path, local_files_only=True, trust_remote_code=True)
    model_kwargs = {
        "device_map": "auto",
        "local_files_only": True,
        "trust_remote_code": True,
    }
    try:
        model = AutoModelForCausalLM.from_pretrained(model_path, dtype=torch.float16, **model_kwargs).eval()
    except TypeError:
        model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, **model_kwargs).eval()
    return tok, model


def model_input_device(model):
    device = getattr(model, "device", None)
    if device is not None:
        return device
    return next(model.parameters()).device


def build_generation_inputs(tok, model, messages: list[dict[str, str]]):
    try:
        encoded = tok.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
        )
    except TypeError:
        encoded = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")

    device = model_input_device(model)
    if hasattr(encoded, "to"):
        encoded = encoded.to(device)

    if hasattr(encoded, "keys") and "input_ids" in encoded.keys():
        input_ids = encoded["input_ids"]
        return {key: encoded[key] for key in encoded.keys()}, input_ids.shape[-1]

    encoded = encoded.to(device)
    return {"input_ids": encoded}, encoded.shape[-1]


def generate_one(tok, model, messages: list[dict[str, str]]) -> str:
    if DUMMY_MODE:
        # Useful for testing CSV shape without downloading weights.
        return json.dumps({"answers": ["DUMMY"]}, ensure_ascii=False)

    import torch

    generation_inputs, prompt_len = build_generation_inputs(tok, model, messages)
    with torch.no_grad():
        out = model.generate(
            **generation_inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=tok.eos_token_id,
        )
    return tok.decode(out[0][prompt_len:], skip_special_tokens=True).strip()


def main() -> None:
    if not INPUT_CSV.exists():
        raise FileNotFoundError(f"Missing input CSV: {INPUT_CSV}")

    df = pd.read_csv(INPUT_CSV, dtype=str).fillna("")
    tok, model = load_model()
    rows: list[dict[str, str]] = []

    for _, row in df.iterrows():
        n_expected = expected_item_count(str(row.get("query", "")))
        messages = build_prompt(row, n_expected)
        raw = generate_one(tok, model, messages)
        answers = normalize_answers(raw, n_expected)

        record: dict[str, str] = {
            "id": str(row["id"]),
            "pred": json.dumps(answers, ensure_ascii=False),
        }
        rows.append(record)

    pd.DataFrame(rows, columns=["id", "pred"]).to_csv(OUTPUT_CSV, index=False)
    print(f"Wrote {OUTPUT_CSV} with {len(rows)} rows")


if __name__ == "__main__":
    main()