"""Token layout for an evidence prefix and independent decision candidates."""

from dataclasses import dataclass

from .candidate_tree import CandidateTree
from .decision_utils import MARKER, literal, locate_markers


@dataclass(frozen=True)
class DecisionLayout:
    tree: CandidateTree
    anchors: tuple[int, ...]
    question_positions: tuple[tuple[int, ...], ...]
    lexical_ids: tuple[tuple[int, ...], ...]
    candidate_questions: tuple[int, ...]
    question_types: tuple[int, ...]

    def to_dict(self) -> dict:
        return {
            "prefix_length": self.tree.prefix_length,
            "branch_lengths": self.tree.branch_lengths,
            "anchors": self.anchors,
            "question_positions": self.question_positions,
            "lexical_ids": self.lexical_ids,
            "candidate_questions": self.candidate_questions,
            "question_types": self.question_types,
        }

    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            CandidateTree(data["prefix_length"], tuple(data["branch_lengths"])),
            tuple(data["anchors"]), tuple(tuple(row) for row in data["question_positions"]),
            tuple(tuple(row) for row in data["lexical_ids"]),
            tuple(data["candidate_questions"]), tuple(data["question_types"]),
        )

    def validate(self, length: int):
        count = len(self.tree.branch_lengths)
        if self.tree.length != length or not (len(self.anchors) == len(self.lexical_ids) == len(self.candidate_questions) == count):
            raise ValueError("Decision layout does not match the packed sequence")
        if not self.question_types or len(self.question_positions) != len(self.question_types):
            raise ValueError("Decision question metadata is incomplete")
        if any(kind not in (0, 1, 2) for kind in self.question_types):
            raise ValueError("Invalid decision question type")
        if set(self.candidate_questions) != set(range(len(self.question_types))):
            raise ValueError("Every question requires candidate branches")
        start = self.tree.prefix_length
        for index, size in enumerate(self.tree.branch_lengths):
            if not start <= self.anchors[index] < start + size or not self.lexical_ids[index]:
                raise ValueError("Invalid decision candidate metadata")
            start += size
        for index, positions in enumerate(self.question_positions):
            branch = self.candidate_questions.index(index)
            start = self.tree.prefix_length + sum(self.tree.branch_lengths[:branch])
            end = start + self.tree.branch_lengths[branch]
            if not positions or any(position < start or position >= end for position in positions):
                raise ValueError("Question positions must belong to its first candidate branch")


def pack_candidates(tokenizer, prefix: list[int], questions: dict) -> tuple[list[int], DecisionLayout]:
    ids = list(prefix)
    lengths, anchors, positions, lexical, owners, kinds = [], [], [], [], [], []
    marker = tokenizer.encode(MARKER.lstrip("\n"), add_special_tokens=False)
    for question_index, question in enumerate(questions.values()):
        kinds.append({"predicate": 0, "choice": 1, "score": 2}[question["type"]])
        for candidate_index, description in enumerate(question["criteria"].values()):
            text = literal({"question": question["instructions"], "criterion": description}) + MARKER
            encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
            start = len(ids)
            branch = encoded["input_ids"]
            ids.extend(branch)
            lengths.append(len(branch))
            anchors.append(start + locate_markers(branch, marker, 1)[0])
            owners.append(question_index)
            lexical.append(tuple(tokenizer.encode(literal(description), add_special_tokens=False)))
            if candidate_index == 0:
                phrase = literal(question["instructions"])
                begin = len('{"question": ')
                end = begin + len(phrase)
                positions.append(tuple(start + i for i, (a, b) in enumerate(encoded["offset_mapping"]) if a < end and b > begin))
    layout = DecisionLayout(CandidateTree(len(prefix), tuple(lengths)), tuple(anchors), tuple(positions), tuple(lexical), tuple(owners), tuple(kinds))
    layout.validate(len(ids))
    return ids, layout