"""Shared pieces for fitting and answering the `decisions` config of JS/TS Code Review Decisions.

Every model script reads cases with `load_split` and writes answers with `PredictionWriter`, in the JSONL shape that
`score_decisions.py` scores: {"id", "answers": {name: System One answer}, "latency_ms"}.
"""
from __future__ import annotations

import json
import os
import random
import time
from pathlib import Path

import pyarrow.parquet as pq

HERE = Path(__file__).resolve().parent
DATA = Path(os.environ.get("DECISIONS_DATA", HERE.parent / "js-ts-code-review-decisions-3000"))


def load_split(split: str, limit: int | None = None, seed: int = 0) -> list[dict]:
    """Cases of one split: id, task_slug, state (task, language, code), questions, and gold label per question.
    `limit` keeps whole tasks (all 4 variants) chosen at random, so a small run still has clean and flawed code."""
    table = pq.read_table(DATA / "decisions" / f"{split}.parquet", columns=["id", "task_slug", "state", "questions", "gold"])
    rows = table.to_pylist()
    if limit is not None and limit < len(rows):
        tasks = sorted({r["task_slug"] for r in rows})
        random.Random(seed).shuffle(tasks)
        keep = set(tasks[: max(1, limit // 4)])
        rows = [r for r in rows if r["task_slug"] in keep]
    return [{"id": r["id"], "task_slug": r["task_slug"], "state": json.loads(r["state"]),
             "questions": json.loads(r["questions"]),
             "gold": {name: g["label"] for name, g in json.loads(r["gold"]).items()}} for r in rows]


def options_of(question: dict) -> list[str]:
    """The option labels `score_decisions.py` scores, in order: true/false, choice keys, or level indexes."""
    if question["type"] == "noul":
        return ["true", "false"]
    if question["type"] == "choice":
        return list(question["criteria"])
    return [str(i) for i in range(len(question["criteria"]))]


def gold_index(question: dict, label: str) -> int:
    return options_of(question).index(label)


def to_answer(question: dict, probs: list[float]) -> dict:
    """A System One answer from probabilities in `options_of` order (a noul's are [true, false])."""
    probs = [float(p) for p in probs]
    if question["type"] == "noul":
        return {"type": "noul", "noul": probs[0]}
    return {"type": question["type"], "probabilities": dict(zip(options_of(question), probs))}


def state_text(state: dict) -> str:
    """One string for models that take plain text: the task, the language, then the code."""
    return f"Task: {state['task']}\nLanguage: {state['language']}\n\n{state['code']}"


class PredictionWriter:
    """Appends one answered case per line and flushes it, so a stopped run resumes where it left off."""

    def __init__(self, path: str | Path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.done: set[str] = set()
        if self.path.exists():
            for line in self.path.read_text().splitlines():
                try:
                    self.done.add(json.loads(line)["id"])
                except (json.JSONDecodeError, KeyError):
                    break  # a line cut short by a stop; everything before it is kept
        self._handle = self.path.open("a")

    def pending(self, cases: list[dict]) -> list[dict]:
        return [c for c in cases if c["id"] not in self.done]

    def write(self, case_id: str, answers: dict, latency_ms: float | None = None) -> None:
        record = {"id": case_id, "answers": answers}
        if latency_ms is not None:
            record["latency_ms"] = round(latency_ms, 3)
        self._handle.write(json.dumps(record) + "\n")
        self._handle.flush()
        self.done.add(case_id)

    def close(self) -> None:
        self._handle.close()


def log(path: str | Path | None, **fields) -> None:
    """Prints a JSON line and appends it to `path` when given (training curves, timings)."""
    line = json.dumps({"time": round(time.time(), 1), **fields})
    print(line, flush=True)
    if path:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "a") as handle:
            handle.write(line + "\n")