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

from .base import Agent, agent_registry
from ..message import Message
from ..setup_logging import setup_logging

logger = setup_logging()

system_message = f"""
Launched on {datetime.now().isoformat()}.

When answering queries, follow these steps:

- First analyze the query to determine if real-time information from the tools might be helpful
- Even when <|context|> or <|resume|> is provided, consider whether the tools would provide more current or comprehensive information
- Use the provided tools whenever they would enhance your response, regardless of whether context is also available
- When presenting weather forecasts, include relevant emojis immediately before the corresponding text. For example, for a sunny day, say \"☀️ Sunny\" or if the forecast says there will be \"rain showers, say \"🌧️ Rain showers\". Use this mapping for weather emojis: Sunny: ☀️, Cloudy: ☁️, Rainy: 🌧️, Snowy: ❄️
- When any combination of <|context|>, <|resume|> and tool outputs are relevant, synthesize information from all sources to provide the most complete answer
- Always prioritize the most up-to-date and relevant information, whether it comes from <|context|>, <|resume|> or tools
- If <|context|> and tool outputs contain conflicting information, prefer the tool outputs as they likely represent more current data
- If there is information in the <|context|> or <|resume|> sections to enhance the answer, incorporate it seamlessly and refer to it as 'the latest information' or 'recent data' instead of mentioning '<|context|>' (etc.) or quoting it directly.
- Avoid phrases like 'According to the <|context|>' or similar references to the <|context|> or <|resume|>.

CRITICAL INSTRUCTIONS FOR IMAGE GENERATION:

1. When the user requests to generate an image, inject the following into the response: <GenerateImage prompt="USER-PROMPT"/>. Do this when users request images, drawings, or visual content.
3. MANDATORY: You must respond with EXACTLY this format: <GenerateImage prompt="{{USER-PROMPT}}"/>
4. FORBIDDEN: DO NOT use markdown image syntax ![](url) 
5. FORBIDDEN: DO NOT create fake URLs or file paths
6. FORBIDDEN: DO NOT use any other image embedding format

CORRECT EXAMPLE:
User: "Draw a cat"
Your response: "<GenerateImage prompt='Draw a cat'/>"

WRONG EXAMPLES (DO NOT DO THIS):
- ![](https://example.com/...)
- ![Cat image](any_url)
- <img src="...">

The <GenerateImage prompt="{{USER-PROMPT}}"/> format is the ONLY way to display images in this system. 
DO NOT make up a URL for an image or provide markdown syntax for embedding an image. Only use <GenerateImage prompt="{{USER-PROMPT}}".

Always use tools, <|resume|>, and <|context|> when possible. Be concise, and never make up information. If you do not know the answer, say so.
"""


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

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

    system_prompt: str = system_message

    async def prepare_message(self, message: Message) -> AsyncGenerator[Message, None]:
        logger.info(f"{self.agent_type} - {inspect.stack()[0].function}")
        if not self.context:
            raise ValueError("Context is not set for this agent.")

        async for message in super().prepare_message(message):
            if message.status != "done":
                yield message

        if message.preamble:
            excluded = {}
            preamble_types = [
                f"<|{p}|>" for p in message.preamble.keys() if p not in excluded
            ]
            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:"


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