from abc import ABC, abstractmethod
from collections.abc import Generator, Mapping
from typing import Any, Self, final

from typing_extensions import deprecated
from werkzeug import Request

from dify_plugin.core.runtime import Session
from dify_plugin.entities import ParameterOption
from dify_plugin.entities.invoke_message import InvokeMessage
from dify_plugin.entities.oauth import ToolOAuthCredentials
from dify_plugin.entities.provider_config import LogMetadata
from dify_plugin.entities.tool import (
    ToolInvokeMessage,
    ToolParameter,
    ToolRuntime,
    ToolSelector,
)
from dify_plugin.file.constants import DIFY_FILE_IDENTITY, DIFY_TOOL_SELECTOR_IDENTITY
from dify_plugin.file.entities import FileType
from dify_plugin.file.file import File
from dify_plugin.protocol.oauth import OAuthCredentials


class ToolLike[T: InvokeMessage](ABC):
    response_type: type[T]

    ############################################################
    #            For plugin implementation use only            #
    ############################################################

    def create_text_message(self, text: str) -> T:
        return self.response_type(
            type=InvokeMessage.MessageType.TEXT,
            message=InvokeMessage.TextMessage(text=text),
        )

    def create_json_message(self, json: Mapping | list) -> T:
        return self.response_type(
            type=InvokeMessage.MessageType.JSON,
            message=InvokeMessage.JsonMessage(json_object=json),
        )

    def create_image_message(self, image_url: str) -> T:
        """
        create an image message

        :param image: the url of the image
        :return: the image message

        Returns:
            The return value.
        """
        return self.response_type(
            type=InvokeMessage.MessageType.IMAGE,
            message=InvokeMessage.TextMessage(text=image_url),
        )

    def create_link_message(self, link: str) -> T:
        """
        create a link message

        :param link: the url of the link
        :return: the link message

        Returns:
            The return value.
        """
        return self.response_type(
            type=InvokeMessage.MessageType.LINK,
            message=InvokeMessage.TextMessage(text=link),
        )

    def create_blob_message(self, blob: bytes, meta: dict | None = None) -> T:
        """
        create a blob message

        :param blob: the blob
        :return: the blob message

        Returns:
            The return value.
        """
        return self.response_type(
            type=InvokeMessage.MessageType.BLOB,
            message=InvokeMessage.BlobMessage(blob=blob),
            meta=meta,
        )

    def create_variable_message(self, variable_name: str, variable_value: object) -> T:
        """
        create a variable message

        :param variable_name: the name of the variable
        :param variable_value: the value of the variable
        :return: the variable message

        Returns:
            The return value.
        """
        return self.response_type(
            type=InvokeMessage.MessageType.VARIABLE,
            message=InvokeMessage.VariableMessage(
                variable_name=variable_name, variable_value=variable_value
            ),
        )

    def create_stream_variable_message(
        self, variable_name: str, variable_value: str
    ) -> T:
        """
        create a variable message that will be streamed to the frontend

        NOTE: variable value should be a string, only string is streaming supported now

        :param variable_name: the name of the variable
        :param variable_value: the value of the variable
        :return: the variable message

        Returns:
            The return value.
        """
        return self.response_type(
            type=InvokeMessage.MessageType.VARIABLE,
            message=InvokeMessage.VariableMessage(
                variable_name=variable_name,
                variable_value=variable_value,
                stream=True,
            ),
        )

    def create_log_message(
        self,
        label: str,
        data: Mapping[str, Any],
        status: InvokeMessage.LogMessage.LogStatus = (
            InvokeMessage.LogMessage.LogStatus.SUCCESS
        ),
        parent: T | None = None,
        metadata: Mapping[LogMetadata, Any] | None = None,
    ) -> T:
        """
        create a log message with status "start"
        """
        return self.response_type(
            type=InvokeMessage.MessageType.LOG,
            message=InvokeMessage.LogMessage(
                label=label,
                data=data,
                status=status,
                parent_id=parent.message.id
                if parent and isinstance(parent.message, InvokeMessage.LogMessage)
                else None,
                metadata=metadata,
            ),
        )

    def create_retriever_resource_message(
        self,
        retriever_resources: list[
            InvokeMessage.RetrieverResourceMessage.RetrieverResource
        ],
        context: str,
    ) -> T:
        """
        create a retriever resource message
        """
        return self.response_type(
            type=InvokeMessage.MessageType.RETRIEVER_RESOURCES,
            message=InvokeMessage.RetrieverResourceMessage(
                retriever_resources=retriever_resources,
                context=context,
            ),
        )

    def finish_log_message(
        self,
        log: T,
        status: InvokeMessage.LogMessage.LogStatus = (
            InvokeMessage.LogMessage.LogStatus.SUCCESS
        ),
        error: str | None = None,
        data: Mapping[str, Any] | None = None,
        metadata: Mapping[LogMetadata, Any] | None = None,
    ) -> T:
        """
        mark log as finished
        """
        if not isinstance(log.message, InvokeMessage.LogMessage):
            msg = "Log message payload is required"
            raise TypeError(msg)
        return self.response_type(
            type=InvokeMessage.MessageType.LOG,
            message=InvokeMessage.LogMessage(
                id=log.message.id,
                label=log.message.label,
                data=data or log.message.data,
                status=status,
                parent_id=log.message.parent_id,
                error=error,
                metadata=metadata or log.message.metadata,
            ),
        )

    @deprecated(
        "This feature is deprecated, will soon be replaced by dynamic select parameter"
    )
    def _get_runtime_parameters(self) -> list[ToolParameter]:
        """
        get the runtime parameters of the tool

        :return: the runtime parameters

        Returns:
            The return value.
        """
        return []

    @classmethod
    def has_runtime_parameters(cls) -> bool:
        """
        check if the _get_runtime_parameters method is overridden by subclass

        :return: True if overridden, False otherwise

        Returns:
            The return value.
        """
        return "_get_runtime_parameters" in cls.__dict__

    @classmethod
    def _is_get_runtime_parameters_overridden(cls) -> bool:
        """
        check if the _get_runtime_parameters method is overridden by subclass

        :return: True if overridden, False otherwise

        Returns:
            The return value.
        """
        return cls.has_runtime_parameters()

    @classmethod
    def _convert_parameters(cls, tool_parameters: dict) -> dict:
        """
        convert parameters into correct types
        """
        for parameter, value in tool_parameters.items():
            if (
                isinstance(value, dict)
                and value.get("dify_model_identity") == DIFY_FILE_IDENTITY
            ):
                tool_parameters[parameter] = File(
                    url=value["url"],
                    mime_type=value.get("mime_type"),
                    type=FileType(value.get("type")),
                    filename=value.get("filename"),
                    extension=value.get("extension"),
                    size=value.get("size"),
                )
            elif isinstance(value, list) and all(
                isinstance(item, dict)
                and item.get("dify_model_identity") == DIFY_FILE_IDENTITY
                for item in value
            ):
                tool_parameters[parameter] = [
                    File(
                        url=item["url"],
                        mime_type=item.get("mime_type"),
                        type=FileType(item.get("type")),
                        filename=item.get("filename"),
                        extension=item.get("extension"),
                        size=item.get("size"),
                    )
                    for item in value
                ]
            elif (
                isinstance(value, dict)
                and value.get("dify_model_identity") == DIFY_TOOL_SELECTOR_IDENTITY
            ):
                tool_parameters[parameter] = ToolSelector.model_validate(value)
            elif isinstance(value, list) and all(
                isinstance(item, dict)
                and item.get("dify_model_identity") == DIFY_TOOL_SELECTOR_IDENTITY
                for item in value
            ):
                tool_parameters[parameter] = [
                    ToolSelector.model_validate(item) for item in value
                ]

        return tool_parameters


