"""
Agent adapter module for integrating custom agents with the Galtea conversation simulation framework.

This module provides the abstract base class that users must implement to integrate
their existing agents with the Galtea conversation simulation framework. The adapter pattern allows
any agent implementation to work with the framework regardless of its underlying
architecture or API.

Additionally, this module supports plain agent functions as a simpler alternative to the
Agent class. Agent functions can be sync or async and support three signatures:
    - (user_message: str) -> str — receives just the last user message
    - (messages: List[Dict[str, str]]) -> str — receives the full conversation history
      in the OpenAI message standard format (keys "role" and "content")
    - (input_data: AgentInput) -> AgentResponse — receives structured input (with session
      context, helper methods) and returns structured output (with optional usage/cost tracking)
"""

from abc import ABC, abstractmethod
from typing import Any, Awaitable, Callable, Dict, List, Optional, Union

from galtea.domain.models.inference_result import CostInfoProperties, UsageInfoProperties
from galtea.utils.from_camel_case_base_model import FromCamelCaseBaseModel


class ConversationMessage(FromCamelCaseBaseModel):
    """
    Represents a single message in a conversation.

    Attributes:
        role (str): The role of the message sender (e.g., "user", "assistant")
        content (str): The message content
        metadata (Optional[Dict[str, Any]]): Additional metadata for the message
    """

    role: str
    content: str
    retrieval_context: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None


class AgentInput(FromCamelCaseBaseModel):
    """
    Input data provided to the agent for processing.

    Attributes:
        messages (List[ConversationMessage]): The conversation history
        session_id (str): The current session identifier
        context (Optional[str]): Additional context information
        metadata (Optional[Dict[str, Any]]): Additional metadata
        inference_result_id (Optional[str]): The inference result ID for this execution.
            Useful for remote agents that need to correlate traces across HTTP boundaries.
            When your agent runs on a remote server, forward this ID and use
            ``set_context(inference_result_id=...)`` on the server to attach traces.
    """

    messages: List[ConversationMessage]
    session_id: str
    context: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    inference_result_id: Optional[str] = None

    def last_user_message(self) -> Optional[ConversationMessage]:
        """
        Get the last user message from the conversation.

        Returns:
            Optional[ConversationMessage]: The last user message or None if no user message exists
        """
        for message in reversed(self.messages):
            if message.role == "user":
                return message
        return None

    def last_user_message_str(self) -> Optional[str]:
        """
        Get the content of the last user message.

        Returns:
            Optional[str]: The content of the last user message or None if no user message exists
        """
        last_message: Optional[ConversationMessage] = self.last_user_message()
        return last_message.content if last_message else None


class AgentResponse(FromCamelCaseBaseModel):
    """
    Response from the agent after processing input.

    Attributes:
        content (str): The response content
        retrieval_context (Optional[str]): Context retrieved for RAG systems
        metadata (Optional[Dict[str, Any]]): Additional metadata for the response
        usage_info (Optional[Dict[str, int]]): Token usage information from the LLM call
            Keys: 'input_tokens', 'output_tokens', 'cache_read_input_tokens'
        cost_info (Optional[Dict[str, float]]): Cost information from the LLM call
            Keys: 'cost_per_input_token', 'cost_per_output_token', 'cost_per_cache_read_input_token'
    """

    content: str
    retrieval_context: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    usage_info: Optional[UsageInfoProperties] = None
    cost_info: Optional[CostInfoProperties] = None


class Agent(ABC):
    """
    Abstract base class for integrating custom agents with the Galtea conversation simulation framework.

    This adapter pattern allows you to wrap any existing agent implementation
    (LLM calls, agent frameworks, or complex multi-step systems) to work with
    the Galtea conversation simulation framework. The adapter receives structured input about
    the conversation state and returns responses in a standardized format.

    Example:
        ```python
        import galtea
        from my_agent import MyCustomAgent

        class MyAgentAdapter(galtea.Agent):
            def __init__(self):
                self.agent = MyCustomAgent()

            def call(self, input_data: galtea.AgentInput) -> galtea.AgentResponse:
                # Get the latest user message
                user_message = input_data.last_user_message_str()

                # Call your existing agent
                response = self.agent.process(
                    message=user_message,
                    history=input_data.messages,
                    session_id=input_data.session_id
                )

                # Return the response
                return galtea.AgentResponse(content=response)

        # Use in a simulation
        client = galtea.Galtea(api_key="your_api_key")
        my_agent = MyAgentAdapter()

        result = client.simulator.simulate(
            session_id="session_123",
            agent=my_agent,
            max_turns=10
        )
        ```

    Note:
        - The call method must return an AgentResponse object
        - For stateful agents, use input_data.session_id to maintain conversation context
        - For stateless agents, use input_data.messages for the full conversation history
    """

    @abstractmethod
    def call(self, input_data: AgentInput) -> AgentResponse:
        """
        Process the input and generate a response.

        This is the main method that your agent implementation must provide.
        It receives structured information about the current conversation state
        and must return a response in the standardized format.

        Args:
            input_data (AgentInput): Input containing conversation history, session context, and metadata

        Returns:
            AgentResponse: The agent's response containing the content and optional metadata

        Example:
            ```python
            def call(self, input_data: AgentInput) -> AgentResponse:
                # Simple string response
                user_msg = input_data.last_user_message_str()
                response_content = f"I understand you said: {user_msg}"

                return AgentResponse(
                    content=response_content,
                    metadata={"processing_time": 0.1}
                )
            ```
        """
        pass


