"""Audited segmented formatting; tokenizer is supplied by the backend.

Derived from Apache-2.0 Decision 2.0 Lux native segments/encode. See NOTICE.
Labels, provenance and rationale are never accepted by this interface.
"""
from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Protocol

from .contracts import DecisionInput


def canonical(value: object) -> str:
    return json.dumps(value, sort_keys=True, ensure_ascii=False, separators=(",", ":"), allow_nan=False)


class Tokenizer(Protocol):
    def encode(self, text: str, *, add_special_tokens: bool = False) -> list[int]: ...


def segments(item: DecisionInput) -> tuple[str, list[str], str]:
    state = item.state if isinstance(item.state, str) else canonical(item.state)
    prefix = f"Context:\n{state}\n\nTask type: {item.task_type}\nQuestion:\n{item.instructions}\nOptions:"
    options = ["\n<option>\n" + canonical(o.model_dump()) + "\n</option>" for o in item.options]
    suffix = "\n\nSelect the single option best supported by the context and instructions.\nDecision:"
    return prefix, options, suffix


@dataclass(frozen=True)
class EncodedInput:
    ids: tuple[int, ...]
    candidate_positions: tuple[int, ...]
    query_position: int
    keys: tuple[str, ...]


def encode(item: DecisionInput, tokenizer: Tokenizer, max_length: int) -> EncodedInput:
    prefix, options, suffix = segments(item)
    ids = list(tokenizer.encode(prefix, add_special_tokens=False))
    positions = []
    for option in options:
        part = tokenizer.encode(option, add_special_tokens=False)
        if not part:
            raise ValueError("empty tokenized option")
        ids.extend(part)
        positions.append(len(ids) - 1)
    tail = tokenizer.encode(suffix, add_special_tokens=False)
    if not tail:
        raise ValueError("empty tokenized query")
    ids.extend(tail)
    if len(ids) > max_length:
        raise ValueError(f"rendered input has {len(ids)} tokens; limit is {max_length}; no truncation")
    return EncodedInput(tuple(ids), tuple(positions), len(ids) - 1, tuple(o.key for o in item.options))