from typing import List, Dict
from pic_prompt.core import PromptMessage, PromptConfig
from pic_prompt.providers import ProviderFactory, Provider
from pic_prompt.images import ImageData
from pic_prompt.images.image_registry import ImageRegistry
from pic_prompt.providers.provider_names import ProviderNames
from pic_prompt.utils.logger import setup_logger
from typing import Optional

logger = setup_logger(__name__)


class PicPrompt:
    """
    Main Builder Class for constructing image-based prompts for OpenAI.
    To use other providers, use the prompt generated by this class
    as input to LiteLLM.

    - Managing provider-specific configurations
    - Building message sequences including text and images
    - Maintaining an image registry for downloaded images
    - Caching generated prompts
    - Converting between different provider message formats

    Given this JSON structure:
    {
      "model": "gpt-4o",
      "messages": [
        {
          "role": "developer",
          "content": "You are a helpful assistant."
        },
        {
          "role": "user",
          "content": "Hello!"
        }
      ]
    }

    This class represents the whole block.
    """

    def __init__(self):
        """
        Initialize a new PicPrompt instance.

        This constructor sets up:
        - Configuration dictionary for different providers
        - Message lists for system, user and image messages
        - Provider factory and initialized providers
        - Image registry for downloaded image data
        - Prompt cache
        """
        self.configs: Dict[str, PromptConfig] = {}

        # These are the messages that will be used to build the prompt
        self.messages: List[PromptMessage] = []
        self.user_messages: List[PromptMessage] = []
        self.image_messages: List[PromptMessage] = []

        # This is the factory that will be used to get the provider helper
        self.provider_factory = ProviderFactory()

        self.providers: Dict[str, Provider] = {}
        self.init_all_providers()

        # This is the registry of all the downloaded image data
        self.image_registry = ImageRegistry()

        # This is the cache of all the prompts
        self.prompts: Dict[str, str] = {}

    def init_all_providers(self) -> None:
        prompt_config_openai = PromptConfig(
            provider_name="openai",
            model="gpt-4o",
            max_tokens=3000,
            temperature=0.0,
        )
        # prompt_config_anthropic = PromptConfig(
        #     provider_name="anthropic",
        #     model="claude-3-opus-20240229",
        #     max_tokens=3000,
        #     temperature=0.5,
        # )
        # prompt_config_gemini = PromptConfig(
        #     provider_name="gemini",
        #     model="gemini-2.0-flash",
        #     max_tokens=3000,
        #     temperature=0.0,
        # )
        self._add_config(prompt_config_openai)
        # self.add_config(prompt_config_anthropic)
        # self.add_config(prompt_config_gemini)

    # Message Methods
    def add_system_message(self, message: str) -> None:
        """
        Add a system message to the list of messages.

        Args:
            message (str): The system message text to add
        """
        pm = PromptMessage(role="system")
        pm.add_text(message)
        self.messages.append(pm)

    def add_user_message(self, message: str) -> None:
        """
        Add a user message to the list of user messages.

        Args:
            message (str): The user message text to add
        """
        pm = PromptMessage(role="user")
        pm.add_text(message)
        self.user_messages.append(pm)

    def set_user_message(self, message: str) -> None:
        """
        Set a single user message, clearing any existing user messages.

        Args:
            message (str): The user message text to set
        """
        self.user_messages = []
        self.add_user_message(message)

    def add_assistant_message(self, message: str) -> None:
        """
        Add an assistant message to the list of messages.

        Args:
            message (str): The assistant message text to add
        """
        pm = PromptMessage(role="assistant")
        pm.add_text(message)
        self.messages.append(pm)

    def add_image_message(self, image_path: str) -> None:
        """
        Add an image message to the list of image messages.
        The image path is stored initially and processed later when building the prompt.

        Args:
            image_path (str): Path to the image file to add
        """
        pm = PromptMessage(role="user")
        pm.add_image(image_path)
        self.image_registry.add_image_path(image_path)
        self.image_messages.append(pm)

    def add_image_data(self, image_data: ImageData) -> None:
        """
        Add an image message using pre-downloaded image data.

        Args:
            image_data (ImageData): The pre-downloaded image data to add
        """
        self.image_registry.add_image_data(image_data)
        pm = PromptMessage(role="user")
        pm.add_image(image_data.image_path)
        self.image_messages.append(pm)

    def add_image_messages(self, image_paths: List[str]) -> None:
        """
        Add multiple image messages to the list of image messages.

        Args:
            image_paths (List[str]): List of paths to image files to add
        """
        for path in image_paths:
            self.add_image_message(path)

    # Config Methods
    def _add_config(self, config: PromptConfig) -> None:
        """
        Add a provider configuration to the prompt builder.

        Args:
            config (PromptConfig): The provider configuration to add

        Raises:
            ValueError: If the provider specified in the config is not supported
        """
        if config.provider_name not in ProviderNames.get_all_names():
            raise ValueError(f"Provider {config.provider_name} is not supported")
        self.configs[config.provider_name] = config
        # Reset providers list to force re-initialization with new config
        self.providers = {}

    def _encode_image_data(self) -> ImageRegistry:
        """
        Encode image data according to provider requirements.

        For each image in the registry and each configured provider:
        - If the provider requires base64 encoding or the image is local
        - Resize and encode the image according to provider's max size

        Returns:
            ImageRegistry: The image registry with encoded images
        """
        for image_data in self.image_registry.get_all_image_data():
            for provider in self._get_providers().values():
                if (
                    provider.get_image_config().requires_base64
                    or image_data.is_local_image()
                ):
                    image_data.resize_and_encode(
                        provider.get_image_config().max_size,
                        provider.get_provider_name(),
                    )
        return self.image_registry

    def _get_providers(self) -> Dict[str, Provider]:
        """
        Get or initialize providers based on current configs.

        Lazily initializes providers if they haven't been created yet.
        Each provider is created using the provider factory with its corresponding config.

        Returns:
            Dict[str, Provider]: Dictionary mapping provider names to provider instances
        """
        if len(self.providers) == 0:
            for provider_name, config in self.configs.items():
                helper = self.provider_factory.get_provider(config.provider_name)
                self.providers[provider_name] = helper
        return self.providers

    def build(self):
        """
        Builds prompts for all configured providers.

        This method:
        1. Downloads any required image data if needed
        2. Encodes the image data according to each provider's requirements

        The formatted prompts can be retrieved using get_content_for().
        """
        self.image_registry.download_image_data()
        self._encode_image_data()

    # def get_content_for(self, provider_name: str, preview=False) -> str:
    #     # if len(self.prompts) == 0:
    #     self.build()
    #     provider = self._get_providers().get(provider_name)
    #     if provider is None:
    #         raise ValueError(f"Provider {provider_name} not found")
    #     messages = self.messages + self.user_messages + self.image_messages
    #     return provider.format_messages(messages, self.image_registry, preview)

    def get_prompt(self, preview=False) -> str:
        """
        Get the formatted prompt for the OpenAI provider.

        This method:
        1. Builds the prompt by downloading and encoding image data
        2. Gets the OpenAI provider instance
        3. Combines all messages (base, user, and image messages)
        4. Formats the messages according to OpenAI's requirements

        Args:
            preview (bool, optional): Whether to generate a preview version. Defaults to False.

        Returns:
            str: The formatted prompt string for OpenAI

        Raises:
            ValueError: If the OpenAI provider is not found
        """
        # if len(self.prompts) == 0:
        self.build()
        provider = self._get_providers().get("openai")
        if provider is None:
            raise ValueError(f"Provider openai not found")
        messages = self.messages + self.user_messages + self.image_messages
        return provider.format_messages(messages, self.image_registry, preview)

    def clear(self) -> None:
        """
        Clears all messages from the prompt builder.

        This method:
        1. Clears the base messages list
        2. Clears the user messages list
        3. Clears the image messages list
        """
        self.messages = []
        self.user_messages = []
        self.image_messages = []

    def __repr__(self) -> str:
        return f"<PromptBuilder messages={self.messages}>"
