#!/usr/bin/env python3
# Historical source snapshot. Comments about accuracy, data origin, or compliance
# are author claims, not verification; see training/README.md and REPRODUCIBILITY.md.
"""
Baseline Inference Runner for Sinhala MMLU Shared Task
Loads all questions from data/ (easy, medium, hard) and outputs predictions.jsonl
Uses Ollama /api/chat with strict structured exam prompting.
Supports:
- Fast direct mode (default), Cross-Lingual CoT mode (--cot), and <think> mode (--think)
- Rapid local validation benchmark (--val_only)
- Automatic checkpointing and resuming
"""

import json
import argparse
import sys
import re
import os
import random
from pathlib import Path
import urllib.request

# Add src to sys.path so prompts can be imported from anywhere
sys.path.insert(0, str(Path(__file__).parent))
from prompts import build_prompt


def extract_answer_digit(text: str) -> str:
    """Extracts the chosen option index (1, 2, 3, 4) from model output text."""
    if not text:
        return "1"
    # If '</think>' is present in content, only parse after '</think>'
    if "</think>" in text:
        text = text.split("</think>")[-1]
    text = text.strip()

    # 1. Explicit 'Answer: X' or 'පිළිතුර: X' or 'Option X' (take last match if CoT text precedes it)
    explicit_matches = re.findall(
        r"(?:පිළිතුර|Answer|Final Answer|Correct Option|නිවැරදි පිළිතුර|තෝරාගැනීම)\s*(?:is|වන්නේ)?\s*[:\s*#(\[]*(?:Option\s*)?([1-4])\b",
        text,
        re.IGNORECASE,
    )
    if explicit_matches:
        return explicit_matches[-1]

    # 2. Match leading digit if response is just a digit
    m = re.match(r"^[(\[]?\s*([1-4])\b", text)
    if m and len(text) <= 8:
        return m.group(1)

    # 3. Match the last standalone digit 1-4 in the response
    matches = re.findall(r"\b([1-4])\b", text)
    if matches:
        return matches[-1]

    m = re.search(r"[1-4]", text)
    if m:
        return m.group(0)

    # Fallback default
    return "1"


def load_all_questions(data_dir: Path):
    """Loads all questions from easy, medium, and hard folders in consistent order."""
    questions = []
    levels = ["easy", "medium", "hard"]

    for lvl in levels:
        folder = data_dir / lvl
        if not folder.exists():
            continue
        files = sorted(list(folder.glob("*.json")))
        for f in files:
            with open(f, "r", encoding="utf-8") as fp:
                data = json.load(fp)
                for item in data:
                    item["_level"] = lvl
                    item["_source_file"] = f.name
                    questions.append(item)

    return questions