class ToolProvider:
    def validate_credentials(self, credentials: dict) -> None:
        return self._validate_credentials(credentials)

    def _validate_credentials(self, credentials: dict) -> None:
        msg = (
            "The tool you are using does not support credentials validation, "
            "please implement `_validate_credentials` method"
        )
        raise NotImplementedError(msg)

    def oauth_get_authorization_url(
        self, redirect_uri: str, system_credentials: Mapping[str, Any]
    ) -> str:
        """
        Get the authorization url

        :param redirect_uri: redirect uri
        :param system_credentials: system credentials
        :return: authorization url

        Returns:
            The return value.
        """
        return self._oauth_get_authorization_url(redirect_uri, system_credentials)

    def _oauth_get_authorization_url(
        self, redirect_uri: str, system_credentials: Mapping[str, Any]
    ) -> str:
        msg = (
            "The tool you are using does not support OAuth, please implement "
            "`_oauth_get_authorization_url` method"
        )
        raise NotImplementedError(msg)

    def oauth_get_credentials(
        self, redirect_uri: str, system_credentials: Mapping[str, Any], request: Request
    ) -> OAuthCredentials:
        """
        Get the credentials

        :param redirect_uri: redirect uri
        :param system_credentials: system credentials
        :param request: raw http request
        :return: credentials

        Returns:
            The return value.
        """
        tool_oauth_credentials = self._oauth_get_credentials(
            redirect_uri, system_credentials, request
        )
        return OAuthCredentials(
            credentials=tool_oauth_credentials.credentials,
            expires_at=tool_oauth_credentials.expires_at,
        )

    def _oauth_get_credentials(
        self, redirect_uri: str, system_credentials: Mapping[str, Any], request: Request
    ) -> ToolOAuthCredentials:
        msg = (
            "The tool you are using does not support OAuth, please implement "
            "`_oauth_get_credentials` method"
        )
        raise NotImplementedError(msg)

    def oauth_refresh_credentials(
        self,
        redirect_uri: str,
        system_credentials: Mapping[str, Any],
        credentials: Mapping[str, Any],
    ) -> OAuthCredentials:
        """
        Refresh the credentials

        :param redirect_uri: redirect uri
        :param system_credentials: system credentials
        :param credentials: credentials
        :return: refreshed credentials

        Returns:
            The return value.
        """
        tool_oauth_credentials = self._oauth_refresh_credentials(
            redirect_uri, system_credentials, credentials
        )
        return OAuthCredentials(
            credentials=tool_oauth_credentials.credentials,
            expires_at=tool_oauth_credentials.expires_at,
        )

    def _oauth_refresh_credentials(
        self,
        redirect_uri: str,
        system_credentials: Mapping[str, Any],
        credentials: Mapping[str, Any],
    ) -> ToolOAuthCredentials:
        msg = (
            "The tool you are using does not support OAuth, please implement "
            "`_oauth_refresh_credentials` method"
        )
        raise NotImplementedError(msg)


