from gevent import monkey

from dify_plugin.core.server.aws_stream import AWSLambdaStream
from dify_plugin.core.server.stdio_stream import StdioStream
from dify_plugin.core.server.tcp_stream import TCPStream
from dify_plugin.stream.io_stream import PluginIOStream

# patch all the blocking calls
monkey.patch_all(sys=True)

from collections.abc import Generator  # noqa: E402
import logging  # noqa: E402

from dify_plugin.config.config import DifyPluginEnv, InstallMethod  # noqa: E402

from dify_plugin.core.runtime.entities.plugin.request import (  # noqa: E402
    ModelActions,
    PluginInvokeType,
    ToolActions,
    WebhookActions,
)
from dify_plugin.core.runtime.session import Session  # noqa: E402
from dify_plugin.core.server.io_server import IOServer  # noqa: E402
from dify_plugin.core.server.router import Router  # noqa: E402
from dify_plugin.logger_format import plugin_logger_handler  # noqa: E402

from dify_plugin.plugin_executor import PluginExecutor  # noqa: E402
from dify_plugin.plugin_registration import PluginRegistration  # noqa: E402


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.addHandler(plugin_logger_handler)


class Plugin(IOServer, Router):
    def __init__(self, config: DifyPluginEnv) -> None:
        """
        Initialize plugin
        """
        # load plugin configuration
        self.registration = PluginRegistration(config)

        if config.INSTALL_METHOD == InstallMethod.Local:
            io_stream = self._launch_local_stream(config)
        elif config.INSTALL_METHOD == InstallMethod.Remote:
            io_stream = self._launch_remote_stream(config)
        elif config.INSTALL_METHOD == InstallMethod.AWSLambda:
            io_stream = self._launch_aws_stream(config)
        else:
            raise ValueError("Invalid install method")

        # initialize plugin executor
        self.plugin_executer = PluginExecutor(self.registration)

        IOServer.__init__(self, config, io_stream)
        Router.__init__(self, io_stream)

        # register io routes
        self._register_request_routes()

    def _launch_local_stream(self, config: DifyPluginEnv) -> PluginIOStream:
        """
        Launch local stream
        """
        stdio_stream = StdioStream()
        io_stream = PluginIOStream(stdio_stream, stdio_stream)

        io_stream.write(self.registration.configuration.model_dump_json() + "\n\n")

        self._log_configuration()
        return io_stream

    def _launch_remote_stream(self, config: DifyPluginEnv) -> PluginIOStream:
        """
        Launch remote stream
        """
        if not config.REMOTE_INSTALL_KEY:
            raise ValueError("Missing remote install key")

        tcp_stream = TCPStream(
            config.REMOTE_INSTALL_HOST,
            config.REMOTE_INSTALL_PORT,
            config.REMOTE_INSTALL_KEY,
            on_connected=lambda: tcp_stream.write(
                self.registration.configuration.model_dump_json() + "\n\n"
            )
            and self._log_configuration(),
        )

        io_stream = PluginIOStream(tcp_stream, tcp_stream)
        tcp_stream.launch()

        return io_stream

    def _launch_aws_stream(self, config: DifyPluginEnv) -> PluginIOStream:
        """
        Launch AWS stream
        """
        aws_stream = AWSLambdaStream(config.AWS_LAMBDA_PORT, config.MAX_REQUEST_TIMEOUT)
        io_stream = PluginIOStream(aws_stream, aws_stream)
        aws_stream.launch()

        return io_stream

    def _log_configuration(self):
        """
        Log plugin configuration
        """
        for tool in self.registration.tools_configuration:
            logger.info(f"Installed tool: {tool.identity.name}")
        for model in self.registration.models_configuration:
            logger.info(f"Installed model: {model.provider}")

    def _register_request_routes(self):
        """
        Register routes
        """
        self.register_route(
            self.plugin_executer.invoke_tool,
            lambda data: data.get("type") == PluginInvokeType.Tool.value
            and data.get("action") == ToolActions.InvokeTool.value,
        )

        self.register_route(
            self.plugin_executer.validate_tool_provider_credentials,
            lambda data: data.get("type") == PluginInvokeType.Tool.value
            and data.get("action") == ToolActions.ValidateCredentials.value,
        )

        self.register_route(
            self.plugin_executer.invoke_llm,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.InvokeLLM.value,
        )

        self.register_route(
            self.plugin_executer.invoke_text_embedding,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.InvokeTextEmbedding.value,
        )

        self.register_route(
            self.plugin_executer.invoke_rerank,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.InvokeRerank.value,
        )

        self.register_route(
            self.plugin_executer.invoke_tts,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.InvokeTTS.value,
        )

        self.register_route(
            self.plugin_executer.invoke_speech_to_text,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.InvokeSpeech2Text.value,
        )

        self.register_route(
            self.plugin_executer.invoke_moderation,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.InvokeModeration.value,
        )

        self.register_route(
            self.plugin_executer.validate_model_provider_credentials,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.ValidateProviderCredentials.value,
        )

        self.register_route(
            self.plugin_executer.validate_model_credentials,
            lambda data: data.get("type") == PluginInvokeType.Model.value
            and data.get("action") == ModelActions.ValidateModelCredentials.value,
        )

        self.register_route(
            self.plugin_executer.invoke_webhook,
            lambda data: data.get("type") == PluginInvokeType.Webhook.value
            and data.get("action") == WebhookActions.InvokeWebhook.value,
        )

    def _execute_request(self, session_id: str, data: dict):
        """
        accept requests and execute
        :param session_id: session id, unique for each request
        :param data: request data
        """

        session = Session(session_id=session_id, executor=self.executer)
        response = self.dispatch(session, data)
        if response:
            if isinstance(response, Generator):
                for message in response:
                    self.io_stream.session_message(
                        session_id=session_id,
                        data=self.io_stream.stream_object(data=message),
                    )
            else:
                self.io_stream.session_message(
                    session_id=session_id,
                    data=self.io_stream.stream_object(data=response),
                )