def load_val_questions(gold_val_path: Path):
    """Loads questions from the gold validation file."""
    questions = []
    with open(gold_val_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                d = json.loads(line)
                d["_level"] = d["level"]
                questions.append(d)
    return questions


def make_q_key(record: dict) -> str:
    lvl = str(record.get("level", "")).lower().strip()
    subj = str(record.get("subject", "")).strip()
    q_no = str(record.get("q_no", "")).strip()
    return f"{lvl}_{subj}_{q_no}"


def load_existing_checkpoint(output_path: Path):
    """Loads existing predictions from output_file if resuming, indexed by composite key."""
    existing = {}
    if output_path.exists():
        with open(output_path, "r", encoding="utf-8") as fp:
            for line in fp:
                line = line.strip()
                if line:
                    try:
                        record = json.loads(line)
                        key = make_q_key(record)
                        existing[key] = record
                    except Exception:
                        pass
    return existing


SYSTEM_EXAM_PROMPT = (
    "You are an examination AI for the Sri Lankan school curriculum (Sinhala medium). "
    "For every multiple-choice question, analyze the options carefully and answer by replying "
    "with ONLY the single digit of the best option (1, 2, 3, or 4). "
    "Do NOT provide explanations, do NOT include option text, and do NOT write anything other than the single digit."
)

SYSTEM_COT_PROMPT = (
    "You are an expert examiner for the Sri Lankan school curriculum (Sinhala medium). "
    "For each multiple-choice question, write a concise 1-2 sentence English translation and fact-check ('Analysis:'), "
    "then end on a new line with strictly 'Answer: <digit>' (where <digit> is 1, 2, 3, or 4)."
)

SYSTEM_THINK_PROMPT = (
    "You are an expert teacher of the Sri Lankan national school curriculum (grades 6-13), "
    "fluent in Sinhala. You will receive a multiple-choice question written in Sinhala with four "
    "numbered options. Exactly one option is correct. Read the Sinhala carefully, including "
    "negations such as 'නොවන' / 'නොවේ' (NOT). You may reason in English. "
    "Finish with the final answer in the form 'Answer: <number>'."
)

ANSWER_THINK_RE = re.compile(r"(?:Answer|පිළිතුර)\s*[:：]?\s*\**\s*\(?([1-4])\)?", re.IGNORECASE)


def build_think_user_prompt(q: dict) -> str:
    """Builds the Zero-Shot Curriculum Teacher user prompt with Subject & Grade context (72.82% val accuracy)."""
    choices = q.get("choices")
    if choices is None:
        choices = q.get("options", [])
    if isinstance(choices, dict):
        keys = sorted(choices.keys(), key=lambda x: str(x))
        opt_list = [choices[k] for k in keys]
    else:
        opt_list = list(choices) if choices else []

    opts = "\n".join(f"{i}. {c}" for i, c in enumerate(opt_list, 1))
    lvl = str(q.get("_level", q.get("level", "medium"))).lower().strip()
    grade = "6-9" if lvl == "easy" else ("10-11 (O/L)" if lvl == "medium" else "12-13 (A/L)")
    subj = q.get("subject", "General")
    qtext = q.get("question", "").strip()
    return (
        f"විෂය (Subject): {subj} | ශ්‍රේණිය (Grade): {grade}\n\n"
        f"ප්‍රශ්නය (Question):\n{qtext}\n\n"
        f"විකල්ප (Options):\n{opts}\n\n"
        "Which option (1, 2, 3 or 4) is correct? End your reply with 'Answer: <number>'."
    )


_NO_PROXY_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))


def query_ollama_chat(
    model_name: str,
    system_prompt: str,
    user_prompt: str,
    think: bool = False,
    cot: bool = False,
    timeout: int = 240
) -> tuple[str, str]:
    """Sends a chat request to local Ollama /api/chat and returns (raw_text, extracted_digit)."""
    if think:
        num_predict = 2048
    elif cot:
        num_predict = 260
    else:
        num_predict = 15

    options = {
        "temperature": 0.0,
        "num_predict": num_predict,
    }
    if not think:
        options["stop"] = ["\nප්‍රශ්නය:", "\nQuestion:", "<|im_end|>", "<|endoftext|>", "<|im_start|>"]

    req_data = json.dumps({
        "model": model_name,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        "stream": False,
        "think": think,
        "options": options
    }).encode("utf-8")

    req = urllib.request.Request(
        "http://127.0.0.1:11434/api/chat",
        data=req_data,
        headers={"Content-Type": "application/json"}
    )

    with _NO_PROXY_OPENER.open(req, timeout=timeout) as resp:
        res = json.loads(resp.read().decode("utf-8"))
        msg = res.get("message", {})
        content = msg.get("content", "")
        thinking = msg.get("thinking", "")
        if think:
            m = ANSWER_THINK_RE.findall(content) or ANSWER_THINK_RE.findall(thinking)
            if m:
                return content or thinking, m[-1]
            combined = (thinking + "\n" + content).strip()
            return combined, extract_answer_digit(content if content.strip() else thinking)
        return content, extract_answer_digit(content)


def save_merged_full_predictions(
    canonical_questions: list,
    completed_map: dict,
    fallback_map: dict,
    merged_path: Path
):
    """Writes a complete 2,238-line prediction file in canonical order using completed_map + fallback_map."""
    with open(merged_path, "w", encoding="utf-8") as fp:
        for q in canonical_questions:
            q_no = q["q_no"]
            subj = q.get("subject", "general")
            cat = q.get("category", "general")
            lvl = q["_level"].lower()
            q_key = f"{lvl}_{subj}_{q_no}"
            if q_key in completed_map:
                rec = completed_map[q_key]
            elif q_key in fallback_map:
                rec = fallback_map[q_key]
            else:
                rec = {
                    "q_no": q_no,
                    "subject": subj,
                    "category": cat,
                    "level": lvl,
                    "answer": "1"
                }
            fp.write(json.dumps(rec, ensure_ascii=False) + "\n")


