from typing import Optional, List, Dict, Union, Iterator
from pydantic import BaseModel, Field, ConfigDict
from PIL.Image import Image
import requests
import io
from .rag import RAG
from .llm.base_llm import BaseLLM
from .knowledge_base.retriever import Retriever
from .knowledge_base.vector_store import VectorStore

class Assistant(BaseModel):
    # -*- Agent settings
    name: Optional[str] = Field(None, description="Name of the assistant.")
    description: Optional[str] = Field(None, description="Description of the assistant's role.")
    instructions: Optional[List[str]] = Field(None, description="List of instructions for the assistant.")
    model: Optional[str] = Field(None, description="The LLM model to use (e.g., Groq, OpenAI).")
    show_tool_calls: bool = Field(False, description="Whether to show tool calls in the response.")
    markdown: bool = Field(False, description="Whether to format the response in markdown.")
    tools: Optional[List[Union[str, Dict]]] = Field(None, description="List of tools available to the assistant.")
    user_name: Optional[str] = Field("User", description="Name of the user interacting with the assistant.")
    emoji: Optional[str] = Field(":robot:", description="Emoji to represent the assistant in the CLI.")
    rag: Optional[RAG] = Field(None, description="RAG instance for context retrieval.")
    knowledge_base: Optional[Dict] = Field(None, description="Knowledge base for domain-specific information.")
    llm: Optional[BaseLLM] = Field(None, description="The LLM instance to use.")

    # Allow arbitrary types
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Initialize the model and tools here if needed
        self._initialize_model()
        # Initialize RAG if not provided
        if self.rag is None:
            self.rag = self._initialize_default_rag()

    def _initialize_model(self):
        """Initialize the model based on the provided configuration."""
        if self.llm is not None:
            return  # LLM is already initialized, do nothing
        if self.model is None:
            raise ValueError("Model must be specified.")
        if self.model.lower() == "groq":
            from .llm.groq import GroqLlm
            self.llm = GroqLlm(api_key=self.llm.api_key if self.llm else None)
        elif self.model.lower() == "openai":
            from .llm.openai import OpenAILlm
            self.llm = OpenAILlm()
        else:
            raise ValueError(f"Unsupported model: {self.model}")

    def _initialize_default_rag(self) -> RAG:
        """Initialize a default RAG instance with a dummy vector store."""
        vector_store = VectorStore()
        retriever = Retriever(vector_store)
        return RAG(retriever)

    def load_image_from_url(self, image_url: str) -> Image:
        """Load an image from a URL and return it as a PIL Image."""
        response = requests.get(image_url)
        image_bytes = response.content
        return Image.open(io.BytesIO(image_bytes))

    def print_response(
        self,
        message: Optional[Union[str, Image, List, Dict]] = None,
        stream: bool = False,
        markdown: bool = False,
        **kwargs,
    ):
        """Print the assistant's response to the console."""
        if isinstance(message, Image):
            # Handle image input
            message = self._process_image(message)

        if stream:
            for chunk in self._stream_response(message, markdown=markdown, **kwargs):
                print(chunk)
        else:
            response = self._generate_response(message, markdown=markdown, **kwargs)
            print(response)

    def _process_image(self, image: Image) -> str:
        """Process the image and return a string representation."""
        # Convert the image to text or extract relevant information
        # For now, we'll just return a placeholder string
        return "Image processed. Extracted text: [Placeholder]"

    def _stream_response(self, message: str, markdown: bool = False, **kwargs) -> Iterator[str]:
        """Stream the assistant's response."""
        # Simulate streaming by yielding chunks of the response
        response = self._generate_response(message, markdown=markdown, **kwargs)
        for chunk in response.split():
            yield chunk + " "

    def _generate_response(self, message: str, markdown: bool = False, **kwargs) -> str:
        """Generate the assistant's response."""
        # Retrieve relevant context using RAG
        context = self.rag.retrieve(message) if self.rag else None

        # Generate the response using the LLM
        response = self.llm.generate(prompt=message, context=context, **kwargs)
        if markdown:
            return f"**Response:**\n\n{response}"
        return response

    def cli_app(
        self,
        message: Optional[str] = None,
        exit_on: Optional[List[str]] = None,
        **kwargs,
    ):
        """Run the assistant in a CLI app."""
        from rich.prompt import Prompt

        if message:
            self.print_response(message=message, **kwargs)

        _exit_on = exit_on or ["exit", "quit", "bye"]
        while True:
            message = Prompt.ask(f"[bold] {self.emoji} {self.user_name} [/bold]")
            if message in _exit_on:
                break

            self.print_response(message=message, **kwargs)