import base64
from io import BytesIO
from typing import Any

import numpy as np
from PIL import Image
from sinapsis_chatbots_base.helpers.llm_keys import LLMChatKeys
from sinapsis_chatbots_base.helpers.tags import Tags
from sinapsis_core.data_containers.data_packet import DataContainer, TextPacket
from sinapsis_core.template_base.base_models import OutputTypes, UIPropertiesMetadata

from sinapsis_vllm.helpers.schemas import vLLMMultimodalInitArgs
from sinapsis_vllm.helpers.vllm_keys import vLLMMultiModalKeys
from sinapsis_vllm.templates.vllm_text_completion import vLLMTextCompletion, vLLMTextCompletionAttributes


class vLLMMultiModal(vLLMTextCompletion):
    """Template for multimodal (text + image) completion using vLLM.

    Usage example:

    agent:
      name: my_multimodal_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: vLLMMultiModal
      class_name: vLLMMultiModal
      template_input: InputTemplate
      attributes:
        init_args:
          llm_model_name: 'Qwen/Qwen2.5-VL-7B-Instruct'
          max_model_len: 4096
          dtype: auto
          seed: 0
          gpu_memory_utilization: 0.9
          max_num_seqs: 4
          disable_log_stats: true
        completion_args:
          temperature: 0.2
          top_p: 0.8
          top_k: 20
          min_p: 0
          max_tokens: 1024
        system_prompt: You are a helpful vision-language assistant.
    """

    UIProperties = UIPropertiesMetadata(
        category="vLLM",
        output_type=OutputTypes.TEXT,
        tags=[Tags.CHATBOTS, Tags.LLM, Tags.TEXT_COMPLETION, Tags.TEXT, Tags.MULTIMODAL, Tags.IMAGE_TO_TEXT],
    )

    class AttributesBaseModel(vLLMTextCompletionAttributes):
        """Attributes for vLLM multimodal template.

        Overrides ``init_args`` to use multimodal-specific defaults such as
        ``trust_remote_code=True`` and ``limit_mm_per_prompt={"image": 1}``.

        Attributes:
            init_args (vLLMMultimodalInitArgs): vLLM multimodal engine arguments.
        """

        init_args: vLLMMultimodalInitArgs

    @staticmethod
    def _image_array_to_data_uri(image_array: np.ndarray, fmt: str = "JPEG") -> str:
        """Converts a numpy image array to a base64-encoded data URI.

        Args:
            image_array (np.ndarray): The image as a numpy array (e.g., from an ImagePacket).
            fmt (str): Image format for encoding. Defaults to ``"JPEG"``.

        Returns:
            str: A data URI string in the format ``data:image/<fmt>;base64,...``.
        """
        img = Image.fromarray(image_array)
        buffer = BytesIO()
        img.save(buffer, format=fmt)
        b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
        mime = f"image/{fmt.lower()}"
        return f"data:{mime};base64,{b64}"

    def extract_image_content(self, container: DataContainer) -> list[dict]:
        """Extracts images from the container and formats them as OpenAI-compatible content parts.

        Each image is converted from a numpy array to a base64-encoded data URI
        and wrapped in the ``image_url`` content part format expected by vLLM.

        Args:
            container (DataContainer): The container holding image packets.

        Returns:
            list[dict]: A list of image content part dictionaries.
        """
        content = []
        for image_packet in container.images:
            data_uri = self._image_array_to_data_uri(image_packet.content)
            content.append(
                {
                    vLLMMultiModalKeys.type: vLLMMultiModalKeys.image_url,
                    vLLMMultiModalKeys.image_url: {vLLMMultiModalKeys.url: data_uri},
                }
            )
        return content

    def generate_response(self, container: DataContainer) -> DataContainer:
        """Processes text packets with images, generating multimodal responses.

        Builds OpenAI-compatible multipart content messages that include both
        images (from ``container.images``) and text. Images are shared across
        all packets since they come from the same container.

        Args:
            container (DataContainer): Container with incoming messages and images.

        Returns:
            DataContainer: Updated DataContainer with the LLM responses.
        """
        self.logger.debug("Multimodal chatbot in progress")
        responses: list[TextPacket] = []
        image_content = self.extract_image_content(container)

        for packet in container.texts:
            full_context: list[dict] = []
            user_content: list[dict] = []
            user_id, session_id, prompt = self.prepare_conversation_context(packet)

            if self.system_prompt:
                system_content = [
                    {vLLMMultiModalKeys.type: vLLMMultiModalKeys.text, vLLMMultiModalKeys.text: self.system_prompt}
                ]
                full_context.append(self.generate_dict_msg(LLMChatKeys.system_value, system_content))

            if self.attributes.chat_history_key:
                full_context.extend(packet.generic_data.get(self.attributes.chat_history_key, []))

            user_content.extend(image_content)
            user_content.append({vLLMMultiModalKeys.type: vLLMMultiModalKeys.text, vLLMMultiModalKeys.text: prompt})
            full_context.append(self.generate_dict_msg(LLMChatKeys.user_value, user_content))

            response = self.infer(full_context)
            self.logger.debug("End of interaction.")

            if response:
                generic_data: dict[str, Any] = {}
                if self.attributes.completion_args.response_format.type_ == "json_object":
                    parsed = self._parse_json_response(response)
                    if parsed is not None:
                        generic_data[self.attributes.structured_output_key] = parsed

                responses.append(TextPacket(source=session_id, content=response, id=user_id, generic_data=generic_data))

        container.texts.extend(responses)
        return container
