"""
Self-contained processor shim for trust_remote_code.

This processor:
- loads frames from a video path (OpenCV)
- applies InternVL2.5 `<img>...<IMG_CONTEXT>*N...</img>` token template
- returns `vision_input`, `input_ids`, `attention_mask`, and `question_input_ids`
  compatible with `QTSplusInternLM2_ForCausalLM.generate`.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
from PIL import Image

from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput


def _uniform_indices(num_frames: int, vlen: int) -> List[int]:
    num_frames = max(int(num_frames), 1)
    if vlen <= 0:
        return []
    if num_frames == 1:
        return [max(0, (vlen - 1) // 2)]
    last = vlen - 1
    return [int(round(i * last / (num_frames - 1))) for i in range(num_frames)]


def _load_video_frames_cv2(path: str, num_frames: int = 8) -> List[Image.Image]:
    import cv2

    cap = cv2.VideoCapture(path)
    if not cap.isOpened():
        raise FileNotFoundError(f"Failed to open video: {path}")

    vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    if vlen <= 0:
        # Fallback: decode sequentially and take the first `num_frames`.
        frames: List[Image.Image] = []
        while len(frames) < num_frames:
            ok, frame = cap.read()
            if not ok:
                break
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(Image.fromarray(frame))
        cap.release()
        return frames

    indices = _uniform_indices(num_frames, vlen)
    frames = []
    for idx in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ok, frame = cap.read()
        if not ok:
            continue
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(Image.fromarray(frame))
    cap.release()
    return frames


@dataclass
class QTSplusInternVL2_5_ProcessorKwargs:
    num_frames: int = 8
    system_prompt: str = "You are a helpful assistant."


class QTSplusInternVL2_5_Processor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(self, image_processor=None, tokenizer=None, **kwargs):
        super().__init__(image_processor=image_processor, tokenizer=tokenizer)
        self.img_start_token = "<img>"
        self.img_end_token = "</img>"
        self.img_context_token = "<IMG_CONTEXT>"

        # InternVL2.5 default: (448/14)^2 * (0.5^2) = 256 tokens per image.
        self.num_image_token = int(kwargs.pop("num_image_token", 256))
        self.system_prompt = str(kwargs.pop("system_prompt", "You are a helpful assistant."))

    def _build_image_tokens(self, num_images: int) -> str:
        num_images = max(int(num_images), 1)
        one = self.img_start_token + (self.img_context_token * int(self.num_image_token)) + self.img_end_token
        return one * num_images

    def __call__(
        self,
        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
        images=None,
        videos: Optional[Union[str, List[Image.Image]]] = None,
        return_tensors: Optional[str] = "pt",
        num_frames: int = 8,
        system_prompt: Optional[str] = None,
        **kwargs,
    ) -> BatchFeature:
        if text is None:
            raise ValueError("`text` is required")
        if isinstance(text, list):
            if len(text) != 1:
                raise ValueError("Only single-example processing is supported for now")
            text = text[0]

        if videos is not None and images is not None:
            raise ValueError("Pass only one of `videos` or `images`")

        if videos is not None:
            if isinstance(videos, str):
                frames = _load_video_frames_cv2(videos, num_frames=num_frames)
            elif isinstance(videos, list):
                frames = videos
            else:
                raise ValueError(f"Unsupported `videos` type: {type(videos)}")
            images = frames

        if images is None:
            raise ValueError("Either `videos` or `images` must be provided")
        if isinstance(images, Image.Image):
            images = [images]

        if not isinstance(images, list) or not images:
            raise ValueError("No frames/images loaded")

        img_tokens = self._build_image_tokens(num_images=len(images))
        user_content = f"{img_tokens}\n{text}"
        messages = [
            {"role": "system", "content": system_prompt or self.system_prompt},
            {"role": "user", "content": user_content},
        ]

        if not hasattr(self.tokenizer, "apply_chat_template"):
            raise ValueError("Tokenizer does not support apply_chat_template; missing chat_template.jinja?")

        prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

        text_inputs = self.tokenizer(
            prompt,
            add_special_tokens=False,
            return_tensors=return_tensors,
        )
        question_inputs = self.tokenizer(
            str(text),
            add_special_tokens=False,
            return_tensors=return_tensors,
        )

        vision_inputs = self.image_processor(images=images, return_tensors=return_tensors)
        pixel_values = vision_inputs["pixel_values"]

        return BatchFeature(
            data={
                "input_ids": text_inputs["input_ids"],
                "attention_mask": text_inputs.get("attention_mask"),
                "question_input_ids": question_inputs["input_ids"],
                "vision_input": pixel_values,
            },
            tensor_type=return_tensors,
        )