class Tool(ToolLike[ToolInvokeMessage]):
    runtime: ToolRuntime
    session: Session

    @final
    def __init__(
        self,
        runtime: ToolRuntime,
        session: Session,
    ) -> None:
        """
        Initialize the tool

        NOTE:
        - This method has been marked as final, DO NOT OVERRIDE IT.
        """
        self.runtime = runtime
        self.session = session
        self.response_type = ToolInvokeMessage

    @classmethod
    def from_credentials(
        cls,
        credentials: dict,
        user_id: str | None = None,
    ) -> Self:
        return cls(
            runtime=ToolRuntime(
                credentials=credentials, user_id=user_id, session_id=None
            ),
            session=Session.empty_session(),  # Session is unavailable here.
        )

    ############################################################
    #        Methods that can be implemented by plugin         #
    ############################################################

    @abstractmethod
    def _invoke(
        self, tool_parameters: dict
    ) -> Generator[ToolInvokeMessage, None, None]:
        pass

    def _fetch_parameter_options(self, parameter: str) -> list[ParameterOption]:
        """
        Fetch the parameter options of the tool.

        To be implemented by subclasses.

        Also, it's optional to implement, that's why it's not an abstract method.
        """
        msg = (
            "This plugin should implement `_fetch_parameter_options` method "
            "to enable dynamic select parameter"
        )
        raise NotImplementedError(msg)

    ############################################################
    #                 For executor use only                    #
    ############################################################

    def invoke(self, tool_parameters: dict) -> Generator[ToolInvokeMessage, None, None]:
        # convert parameters into correct types
        tool_parameters = self._convert_parameters(tool_parameters)
        return self._invoke(tool_parameters)

    @deprecated(
        "This feature is deprecated, will soon be replaced by dynamic select parameter"
    )
    def get_runtime_parameters(self) -> list[ToolParameter]:
        return self._get_runtime_parameters()

    def fetch_parameter_options(self, parameter: str) -> list[ParameterOption]:
        """
        Fetch the parameter options of the tool.

        To be implemented by subclasses.

        Returns:
            The return value.
        """
        return self._fetch_parameter_options(parameter)
