import logging
import os
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
from urllib.parse import urlparse
from typing import List, Optional, Union
from transformers import AutoConfig
from transformers import Qwen2_5_VLProcessor
from transformers import Qwen2_5_VLForConditionalGeneration
from qwen_vl_utils import process_vision_info

MAX_LENGTH = 2048
IMAGE_BASE_FACTOR = 14
IMAGE_FACTOR = IMAGE_BASE_FACTOR * 2
MIN_PIXELS = 4 * IMAGE_FACTOR * IMAGE_FACTOR
MAX_PIXELS = 1280 * IMAGE_FACTOR * IMAGE_FACTOR
FPS = 1
MAX_FRAMES = 64
FRAME_MAX_PIXELS = 768 * IMAGE_FACTOR * IMAGE_FACTOR
MAX_TOTAL_PIXELS = 10 * FRAME_MAX_PIXELS
EOS_TOKEN = "<|endoftext|>"

logger = logging.getLogger(__name__)


class VisME(nn.Module):
    def __init__(
        self,
        model_name: str,
        pooling: str = 'last',
        normalize: bool = True,
        max_length: int = MAX_LENGTH,
        min_pixels: int = MIN_PIXELS,
        max_pixels: int = MAX_PIXELS,
        total_pixels: int = MAX_TOTAL_PIXELS,
        fps: float = FPS,
        max_frames: int = MAX_FRAMES,
        processor=None,
        default_instruction: str = "You are a helpful assistant.",
        attn_implementation: Optional[str] = "flash_attention_2",
        **kwargs,
    ):
        """
        Initialize the Qwen2_5_VL embedding model.

        Args:
            model_name: Path to the base model (HuggingFace model ID or local path)
            device: Device to use
            max_length: Maximum sequence length
            attn_implementation: Attention implementation method
            pooling: Pooling strategy ('last' or 'eos')
            normalize: Whether to normalize embeddings
            checkpoint_path: Path to a trained model checkpoint (if provided, will be loaded with higher priority)
            use_custom_model: Whether to use the custom Qwen2_5_VLForConditionalGeneration
            system_prompt: System prompt string
            output_layer: Which layer to extract embeddings from:
                - -1 or None: Last hidden layer (default)
                - int (0 to num_layers-1): Specific hidden layer index
                - tuple (layer_idx, component): e.g., (-1, 'attn') or (-2, 'mlp')
            output_component: Which component of the layer to use (when output_layer is int):
                - 'full': Complete layer output (default, after attn + mlp)
                - 'attn': Self-attention output only
                - 'mlp': MLP output only
            processor: Optional custom processor instance. If None, will load default processor from model_name
            **kwargs: Additional parameters
        """
        super().__init__()
        
        self.pooling = pooling
        self.normalize = normalize
        self.default_instruction = default_instruction
        self.max_length = max_length
        self.min_pixels = min_pixels
        self.max_pixels = max_pixels
        self.total_pixels = total_pixels
        self.fps = fps
        self.max_frames = max_frames
        self.eos_token = kwargs.get("eos_token", EOS_TOKEN)
        
        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
        config._attn_implementation = attn_implementation
        config.use_cache = False
        
        base_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            model_name,
            config=config,
            dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
        )
        self.model = base_model.model
        
        if processor is not None:
            self.processor = processor
        else:
            self.processor = Qwen2_5_VLProcessor.from_pretrained(
                model_name, 
                padding_side='left',
                trust_remote_code=True
            )

    @property
    def device(self):
        return next(self.model.parameters()).device

    @staticmethod
    def is_image_path(path: str) -> bool:
        image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.tiff', '.svg'}
        if path.startswith(('http://', 'https://')):
            parsed_url = urlparse(path)
            clean_path = parsed_url.path
        else:
            clean_path = path
        _, ext = os.path.splitext(clean_path.lower())
        return ext in image_extensions

    @staticmethod
    def is_video_input(video) -> bool:
        if isinstance(video, str):
            return True
        if isinstance(video, list) and len(video) > 0:
            first_elem = video[0]
            if isinstance(first_elem, Image.Image):
                return True
            if isinstance(first_elem, str):
                return VisME.is_image_path(first_elem)
        return False

    def _load_images(self, images):
        if images is None:
            return None
        if not (isinstance(images, list) or isinstance(images, tuple)):
            images = [images]
        if len(images) == 0:
            return images
        if images[0] is None:
            return images
        if isinstance(images[0], (Image.Image, np.ndarray, torch.Tensor)):
            return images
        if isinstance(images[0], str):
            loaded_images = []
            for image in images:
                if image is None:
                    loaded_images.append(None)
                    continue
                if image.startswith(('http://', 'https://', 'file://')):
                    loaded_images.append(image)
                else:
                    loaded_images.append('file://' + image)
            return loaded_images
        raise ValueError("Unsupported image type")

    def _batch_to_device(self, batch):
        _batch = {}
        for key, value in batch.items():
            if isinstance(value, torch.Tensor):
                _batch[key] = value.to(self.device)
            else:
                _batch[key] = value
        return _batch
    
    def _load_input(self, input):
        """
        Args:
            input: dict, {'text': list of text, 'image': list of image file path or PIL images, 'instruction': list of instruction}
        Returns:
            texts: list of text
            images: list of image file path or PIL images
            instructions: list of instruction
        """
        images = input.get('image', [])
        instructions = input.get('instruction', [])
        texts = input.get('text', [])
        if not isinstance(images, list):
            images = [images]
        if not isinstance(instructions, list):
            if instructions is None:
                instructions = ""
            instructions = [instructions]
        if not isinstance(texts, list):
            if texts is None:
                texts = ""
            texts = [texts]

        assert len(texts) > 0 or len(images) > 0, "At least one of text or image must be provided!"

        if len(images) == 0:
            images = [None] * len(texts)
        if len(texts) == 0:
            texts = [""] * len(images)
        if len(instructions) == 0:
            instructions = [""] * len(images)
        if len(instructions) == 1:
            instructions = instructions * len(images)

        return images, instructions, texts

    def format_input(
        self,
        text: list[str] | str | None = None,
        image: list[str | Image.Image] | str | Image.Image | None = None,
        video: list[str | list[str | Image.Image]] | list[str | Image.Image] | None = None,
        instruction: str | None = None,
        fps: float | None = None,
        max_frames: int | None = None,
    ) -> list[dict]:
        content = []
        conversation = [
            {"role": "system", "content": [{"type": "text", "text": instruction or self.default_instruction}]},
            {"role": "user", "content": content}
        ]

        # Normalize text input to list
        if text is None:
            texts = []
        elif isinstance(text, str):
            texts = [text]
        else:
            texts = text
        
        # Normalize image input to list
        if image is None:
            images = []
        elif not isinstance(image, list):
            images = [image]
        else:
            images = image
        
        # Normalize video input to list
        if video is None:
            videos = []
        elif self.is_video_input(video):
            videos = [video]
        else:
            # Assume it's a list of videos
            videos = video

        # Add text, image, or video content to conversation
        if not texts and not images and not videos:
            logger.warning("No text, image, or video content found in input")
            content.append({'type': 'text', 'text': "NULL"})
            return conversation
        
        # Process each video
        for vid in videos:
            video_content = None
            video_kwargs = {'total_pixels': self.total_pixels}
            
            if isinstance(vid, list):
                # Video as frame sequence
                video_content = vid
                if self.max_frames is not None:
                    video_content = _sample_frames(video_content, self.max_frames)
                video_content = [
                    ('file://' + ele if isinstance(ele, str) else ele) 
                    for ele in video_content
                ]
            elif isinstance(vid, str):
                # Video as file path
                video_content = vid if vid.startswith(('http://', 'https://')) else 'file://' + vid
                video_kwargs = {'fps': fps or self.fps, 'max_frames': max_frames or self.max_frames}
            else:
                raise TypeError(f"Unrecognized video type: {type(vid)}")

            # Add video input to content
            if video_content:
                content.append({
                    'type': 'video', 
                    'video': video_content,
                    **video_kwargs
                })

        # Process each image
        for img in images:
            image_content = None
            
            if isinstance(img, Image.Image):
                image_content = img
            elif isinstance(img, str):
                image_content = img if img.startswith(('http://', 'https://')) else 'file://' + img
            else:
                raise TypeError(f"Unrecognized image type: {type(img)}")

            # Add image input to content
            if image_content:
                content.append({
                    'type': 'image', 
                    'image': image_content,
                    "min_pixels": self.min_pixels,
                    "max_pixels": self.max_pixels
                })

        # Process each text
        for txt in texts:
            content.append({'type': 'text', 'text': txt})

        return conversation
    
    def preprocess_input(
        self,
        inputs: Union[dict, List[dict]],
        **kwargs
    ):
        if isinstance(inputs, dict):
            inputs = [inputs]
        
        fps = kwargs.get('fps', self.fps)
        max_frames = kwargs.get('max_frames', self.max_frames)
        conversations = [self.format_input(
            text=ele.get('text'),
            image=ele.get('image'),
            video=ele.get('video'),
            instruction=ele.get('instruction'),
            fps=fps,
            max_frames=max_frames
        ) for ele in inputs]

        text = self.processor.apply_chat_template(conversations, tokenize=False, add_generation_prompt=True)
        text = [t+self.eos_token for t in text]
        try:
            images, video_inputs, video_kwargs = process_vision_info(
                conversations, image_patch_size=IMAGE_BASE_FACTOR,
                return_video_metadata=True, return_video_kwargs=True
            )
        except Exception as e:
            logger.error(f"Error in processing vision info: {e}")
            images = None
            video_inputs = None
            video_kwargs = {'do_sample_frames': False}
            text = self.processor.apply_chat_template(
                [{'role': 'user', 'content': [{'type': 'text', 'text': 'NULL'}]}], 
                add_generation_prompt=True, tokenize=False
            )

        if video_inputs is not None:
            videos, video_metadata = zip(*video_inputs)
            videos = list(videos)
            video_metadata = list(video_metadata)
        else:
            videos, video_metadata = None, None

        inputs = self.processor(
            text=text, images=images, videos=videos, video_metadata=video_metadata, truncation=True, 
            max_length=self.max_length, padding=True, return_tensors='pt',
            **video_kwargs
        )
        return inputs

    def _pooling(self, embeddings_source: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Pool sequence hidden states using ``self.pooling`` (last / eos / mean)."""
        pooling = getattr(self, "pooling", "last")
        if pooling in ("last", "eos"):
            left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
            batch_size = embeddings_source.shape[0]
            if left_padding:
                reps = embeddings_source[torch.arange(batch_size), -1, :]
            else:
                eos_indices = attention_mask.sum(dim=1) - 1
                reps = embeddings_source[
                    torch.arange(batch_size, device=embeddings_source.device), eos_indices
                ]
        elif pooling == "mean":
            mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings_source.size()).float()
            sum_embeddings = torch.sum(embeddings_source * mask_expanded, dim=1)
            sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
            reps = sum_embeddings / sum_mask
        else:
            raise NotImplementedError(f"Pooling method '{pooling}' not implemented")
        return reps

    def _encode(self, input):
        """
        Encode inputs and extract embeddings from specified layer.
        
        Args:
            input: Processed inputs (input_ids, attention_mask, etc.)
            
        Returns:
            Tensor of shape (batch_size, hidden_dim)
        """
        outputs = self.model(
            **input, 
            output_hidden_states=False,
            return_dict=True,
            use_cache=False,
        )
        embeddings_source = outputs.last_hidden_state
        reps = self._pooling(embeddings_source, input['attention_mask'])
        if self.normalize:
            reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
        return reps
        
    def encode_text(
        self,
        texts: List[str],
        instruction: Optional[Union[str, List[str]]] = None,
        **kwargs,
    ) -> torch.Tensor:
        """Convenience method for text-only embeddings"""
        inputs = {"text": texts, "instruction": instruction}
        return self.encode_input(inputs, instruction=instruction, **kwargs)

    def encode_image(
        self,
        images,
        instruction: Optional[Union[str, List[str]]] = None,
        **kwargs,
    ) -> torch.Tensor:
        """Convenience method for image-only embeddings.

        Args:
            images: Can be:
                - List[Image.Image]: Single image per input
                - List[List[Image.Image]]: Multiple images per input
        """
        inputs = {"image": images, "instruction": instruction}
        return self.encode_input(inputs, instruction=instruction, **kwargs)

    def encode_input(
        self,
        inputs: dict,
        **kwargs,
    ) -> torch.Tensor:
        """Batch processing for large collections of texts/images.

        Args:
            texts: List of text inputs (optional)
            images: Can be:
                - List[Image.Image]: Single image per input
                - List[List[Image.Image]]: Multiple images per input
            instruction: Instruction(s) for the model
            batch_size: Number of items to process at once
            show_progress: Whether to display progress bar
        """
        if 'input_ids' in inputs:
            assert isinstance(inputs['input_ids'], torch.Tensor), "input_ids must be a tensor"
            assert isinstance(inputs['attention_mask'], torch.Tensor), "attention_mask must be a tensor"
            inputs = self._batch_to_device(inputs)
        else:
            inputs = self.preprocess_input(inputs, **kwargs)
            inputs = self._batch_to_device(inputs)

        embeddings = self._encode(inputs)
        return embeddings


def _sample_frames(frames: list[str | Image.Image], max_segments: int) -> list[str | Image.Image]:
    duration = len(frames)
    if duration <= max_segments:
        return frames

    frame_id_array = np.linspace(0, duration - 1, max_segments, dtype=int)
    frame_id_list = frame_id_array.tolist()
    sampled_frames = [ frames[frame_idx] for frame_idx in frame_id_list ]
    return sampled_frames