"""
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, List

from synqed.memory import AgentMemory, InboxMessage


def parse_llm_response(response_text: str) -> str:
    """
    Parse LLM response by stripping markdown code blocks.
    
    Many LLMs wrap JSON responses in markdown code blocks like:
    ```json
    {"send_to": "Agent", "content": "..."}
    ```
    
    This function removes those code blocks to extract the actual content.
    
    Args:
        response_text: Raw LLM response text
        
    Returns:
        Cleaned response text with markdown code blocks removed
    """
    if "```" not in response_text:
        return response_text
    
    # Find the JSON content between code blocks
    start_idx = response_text.find("```")
    if start_idx != -1:
        # Skip the opening ``` and optional language identifier
        start_idx = response_text.find("\n", start_idx)
        if start_idx != -1:
            end_idx = response_text.find("```", start_idx)
            if end_idx != -1:
                return response_text[start_idx:end_idx].strip()
    
    # Fallback: return original
    return response_text


def extract_partial_json_content(response_text: str) -> Optional[dict]:
    """
    Extract content from partial/truncated JSON responses.
    
    Handles cases where LLM response was truncated mid-JSON, like:
    {"send_to": "USER", "content": "story text...
    
    Args:
        response_text: Potentially truncated JSON response
        
    Returns:
        Extracted dict with send_to and content if parseable, None otherwise
    """
    # Check if response was INTENDED for specific recipient
    send_to_match = None
    for pattern in ['"send_to": "', '"send_to":"']:
        if pattern in response_text:
            start = response_text.find(pattern) + len(pattern)
            end = response_text.find('"', start)
            if end != -1:
                send_to_match = response_text[start:end]
                break
    
    if not send_to_match:
        return None
    
    # Try to find the content field
    content_match = None
    for pattern in ['"content": "', '"content":"']:
        if pattern in response_text:
            content_start = response_text.find(pattern) + len(pattern)
            # Extract until end or closing quote
            content_text = response_text[content_start:]
            # Try to find natural end (closing quote followed by })
            # But handle truncation gracefully
            content_match = content_text.rstrip('"}')
            break
    
    if content_match is None:
        return None
    
    return {
        "send_to": send_to_match,
        "content": content_match
    }


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
    - send() helper method for convenience
    - get_conversation_history() for formatted conversation history
    - workspace reference for transcript access
    
    WARNING: Agent logic functions must NOT directly use any router.
    Routing is strictly handled by ExecutionEngine + Workspace.
    """
    
    def __init__(
        self, 
        memory: AgentMemory, 
        default_target: Optional[str] = None,
        workspace: Optional[Any] = None,
        agent_name: Optional[str] = None
    ):
        """
        Initialize the logic context.
        
        Args:
            memory: Agent's memory instance
            default_target: Default target agent if logic doesn't specify one
            workspace: Optional workspace reference for transcript access
            agent_name: Optional agent name for conversation history filtering
        """
        self.memory = memory
        self.default_target = default_target
        self.workspace = workspace
        self.agent_name = agent_name
        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 send(self, to: str, content: str) -> dict[str, str]:
        """
        Helper method to create a structured response.
        
        Args:
            to: Target agent name
            content: Message content
            
        Returns:
            Dictionary with "send_to" and "content" keys
        """
        return {"send_to": to, "content": content}
    
    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()
    
    def get_conversation_history(
        self, 
        format: str = "text",
        include_system_messages: bool = False,
        parse_json_content: bool = True
    ) -> str | List[dict]:
        """
        Get formatted conversation history for this agent.
        
        This method extracts the conversation history from the workspace transcript,
        including both messages received and sent by this agent. It automatically
        parses JSON content to extract the actual message content.
        
        Args:
            format: Output format - "text" for formatted string, "raw" for list of dicts
            include_system_messages: Whether to include system messages like [startup]
            parse_json_content: Whether to parse JSON content and extract "content" field
            
        Returns:
            Formatted conversation history as string (if format="text") or list of dicts (if format="raw")
            
        Example:
            ```python
            async def agent_logic(context):
                # Get conversation history as formatted text
                history = context.get_conversation_history()
                
                # Pass to LLM
                response = await llm.chat(history)
                return context.send("Editor", response)
            ```
        """
        if not self.workspace or not self.agent_name:
            # No workspace context, return empty
            return "" if format == "text" else []
        
        transcript = self.workspace.router.get_transcript()
        conversation_parts = []
        raw_messages = []
        
        for entry in transcript:
            sender = entry.get("from", "")
            recipient = entry.get("to", "")
            content = entry.get("content", "")
            
            # Skip system messages unless requested
            if not include_system_messages:
                if content == "[startup]" or content.startswith("[subteam_result]"):
                    continue
            
            # Only include messages involving this agent (received or sent)
            if recipient == self.agent_name or sender == self.agent_name:
                # Parse JSON content if requested
                display_content = content
                if parse_json_content:
                    try:
                        parsed = json.loads(content)
                        if isinstance(parsed, dict) and "content" in parsed:
                            display_content = parsed["content"]
                    except (json.JSONDecodeError, TypeError):
                        # Not JSON or no "content" field, use as-is
                        pass
                
                # Build message entry
                if format == "text":
                    if sender == "USER":
                        conversation_parts.append(f"User: {display_content}")
                    elif sender == self.agent_name:
                        conversation_parts.append(f"{self.agent_name} (you): {display_content}")
                    else:
                        conversation_parts.append(f"{sender}: {display_content}")
                else:
                    raw_messages.append({
                        "sender": sender,
                        "recipient": recipient,
                        "content": display_content,
                        "original_content": content,
                        "timestamp": entry.get("timestamp", "")
                    })
        
        if format == "text":
            return "\n\n".join(conversation_parts)
        else:
            return raw_messages


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.
    
    WARNING: Agent logic functions must NOT directly use any router.
    Routing is strictly handled by ExecutionEngine + Workspace. Any attempt
    to access or use a router within agent logic will raise an assertion error.
    
    Example:
        ```python
        async def writer_logic(context: AgentLogicContext) -> dict:
            latest = context.latest_message
            if not latest:
                return context.send("Editor", "I'm ready!")
            
            # Process message and respond
            return context.send("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)
        """
        if not name:
            raise ValueError("Agent name must be provided")
        
        if logic is None:
            raise ValueError("Agent logic function must be provided")
        
        # Validate logic is async
        import inspect
        if not inspect.iscoroutinefunction(logic):
            raise ValueError("Agent logic must be an async function")
        
        self.name = name
        self.description = description
        self.logic = logic
        self.default_target = default_target
        self.memory = memory or AgentMemory(agent_name=name)
        
        # Ensure memory has agent name set
        if self.memory.agent_name != name:
            self.memory.agent_name = name
    
    def __deepcopy__(self, memo):
        """
        Deep copy the agent to ensure isolation between workspaces.
        
        Creates new instances of mutable objects to prevent shared state.
        Memory is fully deep-copied, but the logic function reference is NOT
        deep-copied (functions are immutable and should be shared).
        """
        import copy
        cls = self.__class__
        result = cls.__new__(cls)
        memo[id(self)] = result
        
        # Copy immutable fields
        result.name = self.name
        result.description = self.description
        result.logic = self.logic  # Function reference must NOT be deepcopied
        result.default_target = self.default_target
        
        # Deep copy memory - ensure no shared mutable fields remain
        result.memory = copy.deepcopy(self.memory, memo)
        
        return result
    
    async def process(self, context: AgentLogicContext) -> dict[str, str]:
        """
        Process the agent's inbox and generate a response.
        
        This method accepts an externally supplied context and calls the
        agent's logic function with it. The context must be created by the
        caller (typically ExecutionEngine or Workspace).
        
        This method:
        1. Validates that context is provided
        2. Ensures agent logic does not directly use router
        3. Calls the agent's logic function with the provided context
        4. Ensures the response is valid JSON with "send_to" and "content"
        5. Wraps non-JSON responses automatically
        
        Args:
            context: AgentLogicContext instance (must be externally supplied)
        
        Returns:
            Dictionary with "send_to" and "content" keys
            
        Raises:
            ValueError: If context is None or response cannot be parsed or structured
            AssertionError: If agent logic attempts to use router directly
        """
        if context is None:
            raise ValueError("Agent.process() requires an externally supplied context")
        
        # Ensure agent logic does not directly use router
        # This is enforced by checking that context doesn't have router access
        # (Router should not be accessible from AgentLogicContext)
        assert not hasattr(context, 'router'), (
            "Agent logic must NOT directly use any router. "
            "Routing is strictly handled by ExecutionEngine + Workspace."
        )
        
        # Call logic function with provided context
        result = await self.logic(context)
        
        # Ensure result is structured JSON
        return self._ensure_structured_response(result, context.default_target)
    
    def _ensure_structured_response(
        self, result: dict[str, str] | str, default_target: Optional[str] = None
    ) -> dict[str, str]:
        """
        Ensure the response is valid structured JSON with strict validation.
        
        This method automatically handles:
        - Markdown code block stripping (```json ... ```)
        - Partial/truncated JSON responses
        - Missing "send_to" fields (uses default_target)
        - Missing "content" fields (converts whole object to JSON string)
        
        Validation rules:
        - If missing "send_to" → use default_target
        - If missing "content" → convert whole object to JSON string
        - If default_target is None → raise descriptive error
        
        Args:
            result: Result from logic function (dict or string)
            default_target: Default target agent (from context)
            
        Returns:
            Valid structured response dictionary
            
        Raises:
            ValueError: If default_target is None and "send_to" is missing
        """
        # Use provided default_target or fall back to instance default_target
        effective_default = default_target or self.default_target
        
        # If already a dict, validate structure
        if isinstance(result, dict):
            send_to = result.get("send_to")
            content = result.get("content")
            
            # If missing "send_to" → use default_target
            if send_to is None:
                if effective_default is None:
                    raise ValueError(
                        "Response missing 'send_to' field and no default_target is set. "
                        "Either provide 'send_to' in the response or set a default_target."
                    )
                send_to = effective_default
            
            # If missing "content" → convert whole object to JSON string
            if content is None:
                content = json.dumps(result)
            
            return {
                "send_to": send_to,
                "content": content
            }
        
        # If string, try to parse as JSON (with automatic cleaning)
        if isinstance(result, str):
            # Step 1: Strip markdown code blocks
            cleaned_result = parse_llm_response(result)
            
            # Step 2: Try to parse as JSON
            try:
                parsed = json.loads(cleaned_result)
                if isinstance(parsed, dict):
                    send_to = parsed.get("send_to")
                    content = parsed.get("content")
                    
                    # If missing "send_to" → use default_target
                    if send_to is None:
                        if effective_default is None:
                            raise ValueError(
                                "Response missing 'send_to' field and no default_target is set. "
                                "Either provide 'send_to' in the response or set a default_target."
                            )
                        send_to = effective_default
                    
                    # If missing "content" → convert whole object to JSON string
                    if content is None:
                        content = json.dumps(parsed)
                    
                    return {
                        "send_to": send_to,
                        "content": content
                    }
                else:
                    # Invalid JSON structure, wrap it
                    if effective_default is None:
                        raise ValueError(
                            "Response is not a valid dict and no default_target is set. "
                            "Either return a dict with 'send_to' and 'content' or set a default_target."
                        )
                    return {
                        "send_to": effective_default,
                        "content": result
                    }
            except json.JSONDecodeError:
                # Step 3: Try to extract partial JSON (handles truncated responses)
                partial = extract_partial_json_content(cleaned_result)
                if partial:
                    return partial
                
                # Step 4: Not JSON, wrap it
                if effective_default is None:
                    raise ValueError(
                        "Response is not valid JSON and no default_target is set. "
                        "Either return a dict with 'send_to' and 'content' or set a default_target."
                    )
                return {
                    "send_to": effective_default,
                    "content": result
                }
        
        # Fallback: convert to string and wrap
        if effective_default is None:
            raise ValueError(
                f"Response type '{type(result).__name__}' is not supported and no default_target is set. "
                "Either return a dict with 'send_to' and 'content' or set a default_target."
            )
        return {
            "send_to": effective_default,
            "content": str(result)
        }
    
    def __repr__(self) -> str:
        """String representation."""
        return f"Agent(name='{self.name}', memory_messages={len(self.memory.messages)})"
