"""Unified training example for the dense-model pipeline.

An example is one typed decision: a state (text) and a question (`qdef`: type, instructions, criteria) with a target
distribution over its options. Everything is plain JSON so pools can be stored as jsonl and screened, mixed and rendered
(answer_slot.SlotCodec) without any model-specific tokenizer.
"""
import gzip
import hashlib
import json
import os
import sys
from typing import Any, Dict, Iterable, Iterator, List, Optional

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from common import QTYPES, serialize_state  # noqa: E402

DATA_ROOT = "/mnt/datasets/dense-data"


def n_options(qdef: Dict[str, Any]) -> int:
    t = qdef["type"]
    if t == "noul":
        return 2
    return len(qdef["criteria"])


def norm_target(target: List[float], k: int) -> List[float]:
    s = float(sum(target))
    return [float(v) / s for v in target] if s > 0 else [1.0 / k] * k


def make_example(source: str, state: Any, qdef: Dict[str, Any], target: List[float], *, split: str = "train",
                 meta: Optional[Dict[str, Any]] = None, ex_id: Optional[str] = None) -> Dict[str, Any]:
    k = n_options(qdef)
    if len(target) != k:
        raise ValueError(f"{source}: target has {len(target)} entries for {k} options")
    target = norm_target(target, k)
    state = "" if state in (None, {}, []) else serialize_state(state)
    ex = {"source": source, "split": split, "state": state, "qdef": qdef, "target": target,
          "label": max(range(k), key=lambda i: target[i]), "meta": meta or {}}
    ex["id"] = ex_id or f"{source}:{content_hash(ex)[:16]}"
    return ex


def content_hash(ex: Dict[str, Any]) -> str:
    blob = json.dumps([ex["state"], ex["qdef"]], ensure_ascii=False, sort_keys=True)
    return hashlib.sha1(blob.encode("utf-8")).hexdigest()


def codec_item(ex: Dict[str, Any]) -> Dict[str, Any]:
    """The view answer_slot.SlotCodec.render expects."""
    return {**ex, "qtype": QTYPES[ex["qdef"]["type"]]}


def write_jsonl(path: str, rows: Iterable[Dict[str, Any]]) -> int:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    opener = gzip.open if path.endswith(".gz") else open
    n = 0
    with opener(path, "wt", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False, separators=(",", ":")) + "\n")
            n += 1
    return n


def read_jsonl(path: str) -> Iterator[Dict[str, Any]]:
    opener = gzip.open if path.endswith(".gz") else open
    with opener(path, "rt", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                yield json.loads(line)