def run_inference(
    data_dir: str,
    output_file: str,
    engine: str,
    model_name: str,
    mock: bool,
    max_questions: int = None,
    resume: bool = True,
    val_only: bool = False,
    gold_val_file: str = "data/gold_val_41.jsonl",
    think: bool = False,
    cot: bool = False,
    priority_disputed: bool = False,
    fallback_file: str = "data/predictions_gemma4_cot.jsonl",
    compare_file: str = "data/predictions_gemma4_direct.jsonl",
    merged_output: str = "data/predictions_gemma4_think_full.jsonl"
):
    if val_only:
        val_path = Path(gold_val_file)
        if not val_path.exists():
            print(f"Error: Gold validation file '{gold_val_file}' not found!")
            sys.exit(1)
        all_questions = load_val_questions(val_path)
        print(f"Loaded {len(all_questions)} validation questions from '{gold_val_file}'.")
    else:
        data_path = Path(data_dir)
        if not data_path.exists():
            print(f"Error: Data directory '{data_dir}' not found!")
            sys.exit(1)
        all_questions = load_all_questions(data_path)
        print(f"Loaded {len(all_questions)} questions across easy, medium, and hard.")

    canonical_questions = list(all_questions)

    fallback_map = {}
    compare_map = {}
    if not val_only and Path(fallback_file).exists():
        fallback_map = load_existing_checkpoint(Path(fallback_file))
    if not val_only and Path(compare_file).exists():
        compare_map = load_existing_checkpoint(Path(compare_file))

    if priority_disputed and not val_only and fallback_map and compare_map:
        lvl_rank = {"hard": 0, "medium": 1, "easy": 2}
        def sort_key(q):
            lvl = q["_level"].lower()
            subj = q.get("subject", "general")
            q_no = q["q_no"]
            k = f"{lvl}_{subj}_{q_no}"
            a1 = fallback_map.get(k, {}).get("answer", "1")
            a2 = compare_map.get(k, {}).get("answer", "1")
            is_disputed = 0 if a1 != a2 else 1
            return (is_disputed, lvl_rank.get(lvl, 3))
        all_questions = sorted(all_questions, key=sort_key)
        disputed_cnt = sum(
            1 for q in all_questions
            if fallback_map.get(f"{q['_level'].lower()}_{q.get('subject','general')}_{q['q_no']}", {}).get("answer")
            != compare_map.get(f"{q['_level'].lower()}_{q.get('subject','general')}_{q['q_no']}", {}).get("answer")
        )
        print(f"Priority-Disputed Mode ON: Processing {disputed_cnt} disputed questions first (Hard -> Medium -> Easy)!")

    if max_questions and max_questions > 0:
        all_questions = all_questions[:max_questions]
        print(f"Limiting inference to first {len(all_questions)} questions as requested.")

    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    merged_path = Path(merged_output)

    existing_preds = {}
    if resume and output_path.exists():
        existing_preds = load_existing_checkpoint(output_path)
        print(f"Resuming: Found {len(existing_preds)} existing predictions in '{output_file}'.")

    predictions = []
    completed_map = dict(existing_preds)

    if mock:
        print("Running in MOCK mode (generating valid random predictions)...")
        for q in all_questions:
            q_no = q["q_no"]
            subj = q.get("subject", "general")
            lvl = q["_level"].lower()
            q_key = f"{lvl}_{subj}_{q_no}"
            if resume and q_key in existing_preds:
                predictions.append(existing_preds[q_key])
                continue
            pred_record = {
                "q_no": q_no,
                "subject": subj,
                "category": q.get("category", "general"),
                "level": lvl,
                "answer": str(random.randint(1, 4))
            }
            predictions.append(pred_record)
            completed_map[q_key] = pred_record
    elif engine == "ollama":
        print(f"Running inference with local Ollama model: '{model_name}' via /api/chat (think={think}, cot={cot})...")
        if think:
            sys_prompt = SYSTEM_THINK_PROMPT
        elif cot:
            sys_prompt = SYSTEM_COT_PROMPT
        else:
            sys_prompt = SYSTEM_EXAM_PROMPT

        for i, q in enumerate(all_questions, 1):
            q_no = q["q_no"]
            subj = q.get("subject", "general")
            cat = q.get("category", "general")
            lvl = q["_level"].lower()

            q_key = f"{lvl}_{subj}_{q_no}"
            if resume and q_key in existing_preds:
                predictions.append(existing_preds[q_key])
                continue

            if think:
                prompt = build_think_user_prompt(q)
            else:
                prompt = build_prompt(q)
                if cot:
                    prompt += "\nAnalysis:"
            raw = ""
            try:
                raw, ans_digit = query_ollama_chat(model_name, sys_prompt, prompt, think=think, cot=cot)
            except Exception as e:
                print(f"[{i}/{len(all_questions)}] Ollama error on Q#{q_no} ({subj}): {e}. Using fallback 1.")
                ans_digit = fallback_map.get(q_key, {}).get("answer", "1")

            record = {
                "q_no": q_no,
                "subject": subj,
                "category": cat,
                "level": lvl,
                "answer": str(ans_digit)
            }
            predictions.append(record)
            completed_map[q_key] = record

            if i % 10 == 0 or i == len(all_questions):
                print(f"Progress: {i}/{len(all_questions)} processed (Total solved={len(completed_map)}). Latest ({lvl}/{subj} Q#{q_no}): => {ans_digit}")
                with open(output_path, "w", encoding="utf-8") as fp:
                    for p in predictions:
                        fp.write(json.dumps(p, ensure_ascii=False) + "\n")
                if not val_only and fallback_map:
                    save_merged_full_predictions(canonical_questions, completed_map, fallback_map, merged_path)
    else:
        print(f"Unsupported engine: {engine}")
        sys.exit(1)

    # Write final predictions
    with open(output_path, "w", encoding="utf-8") as fp:
        for p in predictions:
            fp.write(json.dumps(p, ensure_ascii=False) + "\n")

    if not val_only and fallback_map:
        save_merged_full_predictions(canonical_questions, completed_map, fallback_map, merged_path)
        print(f"Complete 2,238-line merged file updated at: {merged_output}")

    print(f"\nDONE! {len(predictions)} predictions written to {output_file}")
    if not val_only:
        target_sub = merged_output if (priority_disputed and fallback_map) else output_file
        print(f"To package and validate for Codabench at any time, run:")
        print(f"  python3 src/package_submission.py {target_sub} submission_gemma4_think.zip\n")


