from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from typing import Any, Optional

from agents import Agent, Runner, RunConfig, set_tracing_disabled

from diona.project.agent.model import BaseModelProvider

set_tracing_disabled(disabled=True)


class Middleware(ABC):
    def __init__(self, instructions: Optional[str] = None):
        self._model_provider_manager = None
        self._model_provider = None
        self._instructions = instructions or self._default_instructions()
        self._agent: Optional[Agent] = None
        self._metadata: dict = {}
        self._init_agent()

    @property
    def metadata(self) -> dict:
        return self._metadata

    @property
    @abstractmethod
    def name(self) -> str:
        pass

    @property
    def description(self) -> str:
        return "Middleware agent"

    def _default_instructions(self) -> str:
        return "You are a helpful AI assistant."

    def _init_agent(self) -> None:
        try:
            from diona.project.agent.model.manager import get_model_provider_manager
            self._model_provider_manager = get_model_provider_manager()
            self._model_provider = self._model_provider_manager.get_current()
        except Exception:
            self._model_provider = BaseModelProvider.get_default()

        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=self.name,
            instructions=self._instructions,
            tools=[],
            model=model,
        )

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

        run_config = RunConfig(
            model_provider=self._model_provider
        )
        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_agent(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())

    async def execute(self, context: Any) -> tuple[bool, str]:
        json_input = context.to_json_input(self.name)
        output = await self._run_agent(json_input)

        context.add_result(self.name, json_input, output)

        return True, output

    def execute_sync(self, context: Any) -> tuple[bool, str]:
        async def _execute():
            return await self.execute(context)

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