from __future__ import annotations

import asyncio
from typing import Optional

from agents import Agent, Runner, RunConfig, set_tracing_disabled

from diona.project.agent.model import BaseModelProvider
from diona.project.agent.sub_agents.registry import get_sub_agent_registry
from diona.project.agent.handoffs.registry import get_handoff_registry
from diona.project.agent.core.prompt import PromptFactory

set_tracing_disabled(disabled=False)


class CoreAgent:
    def __init__(
        self,
        model_provider: Optional[BaseModelProvider] = None,
        instructions: Optional[str] = None,
    ):
        self._model_provider = model_provider or BaseModelProvider.get_default()
        self._agent: Optional[Agent] = None
        self._instructions = instructions or PromptFactory.get_default_instructions()
        self._init_agent()

    def _init_agent(self) -> None:
        tools = []
        tools.extend(self._register_tools())
        tools.extend(self._register_sub_agents_as_tools())

        handoffs = self._register_handoff_agents()

        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 _register_tools(self) -> list:
        tools = []

        try:
            from diona.project.agent.tools.registry import get_tool_registry
            tool_registry = get_tool_registry()
            tools.extend(tool_registry.get_tools())
        except Exception as e:
            print(f"Warning: Failed to load tools from registry: {e}")

        return tools

    def _register_sub_agents_as_tools(self) -> list:
        tools = []

        try:
            sub_registry = get_sub_agent_registry()
            for agent_info in sub_registry.list_all():
                try:
                    sub_agent = agent_info.create_func(self._model_provider)
                    tool_agent = sub_agent.as_tool(
                        tool_name=f"{agent_info.name}_agent",
                        tool_description=agent_info.description
                    )
                    tools.append(tool_agent)
                except Exception as e:
                    print(f"Warning: Failed to create sub-agent {agent_info.name}: {e}")
        except Exception as e:
            print(f"Warning: Failed to load sub-agents: {e}")

        return tools

    def _register_handoff_agents(self) -> list:
        handoffs = []

        try:
            handoff_registry = get_handoff_registry()
            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}")

        return handoffs

    async def run(self, user_input: str) -> str:
        if not self._agent:
            return "Agent not initialized"

        run_config = RunConfig(
            model_provider=self._model_provider,
            max_turns=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:
        async def _run():
            return await self.run(user_input)

        try:
            loop = asyncio.get_running_loop()
            try:
                import nest_asyncio
                nest_asyncio.apply()
            except ImportError:
                pass
            return loop.run_until_complete(_run())
        except RuntimeError:
            return asyncio.run(_run())

    @property
    def inner_agent(self) -> Agent:
        return self._agent


def get_core_agent(
    model_provider: Optional[BaseModelProvider] = None,
    instructions: Optional[str] = None
) -> CoreAgent:
    return CoreAgent(model_provider=model_provider, instructions=instructions)
