from __future__ import annotations

import asyncio
from typing import Callable, Optional, Any, AsyncIterator

from diona.project.agent.session import SessionManager
from diona.project.agent.model import get_model_provider_manager
from diona.project.agent.middleware import Middleware, MiddlewareContext, get_middleware_registry


class AgentFacade:
    _instance: Optional['AgentFacade'] = None

    def __new__(cls) -> 'AgentFacade':
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return

        self._session_manager = SessionManager()
        self._pre_middleware: list[Middleware] = []
        self._post_middleware: list[Middleware] = []
        self._model_provider_manager = get_model_provider_manager()
        self._init_middleware()
        self._initialized = True

    def _init_middleware(self):
        registry = get_middleware_registry()
        for info in registry.list_all():
            agent = info.create_func()
            agent.metadata.update(info.metadata)
            if agent.metadata.get('post_agent'):
                self._post_middleware.append(agent)
            else:
                self._pre_middleware.append(agent)

    def register_pre_middleware(self, middleware: Middleware) -> None:
        self._pre_middleware.append(middleware)

    def unregister_pre_middleware(self, middleware_name: str) -> bool:
        for i, m in enumerate(self._pre_middleware):
            if m.name == middleware_name:
                self._pre_middleware.pop(i)
                return True
        return False

    def register_post_middleware(self, middleware: Middleware) -> None:
        self._post_middleware.append(middleware)

    def unregister_post_middleware(self, middleware_name: str) -> bool:
        for i, m in enumerate(self._post_middleware):
            if m.name == middleware_name:
                self._post_middleware.pop(i)
                return True
        return False

    def get_pre_middleware(self) -> list[Middleware]:
        return self._pre_middleware.copy()

    def get_post_middleware(self) -> list[Middleware]:
        return self._post_middleware.copy()

    def register_pre_middleware_by_name(self, middleware_name: str) -> bool:
        registry = get_middleware_registry()
        info = registry.get(middleware_name)
        if info:
            middleware = info.create_func()
            self.register_pre_middleware(middleware)
            return True
        return False

    def register_post_middleware_by_name(self, middleware_name: str) -> bool:
        registry = get_middleware_registry()
        info = registry.get(middleware_name)
        if info:
            middleware = info.create_func()
            self.register_post_middleware(middleware)
            return True
        return False

    def register_all_pre_middleware_from_registry(self) -> int:
        registry = get_middleware_registry()
        count = 0
        for info in registry.list_all():
            middleware = info.create_func()
            self.register_pre_middleware(middleware)
            count += 1
        return count

    def register_all_post_middleware_from_registry(self) -> int:
        registry = get_middleware_registry()
        count = 0
        for info in registry.list_all():
            middleware = info.create_func()
            self.register_post_middleware(middleware)
            count += 1
        return count

    async def _run_pre_chain(self, context: MiddlewareContext) -> str:
        for middleware in self._pre_middleware:
            should_continue, output = await middleware.execute(context)

            if not should_continue:
                return f"[拒绝执行] {context.early_return_reason or output}"

            if context.is_early_return:
                return f"[拒绝执行] {context.early_return_reason}"

        return context.processed_input if context.processed_input else context.user_input

    async def _run_post_chain(self, context: MiddlewareContext) -> str:
        for middleware in self._post_middleware:
            should_continue, output = await middleware.execute(context)

            if not should_continue:
                return f"[后处理中断] {context.early_return_reason or output}"

            if context.is_early_return:
                return f"[后处理中断] {context.early_return_reason}"

        return context.current_result or ""

    def _run_pre_chain_sync(self, context: MiddlewareContext) -> str:
        async def _run():
            return await self._run_pre_chain(context)

        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())

    def _run_post_chain_sync(self, context: MiddlewareContext) -> str:
        async def _run():
            return await self._run_post_chain(context)

        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 run_async(
        self,
        user_input: str,
        session_id: Optional[str] = None,
        use_session: bool = True
    ) -> str:
        context = MiddlewareContext(
            user_input=user_input,
            session_id=session_id,
            use_session=use_session
        )

        processed_input = await self._run_pre_chain(context)

        if context.is_early_return:
            return processed_input

        model_provider = self._model_provider_manager.get_current()

        if use_session:
            if session_id:
                self._session_manager.set_current_session(session_id)
            core_result = self._session_manager.run_with_session(processed_input, model_provider)
        else:
            core_result = self._session_manager.run_without_session(processed_input, model_provider)

        context.current_result = core_result

        final_result = await self._run_post_chain(context)

        return final_result

    def run(
        self,
        user_input: str,
        session_id: Optional[str] = None,
        use_session: bool = True
    ) -> str:
        try:
            loop = asyncio.get_running_loop()
            try:
                import nest_asyncio
                nest_asyncio.apply()
            except ImportError:
                pass
            return loop.run_until_complete(self.run_async(user_input, session_id, use_session))
        except RuntimeError:
            return asyncio.run(self.run_async(user_input, session_id, use_session))

    async def run_streamed(
        self,
        user_input: str,
        session_id: Optional[str] = None,
        use_session: bool = True,
        on_token: Optional[Callable[[str], None]] = None
    ) -> AsyncIterator[str]:
        context = MiddlewareContext(
            user_input=user_input,
            session_id=session_id,
            use_session=use_session
        )

        processed_input = await self._run_pre_chain(context)

        if context.is_early_return:
            yield processed_input
            return

        model_provider = self._model_provider_manager.get_current()

        full_response = ""

        if use_session:
            if session_id:
                self._session_manager.set_current_session(session_id)
            async for token in self._session_manager.run_with_session_streamed(
                processed_input, model_provider, on_token
            ):
                full_response += token
                yield token
        else:
            async for token in self._session_manager.run_without_session_streamed(
                processed_input, model_provider, on_token
            ):
                full_response += token
                yield token

        context.current_result = full_response

        await self._run_post_chain(context)

    def create_session(self, name: Optional[str] = None) -> str:
        return self._session_manager.create_session(name)

    def get_session(self, session_id: str) -> Optional[Any]:
        return self._session_manager.get_session(session_id)

    def list_sessions(self) -> list[Any]:
        return self._session_manager.list_sessions()

    def delete_session(self, session_id: str) -> bool:
        return self._session_manager.delete_session(session_id)

    def set_current_session(self, session_id: str) -> bool:
        return self._session_manager.set_current_session(session_id)

    def get_current_session_id(self) -> Optional[str]:
        return self._session_manager.get_current_session_id()

    def clear_all_sessions(self) -> None:
        return self._session_manager.clear_all_sessions()

    async def get_session_items(self, session_id: Optional[str] = None, limit: Optional[int] = None) -> list:
        return await self._session_manager.get_session_items(session_id, limit)

    def get_session_items_sync(self, session_id: Optional[str] = None, limit: Optional[int] = None) -> list:
        return self._session_manager.get_session_items_sync(session_id, limit)


_facade: Optional[AgentFacade] = None


def get_agent_facade() -> AgentFacade:
    global _facade
    if _facade is None:
        _facade = AgentFacade()
    return _facade