# Type aliases for agent functions
# These follow the OpenAI/Anthropic standard chat message format

ChatMessage = Dict[str, str]
"""A single chat message with 'role' and 'content' keys."""

ChatHistory = List[ChatMessage]
"""A list of chat messages representing the conversation history."""

SyncAgentCallable = Callable[[ChatHistory], str]
"""
A synchronous agent function that takes a chat history and returns a response string.

Example:
    ```python
    def my_agent(messages: list[dict]) -> str:
        user_message = messages[-1]["content"]
        return f"Response to: {user_message}"
    ```
"""

AsyncAgentCallable = Callable[[ChatHistory], Awaitable[str]]
"""
An asynchronous agent function that takes a chat history and returns a response string.

Example:
    ```python
    async def my_async_agent(messages: list[dict]) -> str:
        user_message = messages[-1]["content"]
        response = await my_llm.chat(messages)
        return response
    ```
"""

AgentCallable = Union[SyncAgentCallable, AsyncAgentCallable]
"""An agent function (sync or async) that takes chat history and returns a response string."""

# Simple string-based agent functions (for agents that only need the last user message)
SimpleSyncAgentCallable = Callable[[str], str]
"""
A synchronous agent function that takes a single string (last user message) and returns a response.

Example:
    ```python
    def my_agent(user_message: str) -> str:
        return f"You said: {user_message}"
    ```
"""

SimpleAsyncAgentCallable = Callable[[str], Awaitable[str]]
"""
An asynchronous agent function that takes a single string (last user message) and returns a response.

Example:
    ```python
    async def my_agent(user_message: str) -> str:
        response = await my_llm.complete(user_message)
        return response
    ```
"""

SimpleAgentCallable = Union[SimpleSyncAgentCallable, SimpleAsyncAgentCallable]
"""A simple agent function (sync or async) that takes a string and returns a string."""

# Structured agent functions (mimic Agent.call() signature without requiring a class)
SyncStructuredAgentCallable = Callable[["AgentInput"], AgentResponse]
"""
A synchronous agent function that takes an AgentInput and returns an AgentResponse.

This signature mirrors Agent.call() but as a plain function, giving access to
structured input (session context, helper methods) and structured output
(usage/cost tracking, retrieval context) without requiring a class.

Example:
    ```python
    def my_agent(input_data: galtea.AgentInput) -> galtea.AgentResponse:
        user_message = input_data.last_user_message_str()
        return galtea.AgentResponse(
            content=f"Response to: {user_message}",
            usage_info={"input_tokens": 100, "output_tokens": 50},
        )
    ```
"""

AsyncStructuredAgentCallable = Callable[["AgentInput"], Awaitable[AgentResponse]]
"""
An asynchronous agent function that takes an AgentInput and returns an AgentResponse.

Example:
    ```python
    async def my_agent(input_data: galtea.AgentInput) -> galtea.AgentResponse:
        user_message = input_data.last_user_message_str()
        response = await my_llm.complete(user_message)
        return galtea.AgentResponse(content=response)
    ```
"""

StructuredAgentCallable = Union[SyncStructuredAgentCallable, AsyncStructuredAgentCallable]
"""A structured agent function (sync or async) that takes AgentInput and returns AgentResponse."""

AnyAgentCallable = Union[AgentCallable, SimpleAgentCallable, StructuredAgentCallable]
"""Any agent function that can be used as an agent (chat history, simple string, or structured)."""

AgentType = Union[Agent, AnyAgentCallable]
"""
Union type representing any valid agent: an Agent Function or an Agent Class instance.

This type is accepted by `simulator.simulate()` and `inference_results.generate()`.

Examples:
    ```python
    # Option 1: Agent function — receives just the last user message
    def my_agent(user_message: str) -> str:
        return f"You said: {user_message}"

    # Option 1 (alt): Agent function — receives full conversation history
    def my_agent(messages: list[dict]) -> str:
        return f"Response to: {messages[-1]['content']}"

    # Option 1 (alt): Agent function — structured input/output (mirrors Agent.call())
    def my_agent(input_data: galtea.AgentInput) -> galtea.AgentResponse:
        user_message = input_data.last_user_message_str()
        return galtea.AgentResponse(
            content=f"Response to: {user_message}",
            usage_info={"input_tokens": 100, "output_tokens": 50},
        )

    # Async agent functions work too
    async def my_agent(user_message: str) -> str:
        return await get_response(user_message)

    # Option 2: Agent class (for advanced features like usage/cost tracking)
    class MyAgent(galtea.Agent):
        def call(self, input_data: galtea.AgentInput) -> galtea.AgentResponse:
            return galtea.AgentResponse(content="Hello!")

    # Both options can be used with simulate() or generate()
    result = client.simulator.simulate(session_id=s.id, agent=my_agent, max_turns=5)
    ```
"""
