"""Compact ("q3s", short-prompt) record encoding for the Glint student.

Same model inputs as glint/schema.py `_encode_with_boxes` (system prompt, image, state, suffix, spans,
box fields, image_coords), but the schema is one line per field instead of the Clef FIELD/ID/TYPE/
INSTRUCTION/ALLOWED OPTIONS block:

    FIELDS:
    screen: <one-line question> [confirm_dialog|empty_state|...]
    has_error: <one-line question> [true|false]
    box_primary_button: <short element description>

question span = the question text, option span = the option id tokens only (descriptions are not rendered).
Selected per record with record["prompt_style"] == "compact" (written by scripts/make_short.py; data.to_record
passes it through). Records without it are encoded exactly as before.
"""
from __future__ import annotations
from typing import Any
from .joint_schema_model import (EncodedQuestion, QUESTION_TYPES, SYSTEM_PROMPT, _encode_media, _tokens,
                                 question_options, render)
from .box import is_box_question
from .schema import EncodedBoxQuestion, GlintEncodedRecord, IMAGE_PAD_TOKEN, image_token_coords

PROMPT_STYLE = "compact"


def encode_compact_record(tokenizer: Any, record: dict[str, Any], max_length: int = 16384,
                          max_state_tokens: int | None = None, processor: Any | None = None) -> GlintEncodedRecord:
    schema_ids = _tokens(tokenizer, "\n\nFIELDS:\n")
    questions: list[EncodedQuestion] = []
    box_questions: list[EncodedBoxQuestion] = []
    for question_id, question in record["questions"].items():
        schema_ids.extend(_tokens(tokenizer, f"{question_id}: "))
        question_start = len(schema_ids)
        instructions = question.get("instructions")
        if instructions is None or instructions == "":
            instructions = str(question_id)
        schema_ids.extend(_tokens(tokenizer, render(instructions)))
        question_end = len(schema_ids)
        if is_box_question(question):
            schema_ids.extend(_tokens(tokenizer, "\n"))
            box_questions.append(EncodedBoxQuestion(str(question_id), (question_start, question_end)))
            continue
        schema_ids.extend(_tokens(tokenizer, " ["))
        option_spans: list[tuple[int, int]] = []
        option_ids: list[str] = []
        for option_index, (option_id, _description) in enumerate(question_options(question)):
            if option_index:
                schema_ids.extend(_tokens(tokenizer, "|"))
            option_start = len(schema_ids)
            schema_ids.extend(_tokens(tokenizer, option_id))
            option_spans.append((option_start, len(schema_ids)))
            option_ids.append(option_id)
        schema_ids.extend(_tokens(tokenizer, "]\n"))
        questions.append(EncodedQuestion(str(question_id), QUESTION_TYPES[str(question["type"])],
                                         (question_start, question_end), tuple(option_spans), tuple(option_ids)))

    prefix_ids = _tokens(tokenizer, f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\nSTATE:\n")
    suffix_ids = _tokens(tokenizer, "\n<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nJOINT SCHEMA DECISIONS:")
    media_ids, media = _encode_media(processor, record)
    if media is not None:
        media["token_offset"] = len(prefix_ids)
        prefix_ids = prefix_ids + media_ids
    state_ids = _tokens(tokenizer, render(record["state"]))
    if max_state_tokens is not None:
        state_ids = state_ids[:max_state_tokens]
    fixed_length = len(prefix_ids) + len(schema_ids) + len(suffix_ids)
    if fixed_length > max_length:
        raise ValueError(f"schema requires {fixed_length} tokens before state; maximum is {max_length}")
    state_ids = state_ids[: max_length - fixed_length]
    off = len(prefix_ids) + len(state_ids)
    shifted = tuple(EncodedQuestion(q.question_id, q.question_type, (q.question_span[0] + off, q.question_span[1] + off),
                                    tuple((s + off, e + off) for s, e in q.option_spans), q.option_ids) for q in questions)
    shifted_box = tuple(EncodedBoxQuestion(q.question_id, (q.question_span[0] + off, q.question_span[1] + off))
                        for q in box_questions)
    input_ids = tuple(prefix_ids + state_ids + schema_ids + suffix_ids)
    if not input_ids or not (shifted or shifted_box):
        raise ValueError("record produced no model input or questions")
    coords = None
    if media is not None and "image_grid_thw" in media:
        pad_id = tokenizer.convert_tokens_to_ids(IMAGE_PAD_TOKEN)
        merge = getattr(getattr(processor, "image_processor", None), "merge_size", 2) or 2
        coords = image_token_coords(input_ids, pad_id, media["image_grid_thw"], merge)
    return GlintEncodedRecord(input_ids=input_ids, questions=shifted, record_id=str(record.get("id", "unknown")),
                              media=media, box_questions=shifted_box, image_coords=coords)