from __future__ import annotations
from typing import Literal, AsyncGenerator, ClassVar, Optional
import logging
from . base import Agent, registry
from .. message import Message
import inspect

class Chat(Agent):
  """
  Chat Agent
  """
  agent_type: Literal["chat"] = "chat" # type: ignore
  _agent_type: ClassVar[str] = agent_type  # Add this for registration

  async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
    """
    Prepare message with context information in message.preamble
    """
    logging.info(f"{self.agent_type} - {inspect.stack()[1].function}")

    if not self.context:
       raise ValueError("Context is not set for this agent.")
    
    # Generate RAG content if enabled, based on the content
    rag_context = ""
    if message.enable_rag:
        # Gather RAG results, yielding each result
        # as it becomes available
        for message in self.context.generate_rag_results(message):
            logging.info(f"RAG: {message.status} - {message.response}")
            if message.status == "error":
                yield message
                return
            if message.status != "done":
                yield message
            
        if "rag" in message.metadata and message.metadata["rag"]:
            for rag in message.metadata["rag"]:
                for doc in rag["documents"]:
                  rag_context += f"{doc}\n"

    message.preamble = {}

    if rag_context:
        message.preamble["context"] = rag_context

    if self.context.user_resume:
        message.preamble["resume"] = self.context.user_resume

    if message.preamble:
        preamble_types = [f"<|{p}|>" for p in message.preamble.keys()]
        preamble_types_AND = " and ".join(preamble_types)
        preamble_types_OR = " or ".join(preamble_types)
        message.preamble["rules"] = f"""\
- Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
- If there is no information in these sections, answer based on your knowledge, or use any available tools.
- Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
"""
    message.preamble["question"] = "Respond to:"
    
    message.system_prompt = self.system_prompt
    message.status = "done"
    yield message
    return

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