"""One encoding contract for preparation, training and inference.

Only placeholders appear at answer positions. Targets never enter the rendered prompt.
The code vocabulary follows AMOS's letters followed by single-token uppercase pairs.
"""

import itertools
import math
import random
import string

from .common import digest

ANSWER = "<answer>"
SYSTEM = (
    "You make decisions about a case. Read the case, then answer every question "
    "with the code of the option that is most likely right."
)


class Codec:
    def __init__(self, tokenizer, max_length=2048, max_options=255):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.max_options = max_options
        self.original_vocab_size = len(tokenizer)
        if ANSWER not in tokenizer.get_added_vocab():
            tokenizer.add_special_tokens({"additional_special_tokens": [ANSWER]})
        self.marker = tokenizer.convert_tokens_to_ids(ANSWER)
        self.pad = tokenizer.pad_token_id
        if self.pad is None:
            self.pad = tokenizer.eos_token_id
        candidates = list(string.ascii_uppercase + string.ascii_lowercase)
        candidates += ["".join(p) for p in itertools.product(string.ascii_uppercase, repeat=2)]
        self.codes, self.code_ids = [], []
        for code in candidates:
            ids = tokenizer.encode(code, add_special_tokens=False)
            if (
                len(ids) == 1
                and ids[0] not in self.code_ids
                and ids[0] not in tokenizer.all_special_ids
            ):
                self.codes.append(code)
                self.code_ids.append(ids[0])
            if len(self.codes) == max_options:
                break
        if len(self.codes) != max_options:
            raise ValueError(f"Only {len(self.codes)} distinct single-token option codes available")
        self.specials = sorted(tokenizer.get_added_vocab(), key=len, reverse=True)

    def defuse(self, text):
        if not isinstance(text, str):
            raise ValueError("State, instructions, keys and descriptions must be strings")
        for token in self.specials:
            if token in text:
                # Break the token, preserving visible evidence without giving it control-token semantics.
                replacement = (
                    token.replace("<", "‹", 1) if "<" in token else token[0] + "\u200b" + token[1:]
                )
                text = text.replace(token, replacement)
        return text

    def options(self, question):
        typ = question.get("type")
        criteria = question.get("criteria")
        if typ == "noul":
            criteria = criteria or {}
            if not isinstance(criteria, dict):
                raise ValueError("noul criteria must be a map")
            return [
                ("true", "yes" + (": " + str(criteria["true"]) if criteria.get("true") else "")),
                ("false", "no" + (": " + str(criteria["false"]) if criteria.get("false") else "")),
            ]
        if typ == "choice" and isinstance(criteria, dict):
            # Sequential numeric keys are anonymous IDs, already represented by the
            # code-to-key map. Preserve descriptive names alongside distinct rubrics.
            anonymous = set(criteria) == {str(i) for i in range(len(criteria))}
            return [
                (str(k), str(k) if v is None or v == k else str(v) if anonymous else f"{k}: {v}")
                for k, v in criteria.items()
            ]
        if typ == "score" and isinstance(criteria, list):
            return [(str(i), str(value)) for i, value in enumerate(criteria)]
        raise ValueError(f"Unsupported or malformed decision type: {typ}")

    def encode(self, state, questions, targets=None, option_keys=None, seed=None):
        if not isinstance(questions, dict) or not questions:
            raise ValueError("At least one typed question is required")
        rng = random.Random(seed)
        parts = ["Case:\n", self.defuse(state), "\n\nQuestions:\n"]
        keys, aligned, types = [], [], []
        names = list(questions)
        for index, (name, question) in enumerate(questions.items(), 1):
            if not isinstance(question, dict):
                raise ValueError("Each question must be an object")
            pairs = self.options(question)
            typ = question["type"]
            if not 2 <= len(pairs) <= self.max_options:
                raise ValueError("Each decision needs between 2 and 255 options")
            if len({k for k, _ in pairs}) != len(pairs):
                raise ValueError("Duplicate option keys")
            original = [k for k, _ in pairs]
            values = None
            if targets is not None:
                supplied = option_keys[name]
                raw = targets[name]
                if len(raw) != len(supplied) or set(supplied) != set(original):
                    raise ValueError("Option/target keys do not match the question")
                if (
                    any(not math.isfinite(float(v)) or v < 0 for v in raw)
                    or abs(sum(raw) - 1) > 1e-5
                ):
                    raise ValueError("Invalid probability target")
                values = dict(zip(supplied, raw, strict=True))
            if seed is not None and typ == "choice":
                rng.shuffle(pairs)
            hint = " (levels from lowest to highest)" if typ == "score" else ""
            parts.append(f"\nq{index}. {self.defuse(question['instructions'])}{hint}")
            for rank, (_, text) in enumerate(pairs):
                parts.append(f"\n{self.codes[rank]}) {self.defuse(text)}")
            keys.append([k for k, _ in pairs])
            types.append(typ)
            if values is not None:
                aligned.append([float(values[k]) for k, _ in pairs])
        messages = [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": "".join(parts)},
            {
                "role": "assistant",
                "content": "\n".join(f"q{i}: {ANSWER}" for i in range(1, len(names) + 1)),
            },
        ]
        text = self.tokenizer.apply_chat_template(messages, tokenize=False, enable_thinking=False)
        ids = self.tokenizer.encode(text, add_special_tokens=False)
        positions = [i - 1 for i, value in enumerate(ids) if value == self.marker]
        if len(positions) != len(names) or any(p < 0 for p in positions):
            raise ValueError("Exactly one safe answer placeholder is required per question")
        if len(ids) > self.max_length:
            raise ValueError(f"Rendered request has {len(ids)} tokens, limit {self.max_length}")
        return {
            "input_ids": ids,
            "positions": positions,
            "question_names": names,
            "option_keys": keys,
            "types": types,
            "targets": aligned,
            "input_hash": digest({"state": state, "questions": questions}),
        }

    def config(self):
        return {
            "max_length": self.max_length,
            "max_options": self.max_options,
            "marker_id": self.marker,
            "codes": self.codes,
            "code_ids": self.code_ids,
            "format_version": 1,
        }