"""Core agent implementation using OpenAI Agents SDK"""

from __future__ import annotations

import asyncio
from typing import Optional, Callable

from agents import (
    Agent,
    Runner,
    function_tool,
    RunConfig,
    set_tracing_disabled,
)

from diona.project.agent.model import BaseModelProvider
from diona.project.agent.sub_registry import get_sub_agent_registry, SubAgentRegistry


set_tracing_disabled(disabled=True)


class CoreAgent:
    """Core agent that handles user input and delegates to sub-agents"""
    
    def __init__(
        self,
        model_provider: Optional[BaseModelProvider] = None,
        instructions: Optional[str] = None,
    ):
        """Initialize core agent
        
        Args:
            model_provider: Model provider for the agent
            instructions: Custom instructions for the core agent
        """
        self._model_provider = model_provider or BaseModelProvider.get_default()
        self._registry = get_sub_agent_registry()
        self._agent: Optional[Agent] = None
        self._instructions = instructions or self._default_instructions()
        self._init_agent()
    
    def _default_instructions(self) -> str:
        """Generate default instructions for the core agent"""
        sub_agents = self._registry.get_descriptions()
        
        sub_agent_instructions = ""
        if sub_agents:
            sub_agent_instructions = "\n\nAvailable sub-agents:\n"
            for agent_info in sub_agents:
                sub_agent_instructions += f"- {agent_info['name']}: {agent_info['description']}\n"
        
        return f"""You are a helpful AI assistant. You can:
1. Answer questions directly
2. Delegate tasks to sub-agents when appropriate{sub_agent_instructions}

When a user asks you to perform a task that matches a sub-agent's description, you should use the appropriate sub-agent.
If the user explicitly asks to use a specific sub-agent (e.g., "/sub agent_name"), use the sub_agent_tool to call that agent.
"""
    
    def _init_agent(self) -> None:
        """Initialize the core agent with sub-agents"""
        sub_agents = self._registry.list_all()
        
        tools = [self._create_sub_agent_tool(agent) for agent in sub_agents]
        
        # Add global memory tool
        try:
            from diona.project.agent.tools import global_memory_tool
            tools.append(global_memory_tool)
        except Exception as e:
            print(f"Warning: Failed to add global memory tool: {e}")
        
        # Add SQLite tool
        try:
            from diona.project.agent.tools import sqlite_tool
            tools.append(sqlite_tool)
        except Exception as e:
            print(f"Warning: Failed to add SQLite tool: {e}")
        
        # Add file storage tools
        try:
            from diona.project.agent.tools import write_file, read_file
            tools.append(write_file)
            tools.append(read_file)
        except Exception as e:
            print(f"Warning: Failed to add file storage tools: {e}")
        
        # Add handoff agents
        handoffs = []
        try:
            from diona.project.agent import get_handoff_agent_registry
            handoff_registry = get_handoff_agent_registry()
            
            # Get all registered handoff agents
            for agent_info in handoff_registry.list_all():
                try:
                    agent = agent_info.create_func(self._model_provider)
                    handoffs.append(agent)
                except Exception as e:
                    print(f"Warning: Failed to create handoff agent {agent_info.name}: {e}")
        except Exception as e:
            print(f"Warning: Failed to load handoff agents: {e}")
        
        try:
            model = self._model_provider.get_model(None)
        except Exception as e:
            raise ValueError(f"Failed to get model: {e}. Check API key and base URL configuration.")
        
        self._agent = Agent(
            name="CoreAgent",
            instructions=self._instructions,
            tools=tools,
            handoffs=handoffs,
            model=model,
        )
    
    def _create_sub_agent_tool(self, agent_info) -> Callable:
        """Create a function tool for a sub-agent
        
        Args:
            agent_info: SubAgentInfo object
            
        Returns:
            Function tool for the sub-agent
        """
        tool_description = f"Call {agent_info.name}: {agent_info.description}"
        
        @function_tool(description_override=tool_description)
        def sub_agent_tool(input_text: str) -> str:
            """Execute the sub-agent with user input
            
            Args:
                input_text: Input for the sub-agent
                
            Returns:
                Result from the sub-agent
            """
            return agent_info.callable(input_text)
        
        return sub_agent_tool
    
    def register_sub_agent(
        self,
        name: str,
        description: str,
        instructions: str,
        callable: Callable[[str], str],
        metadata: Optional[dict] = None
    ) -> None:
        """Register a new sub-agent and recreate the core agent
        
        Args:
            name: Agent name
            description: Brief description
            instructions: Detailed instructions
            callable: Function to execute
            metadata: Additional metadata
        """
        self._registry.register(name, description, instructions, callable, metadata)
        self._init_agent()
    
    def reload_sub_agents(self) -> None:
        """Reload sub-agents and recreate the core agent"""
        self._init_agent()
    
    async def run(self, user_input: str) -> str:
        """Run the core agent with user input
        
        Args:
            user_input: User input
            
        Returns:
            Agent response
        """
        if not self._agent:
            return "Agent not initialized"
        
        # Create run config with increased max turns limit
        run_config = RunConfig(
            model_provider=self._model_provider,
            max_turns=100  # Increase from default 10 to 100
        )
        result = await Runner.run(self._agent, user_input, run_config=run_config)
        return result.final_output
    
    def run_sync(self, user_input: str) -> str:
        """Run the core agent synchronously
        
        Args:
            user_input: User input
            
        Returns:
            Agent response
        """
        return asyncio.run(self.run(user_input))
    
    @property
    def agent(self) -> Agent:
        """Get the underlying agent"""
        return self._agent


def get_core_agent(
    model_provider: Optional[BaseModelProvider] = None,
    instructions: Optional[str] = None
) -> CoreAgent:
    """Get or create a core agent instance
    
    Args:
        model_provider: Model provider
        instructions: Custom instructions
        
    Returns:
        CoreAgent instance
    """
    return CoreAgent(model_provider=model_provider, instructions=instructions)
