import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import re
import json
import time
import sys
import subprocess

try:
    import bitsandbytes  # noqa: F401
except ImportError:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"], check=True)
    import bitsandbytes

import torch
import pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

MODEL_ID = "."
TIME_LIMIT_S = 28 * 60
START = time.time()

# Try 8-bit first (much better reasoning fidelity than 4-bit). Falls back to 4-bit
# only if it doesn't fit in memory. Precision matters a lot for symbolic pattern work.
USE_8BIT_FIRST = True


def time_left():
    return TIME_LIMIT_S - (time.time() - START)


def load_model():
    if USE_8BIT_FIRST:
        try:
            bnb_config = BitsAndBytesConfig(load_in_8bit=True)
            tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
            model = AutoModelForCausalLM.from_pretrained(
                MODEL_ID,
                quantization_config=bnb_config,
                device_map="auto",
                local_files_only=True,
            ).eval()
            print("Loaded in 8-bit.", flush=True)
            return tok, model
        except Exception as e:
            print(f"8-bit load failed ({e}), falling back to 4-bit.", flush=True)

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
    )
    tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch.float16,
        local_files_only=True,
    ).eval()
    return tok, model


def generate(messages, max_new_tokens=900, do_sample=False, temperature=0.7):
    ids = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    ).to(model.device)
    gen_kwargs = dict(
        max_new_tokens=max_new_tokens,
        pad_token_id=tok.eos_token_id,
    )
    if do_sample:
        gen_kwargs.update(do_sample=True, temperature=temperature, top_p=0.9)
    else:
        gen_kwargs.update(do_sample=False, temperature=None, top_p=None)

    with torch.no_grad():
        out = model.generate(ids, **gen_kwargs)
    return tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()


def count_items(query: str) -> int:
    nums = re.findall(r"(?m)^\s*(\d+)\s*[.)]\s", query)
    if nums:
        return max(int(n) for n in nums)
    lines = [l for l in query.splitlines() if l.strip()]
    return max(1, len(lines))


def extract_json_object(text: str):
    """Pull the LAST valid JSON object out of a text blob (final-answer JSON should
    come after the reasoning, so we scan from the end)."""
    text = text.strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    # find all {...} candidates, try from the last one backwards
    candidates = list(re.finditer(r"\{.*?\}", text, re.DOTALL))
    for m in reversed(candidates):
        try:
            return json.loads(m.group(0))
        except Exception:
            continue
    # try greedy full-span match as last resort
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if m:
        try:
            return json.loads(m.group(0))
        except Exception:
            pass
    return None


def extract_answers_fallback(text: str, n: int):
    lines = [l.strip() for l in text.splitlines() if l.strip()]
    cleaned = []
    for l in lines:
        l = re.sub(r"^\s*\d+[.)]\s*", "", l)
        l = l.strip(" -\t")
        if l:
            cleaned.append(l)
    if len(cleaned) >= n:
        return cleaned[-n:]  # prefer the tail: likely the final-answer block
    parts = [p.strip() for p in text.replace("\n", ",").split(",") if p.strip()]
    if len(parts) >= n:
        return parts[-n:]
    cleaned = cleaned or parts or [text.strip()]
    while len(cleaned) < n:
        cleaned.append(cleaned[-1] if cleaned else "")
    return cleaned[:n]


def fit_answers(answers, n):
    answers = list(answers)
    if len(answers) < n:
        answers = answers + [answers[-1] if answers else ""] * (n - len(answers))
    return answers[:n]


# Single combined system prompt: reasoning + mandatory self-verification + strict
# final JSON block. This replaces the old two-call reasoning->extraction pipeline,
# roughly halving generation cost per row while forcing the model to check its own
# rule against every given example (the single highest-leverage habit for IOL-style
# problems) before committing to an answer.
SOLVE_SYSTEM = (
    "You are an expert at International Linguistics Olympiad (IOL) problems. You will "
    "be given example data (context) in an unfamiliar language and a query asking you "
    "to translate, match, fill blanks, or convert forms.\n\n"
    "Follow this exact procedure:\n"
    "1. REASONING: Identify morphemes, word order, and correspondences using ONLY the "
    "given examples. Do not assume rules from natural languages you know unless the "
    "data supports them.\n"
    "2. VERIFY: Explicitly re-derive every example in the context using your proposed "
    "rules and confirm each one matches. If any mismatch, revise your rule before "
    "continuing.\n"
    "3. ANSWER: Apply the verified rule to the query.\n"
    "4. Output a final line starting with FINAL_JSON: followed by a single JSON object "
    "with exactly two keys:\n"
    '   "answers": a JSON list of strings, one per numbered item in the query, in the '
    "same order as the query items.\n"
    '   "explanation": 3-6 short bullet points (use \\n between them) summarizing the '
    "key rules used.\n"
    "The FINAL_JSON line must be the last thing you output, and must contain nothing "
    "else on that line besides FINAL_JSON: {...}."
)


