"""One joint row per question. Metadata IDs and gold labels never enter tokens."""
import math
import torch

KINDS = ("choice", "noul", "score")
DEFAULT_NO = "No. The statement or question is not satisfied."
DEFAULT_YES = "Yes. The statement or question is satisfied."


def nonempty(value, name):
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{name} must be a nonempty string")
    return value


def view(record, labeled=False):
    q = record["question"]
    kind = str(q["type"]).lower()
    if kind not in KINDS:
        raise ValueError("Unsupported question type")
    nonempty(record["state_text"], "state_text")
    nonempty(q["text"], "question.text")
    if kind == "noul":
        options = [{"id": "no", "text": q.get("false_criterion", DEFAULT_NO)},
                   {"id": "yes", "text": q.get("true_criterion", DEFAULT_YES)}]
    else:
        options = q["options" if kind == "choice" else "levels"]
    if not isinstance(options, list) or not 2 <= len(options) <= 255:
        raise ValueError("Require 2..255 complete candidate descriptions")
    for o in options:
        nonempty(o["id"], "candidate.id")
        nonempty(o["text"], "candidate.text")
    if len({o["id"] for o in options}) != len(options):
        raise ValueError("Duplicate candidate IDs")
    values = [float(o["value"]) for o in options] if kind == "score" else list(range(len(options)))
    if kind == "score" and (not all(math.isfinite(v) for v in values)
                            or any(a >= b for a, b in zip(values, values[1:]))):
        raise ValueError("Score level values must be finite and strictly increasing")
    target = None
    if labeled:
        gold = record["target"]
        if kind == "noul":
            yes = float(gold["probability"])
            target = [1 - yes, yes]
        elif "probabilities" in gold:
            target = [float(x) for x in gold["probabilities"]]
        elif kind == "choice" and "choice_id" in gold:
            if gold["choice_id"] not in {o["id"] for o in options}:
                raise ValueError("Unknown gold choice_id")
            target = [float(o["id"] == gold["choice_id"]) for o in options]
        else:
            raise ValueError("Score requires an explicit target probability distribution")
        if (len(target) != len(options) or not all(math.isfinite(x) and 0 <= x <= 1 for x in target)
                or abs(sum(target) - 1) > 1e-6):
            raise ValueError("Invalid target probability distribution")
    return kind, options, values, target


class MarkerCollator:
    def __init__(self, tokenizer, max_length=8192, state_truncation="error"):
        if isinstance(max_length, bool) or not isinstance(max_length, int) or max_length < 8:
            raise ValueError("Invalid max_length")
        if state_truncation not in ("error", "right"):
            raise ValueError("state_truncation must be error or right")
        self.tokenizer, self.max_length = tokenizer, max_length
        self.state_truncation = state_truncation
        self.bos = tokenizer.cls_token_id if tokenizer.cls_token_id is not None else tokenizer.bos_token_id
        self.sep = tokenizer.sep_token_id if tokenizer.sep_token_id is not None else tokenizer.eos_token_id
        self.pad, self.marker = tokenizer.pad_token_id, tokenizer.mask_token_id
        if any(v is None for v in (self.bos, self.sep, self.pad, self.marker)):
            raise ValueError("Tokenizer must provide existing BOS/CLS, SEP/EOS, PAD and MASK IDs")

    def tokens(self, text):
        ids = self.tokenizer(text, add_special_tokens=False, truncation=False)["input_ids"]
        if not ids:
            raise ValueError("Empty token span")
        return ids

    def encode(self, record, labeled=False):
        kind, options, values, target = view(record, labeled)
        question = self.tokens(f"{kind} question: {record['question']['text']}")
        ids = [self.bos] + question + [self.sep]
        positions, spans = [], []
        for j, option in enumerate(options):
            # Ordinal rank is meaningful only for Score; opaque external IDs never enter text.
            text = f"level {j}: {option['text']}" if kind == "score" else option["text"]
            desc = self.tokens(text)
            positions.append(len(ids))
            ids.append(self.marker)
            spans.append([len(ids), len(ids) + len(desc)])
            ids.extend(desc + [self.sep])
        state = self.tokens(record["state_text"])
        room = self.max_length - len(ids) - 1
        if room < 1:
            raise ValueError("Complete question and candidates leave no room for state")
        if len(state) > room and self.state_truncation == "error":
            raise ValueError(f"Record {record.get('id')} exceeds {self.max_length} tokens; no implicit truncation")
        kept = state[:room]
        state_span = [len(ids), len(ids) + len(kept)]
        ids.extend(kept + [self.sep])
        return {"ids": ids, "positions": positions, "candidate_spans": spans,
                "state_span": state_span, "kind": kind, "target": target, "values": values,
                "candidate_ids": [o["id"] for o in options], "id": record.get("id"),
                "question_id": record["question"].get("id"), "state_tokens_original": len(state),
                "state_tokens_kept": len(kept), "input_tokens": len(ids)}

    def __call__(self, records, labeled=False, device="cpu"):
        encoded = [self.encode(r, labeled) for r in records]
        if not encoded:
            raise ValueError("Empty batch")
        b, length, k = len(encoded), max(len(e["ids"]) for e in encoded), max(len(e["positions"]) for e in encoded)
        ids = torch.full((b, length), self.pad, dtype=torch.long)
        mask = torch.zeros((b, length), dtype=torch.bool)
        positions = torch.zeros((b, k), dtype=torch.long)
        valid = torch.zeros((b, k), dtype=torch.bool)
        target, values = torch.zeros((b, k)), torch.zeros((b, k))
        for i, e in enumerate(encoded):
            n, count = len(e["ids"]), len(e["positions"])
            ids[i, :n] = torch.tensor(e["ids"])
            mask[i, :n] = True
            positions[i, :count] = torch.tensor(e["positions"])
            valid[i, :count] = True
            values[i, :count] = torch.tensor(e["values"])
            if labeled:
                target[i, :count] = torch.tensor(e["target"])
        batch = {"input_ids": ids, "attention_mask": mask, "marker_positions": positions,
                 "valid_candidates": valid, "targets": target, "values": values,
                 "kind_ids": torch.tensor([KINDS.index(e["kind"]) for e in encoded])}
        return {name: value.to(device) for name, value in batch.items()}, encoded