"""Standalone Transformers processor for quicksviewer_8b (one conversation).

pixel_values stays CPU uint8 [sum_frames,3,980,980]. image_spans identify only
64 content embeddings, excluding <image>/</image>. timestamp_spans identify
ordinary digit-token slots to patch after the model determines cube boundaries.
"""
import json
import os
from pathlib import Path
import re
import shutil

import numpy as np
import torch
from transformers import AutoTokenizer
from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin
from transformers.utils.hub import cached_file

from .image_processing_quicksviewer import QuicksviewerImageProcessor, SIGLIP_SIZE
from .video_processing_quicksviewer import (
    QuicksviewerVideoProcessor,
    frame_temporal_ids,
    num_video_cubes,
    timestamp_placeholder_ids,
    timestamp_token_count,
    video_list,
)

VISUAL_START_ID = 151669
VISUAL_END_ID = 151670
VISUAL_CONTENT_ID = 128244
NQ = 64
_MEDIA_PATTERN = re.compile(r"<image>|<video>")
_DEFAULT_TEMPLATE = """{%- for message in messages %}
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\n' }}
{%- if enable_thinking is defined and enable_thinking is false %}
{{- '<think>\n\n</think>\n\n' }}
{%- elif enable_thinking is defined and enable_thinking is true %}
{{- '<think>\n' }}
{%- endif %}
{%- endif %}"""


def _read_json(path, name, hub_kwargs):
    filename = cached_file(path, name, _raise_exceptions_for_missing_entries=False, **hub_kwargs)
    if filename is None:
        return {}
    with open(filename, encoding="utf-8") as stream:
        return json.load(stream)


def _image_list(images):
    if images is None:
        return []
    if isinstance(images, (np.ndarray, torch.Tensor)) and images.ndim == 4:
        return list(images)
    return list(images) if isinstance(images, (list, tuple)) else [images]


def _media_value(item, kind):
    value = item.get(kind, item.get(kind + "_url", item.get("url", item.get("path"))))
    if isinstance(value, dict):
        value = value.get("url", value.get("path"))
    return value


class QuicksviewerProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer", "video_processor"]
    image_processor_class = "AutoImageProcessor"
    video_processor_class = "AutoVideoProcessor"
    tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
    model_input_names = ["input_ids", "attention_mask", "pixel_values", "frames_per_media",
                         "media_is_video", "frame_temporal_ids", "image_spans", "timestamp_spans"]

    def __init__(self, image_processor=None, tokenizer=None, video_processor=None,
                 chat_template=None, vision_config=None, fpq=5, mm_use_thumbnail=True,
                 thumbnail_position="last", emit_text_timestamps=True,
                 disable_cubing=False, max_frames=360, **kwargs):
        if tokenizer is None:
            raise ValueError("A Qwen tokenizer is required")
        vision = vision_config or {}
        if hasattr(vision, "to_dict"):
            vision = vision.to_dict()
        self.fpq = int(vision.get("fpq", fpq))
        self.mm_use_thumbnail = bool(vision.get("mm_use_thumbnail", mm_use_thumbnail))
        self.thumbnail_position = vision.get("thumbnail_position", thumbnail_position)
        self.emit_text_timestamps = bool(vision.get("emit_text_timestamps", emit_text_timestamps))
        self.disable_cubing = bool(vision.get("disable_cubing", disable_cubing))
        if self.fpq < 1 or self.thumbnail_position not in ("first", "last"):
            raise ValueError("fpq must be positive and thumbnail_position must be 'first' or 'last'")
        if int(vision.get("num_queries", NQ)) != NQ:
            raise ValueError("Quicksviewer requires exactly 64 visual queries")
        self.max_frames = int(max_frames)
        super().__init__(image_processor or QuicksviewerImageProcessor(), tokenizer,
                         video_processor or QuicksviewerVideoProcessor(max_frames=max_frames),
                         chat_template=chat_template or tokenizer.chat_template or _DEFAULT_TEMPLATE)
        for text, expected in (("<image>", [VISUAL_START_ID]), ("</image>", [VISUAL_END_ID]),
                               ("<unk>", [VISUAL_CONTENT_ID]), ("0000.0", [15, 15, 15, 15, 13, 15]),
                               ("0123456789", list(range(15, 25)))):
            if tokenizer.encode(text, add_special_tokens=False) != expected:
                raise ValueError(f"Tokenizer does not satisfy the Quicksviewer token contract for {text!r}")

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        # AutoProcessor dynamically imports this class; resolve components here so
        # custom image/video classes never need global Transformers registration.
        hub_keys = ("cache_dir", "force_download", "local_files_only", "revision", "token", "subfolder")
        hub = {key: kwargs[key] for key in hub_keys if key in kwargs}
        settings = _read_json(pretrained_model_name_or_path, "processor_config.json", hub)
        model = _read_json(pretrained_model_name_or_path, "config.json", hub)
        image_settings = _read_json(pretrained_model_name_or_path, "preprocessor_config.json", hub)
        video_settings = _read_json(pretrained_model_name_or_path, "video_preprocessor_config.json", hub)
        tokenizer = kwargs.pop("tokenizer", None)
        if tokenizer is None:
            tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path,
                                                      trust_remote_code=kwargs.get("trust_remote_code", True), **hub)
        template_path = cached_file(pretrained_model_name_or_path, "chat_template.jinja",
                                   _raise_exceptions_for_missing_entries=False, **hub)
        template = Path(template_path).read_text(encoding="utf-8") if template_path else None
        keys = ("fpq", "mm_use_thumbnail", "thumbnail_position", "emit_text_timestamps",
                "disable_cubing", "max_frames")
        options = {key: kwargs.get(key, settings.get(key)) for key in keys if key in kwargs or key in settings}
        if "max_frames" not in kwargs and "max_video_frames" in model:
            options["max_frames"] = model["max_video_frames"]
        options["vision_config"] = kwargs.get("vision_config", model.get("vision_config", {}))
        options["chat_template"] = kwargs.get("chat_template", template or settings.get("chat_template"))
        # Metadata is handled by this loader, not forwarded to preprocessing.
        for component in (image_settings, video_settings):
            for key in ("auto_map", "processor_class", "image_processor_type", "video_processor_type"):
                component.pop(key, None)
        return cls(image_processor=kwargs.get("image_processor") or QuicksviewerImageProcessor(**image_settings),
                   video_processor=kwargs.get("video_processor") or QuicksviewerVideoProcessor(**video_settings),
                   tokenizer=tokenizer, **options)

    def to_dict(self):
        return {
            "processor_class": type(self).__name__,
            "auto_map": {"AutoProcessor": "processing_quicksviewer.QuicksviewerProcessor"},
            **{key: getattr(self, key) for key in ("fpq", "mm_use_thumbnail", "thumbnail_position",
                                                  "emit_text_timestamps", "disable_cubing", "max_frames")},
        }

    def save_pretrained(self, save_directory, **kwargs):
        if kwargs.get("push_to_hub"):
            raise ValueError("Save locally, then upload the resulting repository with the Hugging Face Hub API")
        destination = Path(save_directory)
        destination.mkdir(parents=True, exist_ok=True)
        self.tokenizer.save_pretrained(destination)
        configs = {
            "processor_config.json": self.to_dict(),
            "preprocessor_config.json": {
                "image_processor_type": "QuicksviewerImageProcessor",
                "processor_class": type(self).__name__,
                "auto_map": {"AutoImageProcessor": "image_processing_quicksviewer.QuicksviewerImageProcessor"},
                "size": self.image_processor.size, "do_normalize": False, "do_rescale": False,
            },
            "video_preprocessor_config.json": {
                "video_processor_type": "QuicksviewerVideoProcessor",
                "processor_class": type(self).__name__,
                "auto_map": {"AutoVideoProcessor": "video_processing_quicksviewer.QuicksviewerVideoProcessor"},
                "size": self.video_processor.size, "max_frames": self.max_frames,
                "num_threads": self.video_processor.num_threads, "do_normalize": False, "do_rescale": False,
            },
        }
        for name, value in configs.items():
            (destination / name).write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")
        (destination / "chat_template.jinja").write_text(self.chat_template, encoding="utf-8")
        for name in ("processing_quicksviewer.py", "image_processing_quicksviewer.py", "video_processing_quicksviewer.py"):
            source = Path(__file__).with_name(name)
            if source.resolve() != (destination / name).resolve():
                shutil.copyfile(source, destination / name)
        return [str(destination / name) for name in configs]

    def __call__(self, images=None, text=None, videos=None, return_tensors="pt", **kwargs):
        if isinstance(text, (list, tuple)):
            if len(text) != 1 or not isinstance(text[0], str):
                raise ValueError("Quicksviewer supports B=1 only; text must be a string or a one-element list")
            text = text[0]
        if text is None:
            text = ""
        if not isinstance(text, str):
            raise TypeError("text must be a string")
        image_pool, video_pool = _image_list(images), video_list(videos)
        images_kwargs = dict(kwargs.pop("images_kwargs", {}))
        if images_kwargs:
            raise TypeError(f"Unsupported image options: {sorted(images_kwargs)}")
        video_kwargs = dict(kwargs.pop("videos_kwargs", {}))
        for key in ("max_frames", "fps", "timestamps", "duration", "video_metadata"):
            if key in kwargs:
                video_kwargs[key] = kwargs.pop(key)
        video_kwargs.setdefault("max_frames", self.max_frames)
        text_kwargs = dict(kwargs.pop("text_kwargs", {}))
        text_kwargs.update(kwargs)
        # Truncation could invalidate image spans and media alignment.
        if text_kwargs.pop("truncation", False):
            raise ValueError("Truncation is unsupported because it may cut visual/timestamp spans")
        max_length = text_kwargs.pop("max_length", None)
        text_kwargs.pop("padding", None)
        text_kwargs.pop("add_special_tokens", None)
        text_kwargs.pop("return_attention_mask", None)
        text_kwargs.pop("return_token_type_ids", None)
        if text_kwargs:
            raise TypeError(f"Unsupported text options: {sorted(text_kwargs)}")
        if not _MEDIA_PATTERN.search(text):
            # Raw evaluation prompts conventionally put media before the question.
            text = "<image>" * len(image_pool) + "<video>" * len(video_pool) + text
        ids, frames, counts, video_flags, temporal, image_spans, timestamp_spans = [], [], [], [], [], [], []
        used = {"image": 0, "video": 0}

        def block():
            ids.append(VISUAL_START_ID)
            start = len(ids)
            ids.extend([VISUAL_CONTENT_ID] * NQ)
            image_spans.append((start, len(ids)))
            ids.append(VISUAL_END_ID)

        def media(kind):
            pool = image_pool if kind == "image" else video_pool
            index = used[kind]
            if index >= len(pool):
                raise ValueError(f"More <{kind}> placeholders than supplied {kind} inputs")
            used[kind] += 1
            if kind == "image":
                pixels = self.image_processor.preprocess_one(pool[index]).unsqueeze(0)
                mid = [-1]
                cubes, thumbnail = 1, False
            else:
                options = dict(video_kwargs)
                metadata = options.get("video_metadata")
                if isinstance(metadata, (list, tuple)):
                    if len(metadata) != len(video_pool):
                        raise ValueError("video_metadata must have one entry per video")
                    options["video_metadata"] = metadata[index]
                pixels, ts, duration = self.video_processor.load_video(pool[index], **options)
                mid = frame_temporal_ids(ts, duration=duration or None)
                cubes = ((len(pixels) + self.fpq - 1) // self.fpq if self.disable_cubing
                         else num_video_cubes(len(pixels), self.fpq))
                thumbnail = self.mm_use_thumbnail and not self.disable_cubing
            frames.append(pixels)
            counts.append(len(pixels))
            video_flags.append(kind == "video")
            temporal.extend(mid)
            if thumbnail and self.thumbnail_position == "first":
                block()
            width = timestamp_token_count(max(mid, default=0))
            for _ in range(cubes):
                if kind == "video" and self.emit_text_timestamps:
                    start = len(ids)
                    ids.extend(timestamp_placeholder_ids(width))
                    timestamp_spans.append((start, len(ids)))
                block()
            if thumbnail and self.thumbnail_position == "last":
                block()

        cursor = 0
        for match in _MEDIA_PATTERN.finditer(text):
            ids.extend(self.tokenizer.encode(text[cursor:match.start()], add_special_tokens=False))
            media(match.group()[1:-1])
            cursor = match.end()
        ids.extend(self.tokenizer.encode(text[cursor:], add_special_tokens=False))
        if used != {"image": len(image_pool), "video": len(video_pool)}:
            raise ValueError("Every supplied image/video must have a matching placeholder")
        if max_length is not None and len(ids) > max_length:
            raise ValueError(f"Prompt length {len(ids)} exceeds max_length={max_length}; truncation is unsafe")
        data = {
            "input_ids": torch.tensor([ids], dtype=torch.long),
            "attention_mask": torch.ones((1, len(ids)), dtype=torch.long),
            "pixel_values": torch.cat(frames) if frames else torch.empty((0, 3, SIGLIP_SIZE, SIGLIP_SIZE), dtype=torch.uint8),
            "frames_per_media": torch.tensor(counts, dtype=torch.long),
            "media_is_video": torch.tensor(video_flags, dtype=torch.bool),
            "frame_temporal_ids": torch.tensor(temporal, dtype=torch.long),
            "image_spans": torch.tensor(image_spans, dtype=torch.long).reshape(-1, 2),
            "timestamp_spans": torch.tensor(timestamp_spans, dtype=torch.long).reshape(-1, 2),
        }
        if return_tensors not in (None, "pt", "np"):
            raise ValueError("return_tensors must be 'pt', 'np', or None")
        if return_tensors == "np":
            data = {key: value.numpy() for key, value in data.items()}
        elif return_tensors is None:
            data = {key: (value if key == "pixel_values" else value.tolist()) for key, value in data.items()}
        return BatchFeature(data)

    def apply_chat_template(self, conversation=None, chat_template=None, tokenize=True,
                            return_dict=False, return_tensors=None, add_generation_prompt=False,
                            enable_thinking=False, **kwargs):
        messages = conversation if conversation is not None else kwargs.pop("messages", None)
        if not isinstance(messages, (list, tuple)) or not messages or not isinstance(messages[0], dict):
            raise ValueError("Pass one nonempty conversation; batches (B>1) are unsupported")
        external_images = _image_list(kwargs.pop("images", None))
        external_videos = video_list(kwargs.pop("videos", None))
        consumed = {"image": 0, "video": 0}
        images, videos, metadata, clean = [], [], [], []
        for message in messages:
            content = message.get("content", "")
            if isinstance(content, str):
                clean.append(dict(message))
                continue
            parts = []
            for item in content:
                kind = item.get("type", "")
                if kind == "text":
                    parts.append(item.get("text", ""))
                    continue
                kind = {"image_url": "image", "video_url": "video"}.get(kind, kind)
                if kind not in ("image", "video"):
                    raise ValueError(f"Unsupported content type: {item.get('type')!r}")
                value = _media_value(item, kind)
                pool = external_images if kind == "image" else external_videos
                if value is None:
                    index = consumed[kind]
                    if index >= len(pool):
                        raise ValueError(f"Content entry requires a {kind} value or matching {kind}s= argument")
                    value = pool[index]
                    consumed[kind] += 1
                (images if kind == "image" else videos).append(value)
                if kind == "video":
                    metadata.append({key: item[key] for key in ("fps", "timestamps", "duration") if key in item})
                parts.append(f"<{kind}>")
            clean.append({**message, "content": "".join(parts)})
        # String-content conversations can use explicit markers with images/videos.
        for kind, pool, target in (("image", external_images, images), ("video", external_videos, videos)):
            remaining = pool[consumed[kind]:]
            if remaining:
                if target:
                    raise ValueError(f"Unused external {kind} inputs")
                target.extend(remaining)
        if (images or videos) and not any(_MEDIA_PATTERN.search(message["content"]) for message in clean):
            # External media with plain-string messages belongs inside the last
            # user turn, not before the system/chat boundary tokens.
            user_index = next((i for i in range(len(clean) - 1, -1, -1)
                               if clean[i].get("role") == "user"), None)
            if user_index is None:
                raise ValueError("External media requires a user message or explicit media markers")
            clean[user_index]["content"] = ("<image>" * len(images) + "<video>" * len(videos)
                                             + clean[user_index]["content"])
        template_kwargs = {key: kwargs.pop(key) for key in ("tools", "documents", "continue_final_message")
                           if key in kwargs}
        rendered = self.tokenizer.apply_chat_template(
            clean, chat_template=chat_template or self.chat_template, tokenize=False,
            add_generation_prompt=add_generation_prompt, enable_thinking=enable_thinking,
            **template_kwargs,
        )
        if not tokenize:
            return rendered
        if metadata and any(metadata) and "video_metadata" not in kwargs:
            kwargs["video_metadata"] = metadata
        result = self(images=images, videos=videos, text=rendered,
                      return_tensors=return_tensors, **kwargs)
        if return_dict:
            return result
        ids = result["input_ids"]
        return ids[0] if return_tensors is None else ids

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)

    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=True, **kwargs):
        return self.batch_decode(generated_outputs, skip_special_tokens=skip_special_tokens, **kwargs)


QuicksviewerProcessor.register_for_auto_class("AutoProcessor")