"""IOL-AI 2026 submission entrypoint.

Runs on the platform's T4 (16 GB), offline, within a 30-minute limit. Reads the
hidden test set from /tmp/data/test.csv and writes submission.csv (id,pred) to
the working directory. Model weights ship inside this repo and load from ".".

Baseline model: Qwen/Qwen2.5-1.5B-Instruct. The score lever here is the harness
(task-aware prompts + robust parsing) in iol_harness.py, not the model size.
"""

import os
import time

# Must be set BEFORE importing transformers: no network at run time.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import json

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

import iol_harness as H

MODEL_ID = "."  # weights are shipped in this repo (offline)
TEST_CSV = "/tmp/data/test.csv"
OUT_CSV = "submission.csv"

TIME_LIMIT_S = 30 * 60
SAFETY_S = 26 * 60  # stop generating past this; still write what we have
_START = time.monotonic()

# Optional per-repo harness config (reasoning-model repos ship one to turn on CoT).
CFG = {}
if os.path.exists("harness_cfg.json"):
    try:
        CFG = json.load(open("harness_cfg.json"))
    except Exception:
        CFG = {}
COT = bool(CFG.get("cot", False))
MAX_CAP = int(CFG.get("max_new_cap", 1024))


def elapsed():
    return time.monotonic() - _START


def load_model():
    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,  # T4 has no bfloat16
        device_map="auto",
        trust_remote_code=True,
    ).eval()
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    return tok, model


def generate(tok, model, messages, max_new_tokens):
    inputs = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(model.device)
    with torch.no_grad():
        out = model.generate(
            inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,  # greedy => reproducible
            num_beams=1,
            pad_token_id=tok.pad_token_id,
        )
    return tok.decode(out[0][inputs.shape[-1] :], skip_special_tokens=True).strip()


def main():
    df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
    total = len(df)
    print(f"loaded {total} problems from {TEST_CSV}", flush=True)

    tok, model = load_model()
    print(f"model loaded at {elapsed():.0f}s", flush=True)

    records = []
    for i, row in df.iterrows():
        rid = row["id"]
        context = row.get("context", "")
        query = row.get("query", "")
        task_type = row.get("task_type", "")
        n = H.count_items(context, query, task_type)

        preds = [""] * n  # safe default so every item has a slot
        if elapsed() < SAFETY_S:
            try:
                messages = H.build_messages(context, query, task_type, n, cot=COT)
                budget = H.max_new_tokens_for(task_type, n, cot=COT, cap=MAX_CAP)
                text = generate(tok, model, messages, budget)
                preds = H.parse_answers(text, n)
            except Exception as exc:  # never let one row sink the whole run
                print(f"[warn] id={rid} failed: {exc!r}", flush=True)
        else:
            print(f"[warn] time budget hit; blanking id={rid}", flush=True)

        records.append({"id": rid, "pred": json.dumps(preds, ensure_ascii=False)})
        if (i + 1) % 5 == 0 or (i + 1) == total:
            print(f"{i + 1}/{total} done  ({elapsed():.0f}s)", flush=True)

    pd.DataFrame(records, columns=["id", "pred"]).to_csv(OUT_CSV, index=False)
    print(f"wrote {OUT_CSV} ({len(records)} rows) at {elapsed():.0f}s", flush=True)


if __name__ == "__main__":
    main()