from __future__ import annotations
from typing import Literal, AsyncGenerator, ClassVar, Optional, Any, List
from datetime import datetime, UTC
import inspect

from .base import Agent, agent_registry
from logger import logger

from .registry import agent_registry
from models import ( ChatMessage, ApiStatusType, ChatMessage, ChatMessageError, ChatMessageRagSearch, ChatMessageStatus, ChatMessageStreaming, ChatOptions, ApiMessageType, ChatSenderType, ApiStatusType, ChatMessageMetaData, Candidate, Tunables )
from rag import ( ChromaDBGetResponse )

class Chat(Agent):
    """
    Chat Agent
    """

    agent_type: Literal["rag_search"] = "rag_search"  # type: ignore
    _agent_type: ClassVar[str] = agent_type  # Add this for registration

    async def generate(
        self, llm: Any, model: str, 
        session_id: str, prompt: str, 
        tunables: Optional[Tunables] = None, 
        temperature=0.7
    ) -> AsyncGenerator[ChatMessage | ChatMessageStatus | ChatMessageError | ChatMessageStreaming, None]:
        """
        Generate a response based on the user message and the provided LLM.

        Args:
            llm: The language model to use for generation.
            model: The specific model to use.
            user_message: The message from the user.
            user: Optional user information.
            temperature: The temperature setting for generation.

        Yields:
            ChatMessage: The generated response.
        """
        rag_message = None
        async for rag_message in self.generate_rag_results(session_id, prompt):
            if rag_message.status == ApiStatusType.ERROR:
                yield rag_message
                return
            if rag_message.status != ApiStatusType.DONE:
                yield rag_message

        if not isinstance(rag_message, ChatMessageRagSearch):
            logger.error(f"Expected ChatMessageRagSearch, got {type(rag_message)}")
            error_message = ChatMessageError(
                session_id=session_id,
                content="RAG search did not return a valid response."
            )
            yield error_message
            return
            
        rag_message.status = ApiStatusType.DONE
        yield rag_message

# Register the base agent
agent_registry.register(Chat._agent_type, Chat)
