"""
Agent - Agent with built-in inbox memory and structured response rules.

This module provides the Agent class that wraps agent logic functions
with inbox memory access and structured JSON response handling.
"""

import json
from typing import Any, Callable, Coroutine, Optional

from synqed.memory import AgentMemory, InboxMessage


class ResponseBuilder:
    """
    Helper class for building structured responses.
    
    Agents use this to emit JSON responses with "send_to" and "content" fields.
    
    Example:
        ```python
        builder = ResponseBuilder()
        builder.send_to("Editor").content("Here's my draft")
        response = builder.build()  # {"send_to": "Editor", "content": "..."}
        ```
    """
    
    def __init__(self):
        """Initialize the response builder."""
        self._send_to: Optional[str] = None
        self._content: Optional[str] = None
    
    def send_to(self, agent_name: str) -> "ResponseBuilder":
        """
        Set the target agent name.
        
        Args:
            agent_name: Name of the agent to send the message to
            
        Returns:
            Self for method chaining
        """
        self._send_to = agent_name
        return self
    
    def content(self, text: str) -> "ResponseBuilder":
        """
        Set the message content.
        
        Args:
            text: Message content text
            
        Returns:
            Self for method chaining
        """
        self._content = text
        return self
    
    def build(self) -> dict[str, str]:
        """
        Build the structured response dictionary.
        
        Returns:
            Dictionary with "send_to" and "content" keys
            
        Raises:
            ValueError: If send_to or content is not set
        """
        if self._send_to is None:
            raise ValueError("send_to must be set before building response")
        if self._content is None:
            raise ValueError("content must be set before building response")
        
        return {
            "send_to": self._send_to,
            "content": self._content
        }
    
    def to_json(self) -> str:
        """
        Build and return as JSON string.
        
        Returns:
            JSON string representation of the response
        """
        return json.dumps(self.build())


class AgentLogicContext:
    """
    Context object passed to agent logic functions.
    
    Provides access to:
    - Agent's memory
    - Latest incoming message
    - ResponseBuilder helper
    """
    
    def __init__(self, memory: AgentMemory, default_target: Optional[str] = None):
        """
        Initialize the logic context.
        
        Args:
            memory: Agent's memory instance
            default_target: Default target agent if logic doesn't specify one
        """
        self.memory = memory
        self.default_target = default_target
        self._response_builder = ResponseBuilder()
    
    @property
    def latest_message(self) -> Optional[InboxMessage]:
        """Get the latest incoming message."""
        return self.memory.get_latest_message()
    
    @property
    def response(self) -> ResponseBuilder:
        """Get the response builder helper."""
        return self._response_builder
    
    def build_response(self, send_to: str, content: str) -> dict[str, str]:
        """
        Convenience method to build a response.
        
        Args:
            send_to: Target agent name
            content: Message content
            
        Returns:
            Structured response dictionary
        """
        return ResponseBuilder().send_to(send_to).content(content).build()


class Agent:
    """
    Agent with built-in inbox memory, JSON-structured response rules, and message processing logic.
    
    This class wraps a logic function and provides it with access to:
    - The agent's memory (via AgentLogicContext)
    - The latest incoming message
    - A ResponseBuilder helper for structured responses
    
    The logic function should return JSON with "send_to" and "content" fields.
    If it returns non-JSON, it will be automatically wrapped.
    
    Example:
        ```python
        async def writer_logic(context: AgentLogicContext) -> dict:
            latest = context.latest_message
            if not latest:
                return context.build_response("Editor", "I'm ready!")
            
            # Process message and respond
            return context.build_response("Editor", "Here's my draft")
        
        agent = Agent(
            name="Writer",
            description="Creative writer",
            logic=writer_logic,
            default_target="Editor"
        )
        ```
    """
    
    def __init__(
        self,
        name: str,
        description: str,
        logic: Callable[[AgentLogicContext], Coroutine[Any, Any, dict[str, str] | str]],
        default_target: Optional[str] = None,
        memory: Optional[AgentMemory] = None,
    ):
        """
        Initialize an Agent.
        
        Args:
            name: Agent name
            description: Agent description
            logic: Async function that takes AgentLogicContext and returns
                   dict with "send_to" and "content", or a string (will be wrapped)
            default_target: Default target agent if logic returns non-JSON or
                           doesn't specify send_to
            memory: Optional pre-existing memory instance (creates new one if not provided)
        """
        self.name = name
        self.description = description
        self.logic = logic
        self.default_target = default_target
        self.memory = memory or AgentMemory(agent_name=name)
    
    async def process(self) -> dict[str, str]:
        """
        Process the agent's inbox and generate a response.
        
        This method:
        1. Creates a logic context with access to memory
        2. Calls the agent's logic function
        3. Ensures the response is valid JSON with "send_to" and "content"
        4. Wraps non-JSON responses automatically
        
        Returns:
            Dictionary with "send_to" and "content" keys
            
        Raises:
            ValueError: If response cannot be parsed or structured
        """
        # Create context for logic function
        context = AgentLogicContext(
            memory=self.memory,
            default_target=self.default_target
        )
        
        # Call logic function
        result = await self.logic(context)
        
        # Ensure result is structured JSON
        return self._ensure_structured_response(result)
    
    def _ensure_structured_response(self, result: dict[str, str] | str) -> dict[str, str]:
        """
        Ensure the response is valid structured JSON.
        
        Args:
            result: Result from logic function (dict or string)
            
        Returns:
            Valid structured response dictionary
        """
        # If already a dict, validate structure
        if isinstance(result, dict):
            if "send_to" in result and "content" in result:
                return result
            else:
                # Invalid structure, wrap it
                content = json.dumps(result) if not isinstance(result.get("content"), str) else result.get("content", str(result))
                return {
                    "send_to": self.default_target or "Unknown",
                    "content": content
                }
        
        # If string, try to parse as JSON
        if isinstance(result, str):
            try:
                parsed = json.loads(result)
                if isinstance(parsed, dict) and "send_to" in parsed and "content" in parsed:
                    return parsed
                else:
                    # Invalid JSON structure, wrap it
                    return {
                        "send_to": self.default_target or "Unknown",
                        "content": result
                    }
            except json.JSONDecodeError:
                # Not JSON, wrap it
                return {
                    "send_to": self.default_target or "Unknown",
                    "content": result
                }
        
        # Fallback: convert to string and wrap
        return {
            "send_to": self.default_target or "Unknown",
            "content": str(result)
        }
    
    def __repr__(self) -> str:
        """String representation."""
        return f"Agent(name='{self.name}', memory_messages={len(self.memory.messages)})"