def main():
    parser = argparse.ArgumentParser(description="Sinhala MMLU Baseline Runner")
    parser.add_argument("--data_dir", default="data", help="Path to data directory containing easy, medium, hard")
    parser.add_argument("--output", default="data/predictions.jsonl", help="Output predictions path")
    parser.add_argument("--engine", choices=["ollama"], default="ollama", help="Inference engine")
    parser.add_argument("--model", default="qwen3.5:4b", help="Model name in Ollama")
    parser.add_argument("--mock", action="store_true", help="Generate fast mock predictions to test pipeline")
    parser.add_argument("--max_questions", type=int, default=None, help="Process only first N questions")
    parser.add_argument("--no_resume", action="store_true", help="Do not resume from existing checkpoint")
    parser.add_argument("--val_only", action="store_true", help="Run only on the 41-question local gold benchmark")
    parser.add_argument("--gold_val_file", default="data/gold_val_41.jsonl", help="Gold validation benchmark file")
    parser.add_argument("--think", action="store_true", help="Enable native thinking mode (72.82%% val accuracy on gemma4:e4b)")
    parser.add_argument("--cot", action="store_true", help="Enable concise Cross-Lingual Translate-and-Reason CoT mode")
    parser.add_argument("--priority_disputed", action="store_true", help="Solve disputed questions first and continuously update 2238-line merged output")
    parser.add_argument("--merged_output", default="data/predictions_gemma4_think_full.jsonl", help="Path for continuously updated 2238-line full submission file")
    args = parser.parse_args()

    run_inference(
        args.data_dir,
        args.output,
        args.engine,
        args.model,
        args.mock,
        args.max_questions,
        resume=not args.no_resume,
        val_only=args.val_only,
        gold_val_file=args.gold_val_file,
        think=args.think,
        cot=args.cot,
        priority_disputed=args.priority_disputed,
        merged_output=args.merged_output
    )


if __name__ == "__main__":
    main()