"""BaguettotronVLM processor — self-contained for HuggingFace Hub."""
from __future__ import annotations

import json
import os
from pathlib import Path

from PIL import Image
from transformers import CLIPImageProcessor, ProcessorMixin
from transformers import PreTrainedTokenizerFast

NUM_VISUAL_TOKENS = 256
IMAGE_TOKEN = "<image>"
IMAGE_END_TOKEN = "</image>"


def _assistant_prefix(chat_style: str, enable_thinking: bool | None) -> str:
    """Return the assistant-turn content prefix for a generation prompt.

    chat_style is the default baked into the repo at publish time:
      - "base":   stage 1 — no think tokens
      - "answer": stage 2 — pre-fill </think> so the model skips reasoning
      - "think":  stage 3 — pre-fill <think> to trigger reasoning traces

    enable_thinking overrides chat_style at call-time (stage 3 models can
    toggle thinking on/off dynamically):
      - None  → keep chat_style default
      - True  → "<think>\n"
      - False → "</think>\n"
    """
    if enable_thinking is True:
        return "<think>\n"
    if enable_thinking is False:
        return "</think>\n"
    if chat_style == "think":
        return "<think>\n"
    if chat_style == "answer":
        return "</think>\n"
    return ""


class BaguettotronVLMProcessor(ProcessorMixin):
    """
    Wraps CLIPImageProcessor + Baguettotron tokenizer.

    Expands a single <image> placeholder into NUM_VISUAL_TOKENS consecutive
    <image> token IDs so the model's forward() can replace them with ViT
    features. Builds Qwen-style chat prompts with stage-appropriate
    assistant prefixes.

    Load with::

        from transformers import AutoProcessor
        processor = AutoProcessor.from_pretrained(
            "andreagemelli/Baguettotron-VLM",
            trust_remote_code=True,
        )
    """

    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "CLIPImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(
        self,
        image_processor: CLIPImageProcessor,
        tokenizer: PreTrainedTokenizerFast,
        num_visual_tokens: int = NUM_VISUAL_TOKENS,
        chat_style: str = "answer",
    ):
        super().__init__(image_processor, tokenizer)
        self.num_visual_tokens = num_visual_tokens
        self.chat_style = chat_style
        raw_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
        self.image_token_id: int = raw_id if isinstance(raw_id, int) else int(raw_id[0])

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):  # type: ignore[override]
        processor = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        # Re-register special tokens (some tokenizers drop them on save/load)
        processor.tokenizer.add_special_tokens(
            {"additional_special_tokens": [IMAGE_TOKEN, IMAGE_END_TOKEN]}
        )
        raw_id = processor.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
        processor.image_token_id = raw_id if isinstance(raw_id, int) else int(raw_id[0])

        # Load chat_style from config.json (written by push_to_hub per stage)
        chat_style = "answer"
        try:
            if os.path.isdir(pretrained_model_name_or_path):
                cfg_path = Path(pretrained_model_name_or_path) / "config.json"
                if cfg_path.exists():
                    chat_style = json.loads(cfg_path.read_text()).get(
                        "chat_style", chat_style
                    )
            else:
                from huggingface_hub import hf_hub_download
                cfg_path = hf_hub_download(
                    repo_id=pretrained_model_name_or_path, filename="config.json"
                )
                chat_style = json.loads(Path(cfg_path).read_text()).get(
                    "chat_style", chat_style
                )
        except Exception:
            pass
        processor.chat_style = chat_style
        return processor

    def _format_messages(
        self,
        messages: list[dict],
        add_generation_prompt: bool,
        enable_thinking: bool | None,
    ) -> str:
        parts = [
            f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>" for m in messages
        ]
        text = "\n".join(parts)
        if add_generation_prompt:
            prefix = _assistant_prefix(self.chat_style, enable_thinking)
            text = f"{text}\n<|im_start|>assistant\n{prefix}"
        return text

    def __call__(
        self,
        text: str | None = None,
        messages: list[dict] | None = None,
        image: Image.Image | None = None,
        return_tensors: str = "pt",
        add_generation_prompt: bool = True,
        enable_thinking: bool | None = None,
    ) -> dict:
        """Tokenise text (or a messages list) and optionally preprocess an image.

        Args:
            text: raw prompt string (with a single <image> placeholder).
            messages: alternative to text — list of chat dicts with "role"/"content".
            image: PIL image to preprocess (optional).
            add_generation_prompt: append an <|im_start|>assistant\n prefix.
            enable_thinking: override chat_style for this call.
                - None: keep the stage default (chat_style)
                - True: pre-fill <think>\n (stage 3 reasoning mode)
                - False: pre-fill </think>\n (stage 2 / stage 3 no-think mode)
        """
        if messages is not None:
            text = self._format_messages(
                messages, add_generation_prompt, enable_thinking
            )
        if text is None:
            raise ValueError("Provide either text or messages.")

        expanded = text.replace(IMAGE_TOKEN, IMAGE_TOKEN * self.num_visual_tokens, 1)
        enc = self.tokenizer(
            expanded, return_tensors=return_tensors, add_special_tokens=False
        )
        result = {
            "input_ids": enc["input_ids"],
            "attention_mask": enc["attention_mask"],
        }
        if image is not None:
            pv = self.image_processor(images=image, return_tensors=return_tensors)
            result["pixel_values"] = pv.pixel_values
        else:
            result["pixel_values"] = None
        return result