"""Entrypoint: read /tmp/data/test.csv, write submission.csv (id, pred, explanation)."""

from __future__ import annotations

import os
from pathlib import Path

# ===== CHANGE HERE — your model (must fit the T4's ~15 GB) =====
# "." for the real submission (this repo ships Qwen2.5-14B-Instruct-AWQ at the
# root); a Hub name (e.g. "Qwen/Qwen2.5-14B-Instruct-AWQ") while testing on
# Colab.
MODEL_ID = "."
# None => the pipeline picks per test size (deep mode 2048 for a small test set,
# coverage mode 1024 for a large one). Set an int to force it.
MAX_NEW_TOKENS = None
LLM_BATCH = 6          # puzzles gathered per checkpoint cycle (the client then
                       # sub-batches by token budget to fit the T4)

# Skip the LLM and emit the symbolic-only baseline. Diagnostic; leave False.
SYMBOLIC_ONLY = False

# LLM pass uses a minimal prompt: no scaffold injection, no chain-of-thought.
# Set IOL_LEAN=0 for the scaffolded path.
LEAN_MODE = os.environ.get("IOL_LEAN", "1") == "1"

# Answer match_letters via the free-form LLM pass. Set IOL_MATCH_ASSIGN=1 to use
# the assignment solver instead.
MATCH_ASSIGNMENT = os.environ.get("IOL_MATCH_ASSIGN", "0") == "1"

# Generation batch size. 1 = one prompt at a time, no padding. Larger batches
# are faster but pad to the longest prompt. Set IOL_GEN_BATCH to change.
GEN_BATCH_SIZE = int(os.environ.get("IOL_GEN_BATCH", "1"))

# Light greedy-anchored self-consistency: N sampled passes that can only
# displace the greedy answer on genuine agreement. Budget-gated. 0 disables.
VOTE_SAMPLES = int(os.environ.get("IOL_VOTE_SAMPLES", "2"))
VOTE_TEMP = float(os.environ.get("IOL_VOTE_TEMP", "0.5"))

# Optional segmentation hint in the prompt. Off by default. Set IOL_HINT=1.
HINT = os.environ.get("IOL_HINT", "0") == "1"

# The eval sandbox has no internet; only go offline when loading local
# weights so Colab testing with a Hub MODEL_ID still downloads normally.
if MODEL_ID == "." or Path(MODEL_ID).exists():
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
# reduce CUDA fragmentation on the T4 (must be set before torch initializes)
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import csv
import json
import sys

sys.path.insert(0, str(Path(__file__).resolve().parent))

from solver.budget import Budget
from solver.llm import load_client
from solver.pipeline import run_pipeline

TEST_CSV = "/tmp/data/test.csv"
OUT_CSV = "submission.csv"

csv.field_size_limit(min(sys.maxsize, 2 ** 31 - 1))


def read_rows(path: str):
    with open(path, newline="", encoding="utf-8") as f:
        return [{k: (v or "") for k, v in row.items()} for row in csv.DictReader(f)]


def write_submission(results, out_path: str) -> None:
    """Atomic write (tmp + rename) so a crash mid-write never leaves a
    truncated submission.csv."""
    tmp = out_path + ".tmp"
    with open(tmp, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
        writer.writeheader()
        for r in results:
            writer.writerow({
                "id": r.row_id,
                "pred": json.dumps([str(a).strip() or "?" for a in r.answers],
                                   ensure_ascii=False),
                "explanation": r.explanation,
            })
    os.replace(tmp, out_path)


def main(test_path: str = TEST_CSV, out_path: str = OUT_CSV) -> None:
    budget = Budget()
    rows = read_rows(test_path)
    try:
        if SYMBOLIC_ONLY:
            from solver.llm import NullClient
            client = NullClient()
            print("SYMBOLIC_ONLY: skipping the LLM; submitting the symbolic "
                  "baseline", flush=True)
        else:
            client = load_client(MODEL_ID)
            if hasattr(client, "batch_size"):
                client.batch_size = GEN_BATCH_SIZE
        # checkpoint after the symbolic pass and every LLM batch: a crash at
        # any later point still leaves a complete submission on disk
        results = run_pipeline(rows, client, budget,
                               llm_batch=LLM_BATCH, max_new_tokens=MAX_NEW_TOKENS,
                               checkpoint=lambda rs: write_submission(rs, out_path),
                               lean=LEAN_MODE,
                               use_match_assignment=MATCH_ASSIGNMENT,
                               vote_samples=VOTE_SAMPLES, vote_temp=VOTE_TEMP,
                               hint=HINT)
        write_submission(results, out_path)
        print(f"wrote {out_path}: {len(results)} rows in {budget.elapsed():.1f}s",
              flush=True)
    except BaseException as e:
        # last resort: if the pipeline itself died before the first
        # checkpoint, emit query echoes — an empty pred is a zero row
        if not Path(out_path).exists():
            from solver.pipeline import PuzzleResult
            fallback = [PuzzleResult(str(r.get("id", i)),
                                     [str(r.get("query", "?")).strip() or "?"],
                                     "- fallback")
                        for i, r in enumerate(rows)]
            write_submission(fallback, out_path)
            print(f"pipeline failed ({type(e).__name__}); wrote fallback "
                  f"{out_path}", flush=True)
        raise


if __name__ == "__main__":
    args = sys.argv[1:]
    main(args[0] if args else TEST_CSV, args[1] if len(args) > 1 else OUT_CSV)