"""General decision datasets in JSONL, one decision per line:

    {"state": "hotdog", "question": "Is this food?", "choices": ["yes", "no"], "target": "yes"}

Optional fields:
    teacher_distribution  {"yes": 0.97, "no": 0.03}   soft labels from a teacher (LLM, oracle, ...)
    teacher_probs         [0.97, 0.03]                 same, aligned with `choices`
    score                 0.83                         continuous target for the score head
    group                 any id; splits and dedupe work per group (entity, document, episode)

`target` may be the choice text or its index; it may be missing when only a
teacher distribution is given. The teacher is only needed to *build* the data:
distillation stores its probabilities, the deployed model never calls it.
"""
from __future__ import annotations

import json
import zlib
from pathlib import Path
from typing import Iterable

import numpy as np

from .decision_data import DecisionData

NEG = -1e9


def read_jsonl(path: str | Path) -> list[dict]:
    with open(path, encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def write_jsonl(path: str | Path, rows: Iterable[dict]) -> None:
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def _gid(v) -> int:
    if isinstance(v, (int, np.integer)):
        return int(v)
    return zlib.crc32(str(v).encode())


def teacher_probs(row: dict) -> list[float] | None:
    if row.get("teacher_probs") is not None:
        p = [float(x) for x in row["teacher_probs"]]
    elif row.get("teacher_distribution") is not None:
        td = row["teacher_distribution"]
        p = [float(td.get(c, 0.0)) for c in row["choices"]]
    else:
        return None
    s = sum(p)
    if len(p) != len(row["choices"]) or s <= 0:
        raise ValueError(f"bad teacher distribution for row {row.get('question')!r}")
    return [x / s for x in p]


def rows_to_data(rows: list[dict], task: str = "text", task_cfg: dict | None = None) -> DecisionData:
    from minijev.tasks import get_task

    t = get_task(task, task_cfg or {})
    n = len(rows)
    nmax = max((len(r.get("choices") or []) for r in rows), default=1) or 1
    qvocab, cvocab = {}, {}
    qi = np.zeros(n, np.int32)
    ci = np.full((n, nmax), -1, np.int32)
    target = np.full(n, -1, np.int64)
    teacher = np.full((n, nmax), NEG, np.float32)
    score = np.full(n, np.nan, np.float32)
    has_teacher = has_score = False
    for i, r in enumerate(rows):
        qi[i] = qvocab.setdefault(r.get("question", ""), len(qvocab))
        ch = list(r.get("choices") or [])
        if len(set(ch)) != len(ch):
            raise ValueError(f"duplicate choices in row {i}: {ch}")
        for j, c in enumerate(ch):
            ci[i, j] = cvocab.setdefault(c, len(cvocab))
        tg = r.get("target")
        if tg is not None:
            target[i] = tg if isinstance(tg, int) else ch.index(tg)
        p = teacher_probs(r)
        if p is not None:
            has_teacher = True
            teacher[i, : len(p)] = np.log(np.maximum(np.array(p, np.float32), 1e-8))
        if r.get("score") is not None:
            has_score = True
            score[i] = float(r["score"])
    groups = np.array([_gid(r.get("group", r.get("state") or r.get("question", ""))) for r in rows], np.int64)
    hashes = np.array([_gid(f"{r.get('state') or ''}\x00{r.get('question', '')}") for r in rows], np.int64)
    return DecisionData(task=task, state=t.states_to_arrays([r.get("state") or "" for r in rows]),
                        questions=list(qvocab), question_idx=qi, choices=list(cvocab), choice_idx=ci,
                        target=target, group=groups, teacher_logits=teacher if has_teacher else None,
                        hash=hashes, score=score if has_score else None,
                        meta={"source": "jsonl", "rows": n})


def generate_splits(out_dir: str | Path, samples: int = 0, source: str | None = None, seed: int = 0,
                    fractions=(0.9, 0.05, 0.05), task: str = "text", **_) -> dict:
    """`data.generate: {source: file.jsonl}` -> group-split train/val/test .npz."""
    from .decision_data import split_by_group

    data = rows_to_data(read_jsonl(source), task=task)
    splits = split_by_group(data, fractions, seed=seed)
    for k, v in splits.items():
        v.save(Path(out_dir) / f"{k}.npz")
    return {k: len(v) for k, v in splits.items()}