import re
from banks import Prompt
from banks.types import ContentBlockType as BanksContentBlockType
from copy import deepcopy
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union

from llama_index.core.bridge.pydantic import Field
from llama_index.core.base.llms.base import BaseLLM
from llama_index.core.base.llms.generic_utils import (
    messages_to_prompt as default_messages_to_prompt,
)
from llama_index.core.base.llms.types import ContentBlock, TextBlock, ImageBlock
from llama_index.core.llms import ChatMessage
from llama_index.core.prompts.base import BasePromptTemplate, PromptType
from llama_index.core.prompts.utils import get_template_vars, format_string
from llama_index.core.types import BaseOutputParser


def get_jinja_template_vars(template_str: str) -> list[str]:
    """Gets the template variables from a Jinja template string. Accounts for {var | fn} syntax."""
    vars = re.findall(r"\{\{(.*?)\}\}", template_str)
    return [var.split("|")[0].strip() for var in vars]


class RichPromptTemplate(BasePromptTemplate):
    template_str: str = Field(description="The template string for the prompt.")

    def __init__(
        self,
        template_str: str,
        metadata: Optional[Dict[str, Any]] = None,
        output_parser: Optional[BaseOutputParser] = None,
        template_vars: Optional[List[str]] = None,
        template_var_mappings: Optional[Dict[str, Any]] = None,
        function_mappings: Optional[Dict[str, Callable]] = None,
        **kwargs: Any,
    ):
        template_vars = template_vars or []
        if not template_vars:
            template_vars = get_jinja_template_vars(template_str)

        super().__init__(
            template_str=template_str,
            kwargs=kwargs,
            metadata=metadata or {},
            output_parser=output_parser,
            template_vars=template_vars,
            template_var_mappings=template_var_mappings,
            function_mappings=function_mappings,
        )

    @classmethod
    def from_messages(
        cls,
        message_templates: Union[List[Tuple[str, str]], List[ChatMessage]],
        **kwargs: Any,
    ) -> "RichPromptTemplate":
        """From messages."""
        raise NotImplementedError("RichPromptTemplate does not support from_messages")

    def partial_format(self, **kwargs: Any) -> "RichPromptTemplate":
        prompt = deepcopy(self)
        prompt.kwargs.update(kwargs)
        return prompt

    def format(
        self,
        llm: Optional[BaseLLM] = None,
        messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
        **kwargs: Any,
    ) -> str:
        del llm  # unused
        messages = self.format_messages(**kwargs)

        if messages_to_prompt is not None:
            return messages_to_prompt(messages)

        return default_messages_to_prompt(messages)

    def format_messages(
        self, llm: Optional[BaseLLM] = None, **kwargs: Any
    ) -> List[ChatMessage]:
        del llm  # unused
        """Format the prompt into a list of chat messages."""
        all_kwargs = {
            **self.kwargs,
            **kwargs,
        }
        mapped_all_kwargs = self._map_all_vars(all_kwargs)

        banks_prompt = Prompt(self.template_str)
        banks_messages = banks_prompt.chat_messages(mapped_all_kwargs)

        llama_messages: list[ChatMessage] = []
        for bank_message in banks_messages:
            if isinstance(bank_message.content, str):
                llama_messages.append(ChatMessage(role=bank_message.role, content=bank_message.content))
            elif isinstance(bank_message.content, list):
                llama_blocks: list[ContentBlock] = []
                for bank_block in bank_message.content:
                    if bank_block.type == BanksContentBlockType.text:
                        llama_blocks.append(TextBlock(text=bank_block.text))
                    elif bank_block.type == BanksContentBlockType.image_url:
                        llama_blocks.append(ImageBlock(url=bank_block.image_url.url))
                    else:
                        raise ValueError(f"Unsupported content block type: {bank_block.type}")
                
                llama_messages.append(ChatMessage(role=bank_message.role, content=llama_blocks))
            else:
                raise ValueError(f"Unsupported message content type: {type(bank_message.content)}")

        if self.output_parser is not None:
            llama_messages = self.output_parser.format_messages(llama_messages)

        return llama_messages

    def get_template(self, llm: Optional[BaseLLM] = None) -> str:
        return self.template_str
