"""AutoModel entry point: render preprocessing plus Qwen's existing scorer."""
from pathlib import Path

from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
from transformers.utils.hub import cached_file

from .qwen3_vl_reranker import Qwen3VLReranker
from .rendering import DocumentImageConfig, DocumentImageRenderer, prepare_render_inputs


class _QwenScorer(Qwen3VLReranker):
    def tokenize(self, pairs, **kwargs):
        inputs = super().tokenize(pairs, **kwargs)
        # Transformers 5 adds modality IDs. Align them with the upstream's
        # finalized input IDs, then convert every model input to tensors.
        if "mm_token_type_ids" in inputs:
            inputs["mm_token_type_ids"] = self.processor.create_mm_token_type_ids(
                inputs["input_ids"].tolist()
            )
        return inputs.convert_to_tensors("pt")


class RenderRank(Qwen3VLForConditionalGeneration):
    """Keep the checkpoint architecture; expose the official process API."""

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        loaded = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
        model = loaded[0] if isinstance(loaded, tuple) else loaded
        asset_kwargs = {
            key: kwargs[key] for key in
            ("cache_dir", "revision", "token", "local_files_only", "subfolder", "force_download")
            if key in kwargs
        }
        model.processor = AutoProcessor.from_pretrained(
            pretrained_model_name_or_path, padding_side="left", **asset_kwargs
        )
        model._render_font_path = cached_file(
            pretrained_model_name_or_path, "Roboto-Regular.ttf", **asset_kwargs
        )
        return loaded

    def process(self, inputs):
        """Return scores in document order; render is only an input conversion."""
        if not isinstance(inputs, dict):
            raise TypeError("inputs must be a dictionary")
        renderer = DocumentImageRenderer(DocumentImageConfig(Path(self._render_font_path)))
        prepared = prepare_render_inputs(inputs, renderer)

        # Bind the official wrapper to the already-loaded model (no second load).
        scorer = _QwenScorer.__new__(_QwenScorer)
        scorer.model = self.model
        scorer.processor = self.processor
        scorer.device = self.device
        scorer.max_length = 32768
        scorer.min_pixels = self.processor.image_processor.size["shortest_edge"]
        scorer.max_pixels = self.processor.image_processor.size["longest_edge"]
        scorer.total_pixels = 10 * 768 * 32 * 32
        scorer.fps = 1
        scorer.num_frames = 64
        scorer.max_frames = 64
        scorer.default_instruction = "Retrieve text relevant to the user's query."
        vocab = self.processor.tokenizer.get_vocab()
        scorer.score_linear = scorer.get_binary_linear(self, vocab["yes"], vocab["no"])
        scorer.score_linear = scorer.score_linear.to(device=self.device, dtype=self.dtype).eval()
        return scorer.process(prepared)