import base64
import re

import numpy as np
from google.genai import types

from mm_agents.vlaa_gui.core.engine import (
    LMMEngineAnthropic,
    LMMEngineAnthropicBedrock,
    LMMEngineArk,
    LMMEngineAzureOpenAI,
    LMMEngineGemini,
    LMMEngineQwen,
    LMMEngineHuggingFace,
    LMMEngineOpenAI,
    LMMEngineOpenRouter,
    LMMEngineParasail,
    LMMEnginevLLM,
)

ENGINE_REGISTRY = {
    "openai": LMMEngineOpenAI,
    "anthropic": LMMEngineAnthropic,
    "anthropic_bedrock": LMMEngineAnthropicBedrock,
    "azure": LMMEngineAzureOpenAI,
    "vllm": LMMEnginevLLM,
    "huggingface": LMMEngineHuggingFace,
    "gemini": LMMEngineGemini,
    "open_router": LMMEngineOpenRouter,
    "parasail": LMMEngineParasail,
    "ark": LMMEngineArk,
    "qwen": LMMEngineQwen,
}


class LMMAgent:
    def __init__(self, engine_params=None, system_prompt=None, engine=None):
        if engine is None:
            if engine_params is not None:
                self.engine = self._build_engine(engine_params)
            else:
                raise ValueError("engine_params must be provided")
        else:
            self.engine = engine

        self.messages = []  # Empty messages

        if system_prompt:
            self.add_system_prompt(system_prompt)
        else:
            self.add_system_prompt("You are a helpful assistant.")

    @staticmethod
    def _build_engine(engine_params):
        engine_type = engine_params.get("engine_type")
        engine_class = ENGINE_REGISTRY.get(engine_type)
        if engine_class is None:
            raise ValueError("engine_type is not supported")
        return engine_class(**engine_params)

    @staticmethod
    def _split_thinking_response(response):
        thoughts_match = re.search(r"<thoughts>(.*?)</thoughts>", response, re.DOTALL)
        answer_match = re.search(r"<answer>(.*?)</answer>", response, re.DOTALL)

        if not thoughts_match and not answer_match:
            return response, None

        answer = answer_match.group(1).strip() if answer_match else response
        thoughts = thoughts_match.group(1).strip() if thoughts_match else None
        return answer, thoughts

    def _normalize_response(self, response):
        if isinstance(response, tuple):
            answer = response[0]
            thoughts = response[1] if len(response) > 1 else None
            return answer, thoughts
        if isinstance(response, str):
            return self._split_thinking_response(response)
        return response, None

    def encode_image(self, image_content):
        # if image_content is a path to an image file, check type of the image_content to verify
        if isinstance(image_content, str):
            with open(image_content, "rb") as image_file:
                return base64.b64encode(image_file.read()).decode("utf-8")
        else:
            return base64.b64encode(image_content).decode("utf-8")

    def _has_image_content(self, image_content):
        if image_content is None:
            return False
        if isinstance(image_content, list):
            return len(image_content) > 0
        return True

    def _get_message_role(self, message):
        role = getattr(message, "role", None)
        if role is None and isinstance(message, dict):
            role = message.get("role")
        return role

    def _build_gemini_parts(
        self,
        text_content,
        image_content=None,
        image_detail="high",
        put_text_last=True,
    ):
        parts = [types.Part.from_text(text=text_content)]
        img_detail_mappings = {
            "low": types.MediaResolution.MEDIA_RESOLUTION_LOW,
            "medium": types.MediaResolution.MEDIA_RESOLUTION_MEDIUM,
            "high": types.MediaResolution.MEDIA_RESOLUTION_HIGH,
            "original": types.MediaResolution.MEDIA_RESOLUTION_HIGH,
        }

        if self._has_image_content(image_content):
            images = (
                image_content if isinstance(image_content, list) else [image_content]
            )
            for image in images:
                b64_data = self.encode_image(image)
                parts.append(
                    types.Part.from_bytes(
                        data=base64.b64decode(b64_data),
                        mime_type="image/png",
                        media_resolution=img_detail_mappings.get(
                            image_detail, types.MediaResolution.MEDIA_RESOLUTION_HIGH
                        ),
                    )
                )

        if put_text_last and len(parts) > 1:
            parts.append(parts.pop(0))

        return parts

    def reset(
        self,
    ):

        self.messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": self.system_prompt}],
            }
        ]

    def add_system_prompt(self, system_prompt):
        self.system_prompt = system_prompt
        if len(self.messages) > 0:
            self.messages[0] = {
                "role": "system",
                "content": [{"type": "text", "text": self.system_prompt}],
            }
        else:
            self.messages.append(
                {
                    "role": "system",
                    "content": [{"type": "text", "text": self.system_prompt}],
                }
            )

    def remove_message_at(self, index):
        """Remove a message at a given index"""
        if index < len(self.messages):
            self.messages.pop(index)

    def replace_message_at(
        self, index, text_content, image_content=None, image_detail="high"
    ):
        """Replace a message at a given index"""
        if index < len(self.messages):
            self.messages[index] = {
                "role": self.messages[index]["role"],
                "content": [{"type": "text", "text": text_content}],
            }
            if image_content:
                base64_image = self.encode_image(image_content)
                self.messages[index]["content"].append(
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{base64_image}",
                            "detail": image_detail,
                        },
                    }
                )

    def add_message(
        self,
        text_content,
        image_content=None,
        role=None,
        image_detail="high",
        put_text_last=False,
    ):
        """Add a new message to the list of messages"""

        # API-style inference from OpenAI and AzureOpenAI
        if isinstance(
            self.engine,
            (
                LMMEngineArk,
                LMMEngineOpenAI,
                LMMEngineAzureOpenAI,
                LMMEngineHuggingFace,
                LMMEngineGemini,
                LMMEngineQwen,
                LMMEngineOpenRouter,
                LMMEngineParasail,
            ),
        ):
            # infer role from previous message
            if role != "user":
                if self.messages[-1]["role"] == "system":
                    role = "user"
                elif self.messages[-1]["role"] == "user":
                    role = "assistant"
                elif self.messages[-1]["role"] == "assistant":
                    role = "user"

            message = {
                "role": role,
                "content": [{"type": "text", "text": text_content}],
            }

            if isinstance(image_content, np.ndarray) or image_content:
                # Check if image_content is a list or a single image
                if isinstance(image_content, list):
                    # If image_content is a list of images, loop through each image
                    for image in image_content:
                        base64_image = self.encode_image(image)
                        message["content"].append(
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{base64_image}",
                                    "detail": image_detail,
                                },
                            }
                        )
                else:
                    # If image_content is a single image, handle it directly
                    base64_image = self.encode_image(image_content)
                    message["content"].append(
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}",
                                "detail": image_detail,
                            },
                        }
                    )

            # Rotate text to be the last message if desired
            if put_text_last:
                text_content = message["content"].pop(0)
                message["content"].append(text_content)

            self.messages.append(message)

        # For API-style inference from Anthropic
        elif isinstance(
            self.engine,
            (LMMEngineAnthropic, LMMEngineAnthropicBedrock),
        ):
            # infer role from previous message
            if role != "user":
                if self.messages[-1]["role"] == "system":
                    role = "user"
                elif self.messages[-1]["role"] == "user":
                    role = "assistant"
                elif self.messages[-1]["role"] == "assistant":
                    role = "user"

            message = {
                "role": role,
                "content": [{"type": "text", "text": text_content}],
            }

            if image_content:
                # Check if image_content is a list or a single image
                if isinstance(image_content, list):
                    # If image_content is a list of images, loop through each image
                    for image in image_content:
                        base64_image = self.encode_image(image)
                        message["content"].append(
                            {
                                "type": "image",
                                "source": {
                                    "type": "base64",
                                    "media_type": "image/png",
                                    "data": base64_image,
                                },
                            }
                        )
                else:
                    # If image_content is a single image, handle it directly
                    base64_image = self.encode_image(image_content)
                    message["content"].append(
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": "image/png",
                                "data": base64_image,
                            },
                        }
                    )

            # Rotate text to be the last content block if desired
            if put_text_last:
                text_content_block = message["content"].pop(0)
                message["content"].append(text_content_block)

            self.messages.append(message)

        # Locally hosted vLLM model inference
        elif isinstance(self.engine, LMMEnginevLLM):
            # infer role from previous message
            if role != "user":
                if self.messages[-1]["role"] == "system":
                    role = "user"
                elif self.messages[-1]["role"] == "user":
                    role = "assistant"
                elif self.messages[-1]["role"] == "assistant":
                    role = "user"

            message = {
                "role": role,
                "content": [{"type": "text", "text": text_content}],
            }

            if image_content:
                # Check if image_content is a list or a single image
                if isinstance(image_content, list):
                    # If image_content is a list of images, loop through each image
                    for image in image_content:
                        base64_image = self.encode_image(image)
                        message["content"].append(
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image;base64,{base64_image}"
                                },
                            }
                        )
                else:
                    # If image_content is a single image, handle it directly
                    base64_image = self.encode_image(image_content)
                    message["content"].append(
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image;base64,{base64_image}"},
                        }
                    )

            self.messages.append(message)

        elif isinstance(self.engine, LMMEngineGemini):
            # infer role from previous message
            if role != "user":
                if len(self.messages) == 0:
                    role = "user"
                else:
                    last_message = self.messages[-1]
                    last_role = getattr(last_message, "role", None)
                    if last_role is None and isinstance(last_message, dict):
                        last_role = last_message.get("role")

                    if last_role == "user":
                        role = "model"
                    elif last_role in {"model", "assistant"}:
                        role = "user"
                    else:
                        role = "user"

            parts = self._build_gemini_parts(
                text_content,
                image_content=image_content,
                image_detail=image_detail,
                put_text_last=put_text_last,
            )

            # Append to messages in Google's Content format
            self.messages.append(types.Content(role=role, parts=parts))
        else:
            raise ValueError("engine_type is not supported")

    def get_response(
        self,
        user_message=None,
        messages=None,
        temperature=0.0,
        max_new_tokens=None,
        **kwargs,
    ):
        """Generate the next response based on previous messages"""
        if messages is None:
            messages = self.messages
        if user_message:
            messages.append(
                {"role": "user", "content": [{"type": "text", "text": user_message}]}
            )

        response = self.engine.generate(
            messages,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            **kwargs,
        )
        answer, _ = self._normalize_response(response)
        return answer

    def get_response_with_thinking(
        self,
        user_message=None,
        messages=None,
        temperature=0.0,
        max_new_tokens=None,
        **kwargs,
    ):
        if messages is None:
            messages = self.messages
        if user_message:
            messages.append(
                {"role": "user", "content": [{"type": "text", "text": user_message}]}
            )
        thinking_enabled = getattr(self.engine, "thinking", False)
        generate_with_thinking = getattr(self.engine, "generate_with_thinking", None)
        if thinking_enabled and callable(generate_with_thinking):
            response = generate_with_thinking(
                messages,
                temperature=temperature,
                max_new_tokens=max_new_tokens,
                **kwargs,
            )
            return self._normalize_response(response)