def solve_once(problem_text: str, n_items: int, max_new_tokens: int, do_sample=False, temperature=0.7):
    raw = generate(
        [
            {"role": "system", "content": SOLVE_SYSTEM},
            {"role": "user", "content": problem_text + f"\n\n(The query has {n_items} numbered item(s).)"},
        ],
        max_new_tokens=max_new_tokens,
        do_sample=do_sample,
        temperature=temperature,
    )

    obj = None
    m = re.search(r"FINAL_JSON:\s*(\{.*\})", raw, re.DOTALL)
    if m:
        try:
            obj = json.loads(m.group(1))
        except Exception:
            obj = extract_json_object(m.group(1))
    if obj is None:
        obj = extract_json_object(raw)

    answers, explanation = None, ""
    if isinstance(obj, dict):
        a = obj.get("answers")
        if isinstance(a, list):
            answers = [str(x) for x in a]
        explanation = str(obj.get("explanation", "")).strip()

    if answers is None:
        answers = extract_answers_fallback(raw, n_items)
    if not explanation:
        snippet = raw.strip().replace("\n", " ")
        explanation = (snippet[:400] + "...") if len(snippet) > 400 else snippet

    return fit_answers(answers, n_items), explanation, raw


def majority_vote(all_answers, n_items):
    """Per-position majority vote across N sampled runs."""
    final = []
    for i in range(n_items):
        votes = {}
        for ans_list in all_answers:
            v = ans_list[i] if i < len(ans_list) else ""
            votes[v] = votes.get(v, 0) + 1
        final.append(max(votes.items(), key=lambda kv: kv[1])[0])
    return final


def solve_row(context: str, query: str, n_items: int, row_budget_s: float):
    problem_text = f"{context.strip()}\n\n{query.strip()}"

    # Scale tokens to problem size; cap so one bad row can't eat the whole budget.
    approx_len = len(context) + len(query)
    max_new = 700 if approx_len < 1500 else 1000
    max_new = min(max_new, 1100)

    row_start = time.time()

    def row_time_left():
        return row_budget_s - (time.time() - row_start)

    # First (greedy, most reliable) attempt.
    answers, explanation, _ = solve_once(problem_text, n_items, max_new, do_sample=False)

    # If time allows, take 2 more sampled attempts and majority-vote for robustness.
    # This costs roughly the same as the old broken two-pass pipeline, but spends the
    # extra budget on getting the ANSWER right instead of just reformatting it.
    if row_time_left() > max_new * 0.03:  # crude per-token-time heuristic guard
        try:
            extra_runs = [answers]
            for t in (0.5, 0.8):
                if row_time_left() < 20:
                    break
                a2, _, _ = solve_once(problem_text, n_items, max_new, do_sample=True, temperature=t)
                extra_runs.append(a2)
            if len(extra_runs) > 1:
                answers = majority_vote(extra_runs, n_items)
        except Exception:
            pass  # keep the greedy answer if sampling fails

    return answers, explanation


def main():
    global tok, model
    tok, model = load_model()

    df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
    n_rows = len(df)
    out_rows = []

    for i, r in df.iterrows():
        rid = r["id"]
        context, query = r.get("context", ""), r.get("query", "")
        n_items = count_items(query)

        rows_remaining = n_rows - i
        # Reserve a small safety margin (30s) for writing the CSV at the end.
        safe_time_left = max(time_left() - 30, 0)
        row_budget = safe_time_left / max(rows_remaining, 1)

        if safe_time_left < 20:
            # Truly out of time: still try a fast, cheap single greedy pass rather
            # than silently submitting blanks — a bad guess beats an empty string
            # on chrF and exact-match both.
            try:
                answers, explanation, _ = solve_once(
                    f"{context.strip()}\n\n{query.strip()}", n_items, max_new_tokens=300
                )
            except Exception as e:
                answers = [""] * n_items
                explanation = f"(error: {e})"
        else:
            try:
                answers, explanation = solve_row(context, query, n_items, row_budget)
            except Exception as e:
                answers = [""] * n_items
                explanation = f"(error: {e})"

        out_rows.append(
            {
                "id": rid,
                "pred": json.dumps(answers, ensure_ascii=False),
                "explanation": explanation,
            }
        )
        print(
            f"[{i + 1}/{n_rows}] id={rid} n_items={n_items} "
            f"row_budget={row_budget:.0f}s time_left={time_left():.0f}s",
            flush=True,
        )

        # Checkpoint after every row so a hard timeout still leaves a submission
        # instead of losing everything if the process gets killed mid-loop.
        pd.DataFrame(out_rows, columns=["id", "pred", "explanation"]).to_csv(
            "submission.csv", index=False
        )

    print("wrote submission.csv", flush=True)


if __name__ == "__main__":
    main()