import json
import json
from typing import List, Optional, Any
from .llm import LLMProvider
from .memory import BaseMemory, InMemoryMemory
from .tools import Tool

class Agent:
    """Core orchestrator for an apexagents Agent."""
    
    def __init__(
        self,
        name: str = "Assistant",
        model: str = "gpt-4o-mini",
        instructions: str = "You are a helpful AI assistant.",
        tools: Optional[List[Tool]] = None,
        memory: Optional[BaseMemory] = None,
        **model_kwargs
    ):
        """
        Task: Initializes an Agent instance configuring name, instruction prompts, tool suite, memory backend, and LLM connection.
        
        Input Parameters:
            name (str): Friendly name for the agent (defaults to "Assistant").
            model (str): Name of the LLM to use via litellm (defaults to "gpt-4o-mini").
            instructions (str): The foundational system prompt dictating the agent's behavior.
            tools (Optional[List[Tool]]): List of Tool objects available to the agent for external computation.
            memory (Optional[BaseMemory]): State persistence manager for chat history. Uses InMemoryMemory by default.
            **model_kwargs: Extra LLM configuration arguments passed to LLMProvider.
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.name = name
        self.instructions = instructions
        self.tools = tools or []
        self.tools_map = {tool.name: tool for tool in self.tools}
        self.memory = memory or InMemoryMemory()
        self.llm = LLMProvider(model_name=model, **model_kwargs)
        
        # Initialize memory with system instructions if empty
        if not self.memory.get_messages():
            self.memory.add_message("system", self.instructions)
            
    async def arun(self, user_input: str, response_model: Optional[Any] = None) -> Any:
        """
        Task: Runs the agent loop asynchronously given a new user input string. Executes tools dynamically and returns standard or enforced structural output.
        
        Input Parameters:
            user_input (str): The initial prompt or instruction provided by the user.
            response_model (Optional[Any]): A Pydantic BaseModel class for strictly formatted structured JSON outputs.
            
        Output Parameters / Returns:
            Any: A standard string response, or an instantiated Pydantic model if response_model was specified.
            
        Raises:
            Exception: Bubbles up any unexpected exceptions from the LLMProvider generating process.
        """
        self.memory.add_message("user", user_input)
        
        while True:
            messages = self.memory.get_messages()
            response = await self.llm.agenerate(
                messages=messages, 
                tools=self.tools, 
                response_model=response_model
            )
            
            message = response.choices[0].message
            
            # If the model called tools, handle them
            if hasattr(message, "tool_calls") and message.tool_calls:
                
                # We need to construct a dictionary for the memory that works cleanly with litellm structure
                assistant_msg = {
                    "role": "assistant",
                    "content": message.content,
                    "tool_calls": []
                }
                
                # Reconstruct tool_calls for history
                for tc in message.tool_calls:
                    # litellm objects mimic openai pydantic objects
                    tc_dict = {
                        "id": getattr(tc, "id", None),
                        "type": getattr(tc, "type", "function"),
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments
                        }
                    }
                    assistant_msg["tool_calls"].append(tc_dict)
                    
                self.memory.add_message(
                    role="assistant", 
                    content=message.content, 
                    tool_calls=assistant_msg["tool_calls"]
                )
                
                # Execute each tool call
                for tool_call in message.tool_calls:
                    function_name = tool_call.function.name
                    function_args = json.loads(tool_call.function.arguments)
                    
                    if function_name in self.tools_map:
                        tool_instance = self.tools_map[function_name]
                        try:
                            result = tool_instance.execute(**function_args)
                            import inspect
                            if inspect.isawaitable(result):
                                result = await result
                        except Exception as e:
                            result = f"Error executing tool: {e}"
                    else:
                        result = f"Error: Tool {function_name} not found."
                        
                    self.memory.add_message(
                        role="tool",
                        content=str(result),
                        tool_call_id=tool_call.id,
                        name=function_name
                    )
            else:
                # No more tool calls, final response
                content = message.content
                self.memory.add_message("assistant", str(content))
                
                # If a pydantic model was requested, parse the JSON content into the object
                if response_model and isinstance(content, str):
                    try:
                        return response_model.model_validate_json(content)
                    except Exception as e:
                        # Fallback parsing if litellm doesn't format it perfectly as JSON string
                        try:
                            # litellm sometimes returns raw dicts if response_format is used internally
                            data = json.loads(content)
                            return response_model(**data)
                        except:
                            return content # Return raw string if parsing outright fails
                return content

    def as_tool(self) -> Tool:
        """
        Task: Converts the agent instance into a Tool, enabling Agent Swarms and Task Delegation scenarios by nesting inside another agent's tools array.
        
        Input Parameters: None
            
        Output Parameters / Returns:
            Tool: A callable Tool object containing a synchronized delegation wrapper over this agent's arun.
            
        Raises: None
        """
        import asyncio
        import inspect

        def delegate_task(task_description: str) -> str:
            # We must gracefully run the async arun inside a sync tool executor
            try:
                loop = asyncio.get_running_loop()
                # If we are already in an event loop, we create a task and wait for it
                return loop.create_task(self.arun(task_description))
            except RuntimeError:
                # If no event loop is running (unlikely if the parent is arun), use asyncio.run
                return asyncio.run(self.arun(task_description))
                
        # Manually alter the docstring and name of the generated function so the LLM understands it
        delegate_task.__name__ = f"ask_{self.name.lower().replace(' ', '_')}"
        delegate_task.__doc__ = f"Delegates a complex sub-task to the {self.name} agent. Use this for: {self.instructions}"
        
        return Tool(delegate_task)
