from __future__ import annotations

import json
from datetime import datetime
from enum import StrEnum
from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse

from pydantic import BaseModel

from configs import dify_config
from core.entities.provider_entities import BasicProviderConfig
from core.helper import encrypter
from core.helper.provider_cache import NoOpProviderCredentialCache
from core.mcp.types import OAuthClientInformation, OAuthClientMetadata, OAuthTokens
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_entities import ToolProviderType
from graphon.file import helpers as file_helpers

if TYPE_CHECKING:
    from models.tools import MCPToolProvider

# Constants
CLIENT_NAME = "Dify"
CLIENT_URI = "https://github.com/langgenius/dify"
DEFAULT_TOKEN_TYPE = "Bearer"
DEFAULT_EXPIRES_IN = 3600
MASK_CHAR = "*"
MIN_UNMASK_LENGTH = 6


class MCPSupportGrantType(StrEnum):
    """The supported grant types for MCP"""

    AUTHORIZATION_CODE = "authorization_code"
    CLIENT_CREDENTIALS = "client_credentials"
    REFRESH_TOKEN = "refresh_token"


class MCPAuthentication(BaseModel):
    client_id: str
    client_secret: str | None = None


class MCPConfiguration(BaseModel):
    timeout: float = 30
    sse_read_timeout: float = 300


class MCPProviderEntity(BaseModel):
    """MCP Provider domain entity for business logic operations"""

    # Basic identification
    id: str
    provider_id: str  # server_identifier
    name: str
    tenant_id: str
    user_id: str

    # Server connection info
    server_url: str  # encrypted URL
    headers: dict[str, str]  # encrypted headers
    timeout: float
    sse_read_timeout: float

    # Authentication related
    authed: bool
    credentials: dict[str, Any]  # encrypted credentials
    code_verifier: str | None = None  # for OAuth

    # Tools and display info
    tools: list[dict[str, Any]]  # parsed tools list
    icon: str | dict[str, str]  # parsed icon

    # Timestamps
    created_at: datetime
    updated_at: datetime

    @classmethod
    def from_db_model(cls, db_provider: MCPToolProvider) -> MCPProviderEntity:
        """Create entity from database model with decryption"""

        return cls(
            id=db_provider.id,
            provider_id=db_provider.server_identifier,
            name=db_provider.name,
            tenant_id=db_provider.tenant_id,
            user_id=db_provider.user_id,
            server_url=db_provider.server_url,
            headers=db_provider.headers,
            timeout=db_provider.timeout,
            sse_read_timeout=db_provider.sse_read_timeout,
            authed=db_provider.authed,
            credentials=db_provider.credentials,
            tools=db_provider.tool_dict,
            icon=db_provider.icon or "",
            created_at=db_provider.created_at,
            updated_at=db_provider.updated_at,
        )

    @property
    def redirect_url(self) -> str:
        """OAuth redirect URL"""
        return dify_config.CONSOLE_API_URL + "/console/api/mcp/oauth/callback"

    @property
    def client_metadata(self) -> OAuthClientMetadata:
        """Metadata about this OAuth client."""
        # Get grant type from credentials
        credentials = self.decrypt_credentials()

        # Try to get grant_type from different locations
        grant_type = credentials.get("grant_type", MCPSupportGrantType.AUTHORIZATION_CODE)

        # For nested structure, check if client_information has grant_types
        if "client_information" in credentials and isinstance(credentials["client_information"], dict):
            client_info = credentials["client_information"]
            # If grant_types is specified in client_information, use it to determine grant_type
            if "grant_types" in client_info and isinstance(client_info["grant_types"], list):
                if "client_credentials" in client_info["grant_types"]:
                    grant_type = MCPSupportGrantType.CLIENT_CREDENTIALS
                elif "authorization_code" in client_info["grant_types"]:
                    grant_type = MCPSupportGrantType.AUTHORIZATION_CODE

        # Configure based on grant type
        is_client_credentials = grant_type == MCPSupportGrantType.CLIENT_CREDENTIALS

        grant_types = ["refresh_token"]
        grant_types.append("client_credentials" if is_client_credentials else "authorization_code")

        response_types = [] if is_client_credentials else ["code"]
        redirect_uris = [] if is_client_credentials else [self.redirect_url]

        return OAuthClientMetadata(
            redirect_uris=redirect_uris,
            token_endpoint_auth_method="none",
            grant_types=grant_types,
            response_types=response_types,
            client_name=CLIENT_NAME,
            client_uri=CLIENT_URI,
        )

    @property
    def provider_icon(self) -> dict[str, str] | str:
        """Get provider icon, handling both dict and string formats"""
        if isinstance(self.icon, dict):
            return self.icon
        try:
            return json.loads(self.icon)
        except (json.JSONDecodeError, TypeError):
            # If not JSON, assume it's a file path
            return file_helpers.get_signed_file_url(self.icon)

    def to_api_response(self, user_name: str | None = None, include_sensitive: bool = True) -> dict[str, Any]:
        """Convert to API response format

        Args:
            user_name: User name to display
            include_sensitive: If False, skip expensive decryption operations (for list view optimization)
        """
        response = {
            "id": self.id,
            "author": user_name or "Anonymous",
            "name": self.name,
            "icon": self.provider_icon,
            "type": ToolProviderType.MCP.value,
            "is_team_authorization": self.authed,
            "server_url": self.masked_server_url(),
            "server_identifier": self.provider_id,
            "updated_at": int(self.updated_at.timestamp()),
            "label": I18nObject(en_US=self.name, zh_Hans=self.name).to_dict(),
            "description": I18nObject(en_US="", zh_Hans="").to_dict(),
        }

        # Add configuration
        response["configuration"] = {
            "timeout": str(self.timeout),
            "sse_read_timeout": str(self.sse_read_timeout),
        }

        # Skip expensive operations when sensitive data is not needed (e.g., list view)
        if not include_sensitive:
            response["masked_headers"] = {}
            response["is_dynamic_registration"] = True
        else:
            # Add masked headers
            response["masked_headers"] = self.masked_headers()

            # Add authentication info if available
            masked_creds = self.masked_credentials()
            if masked_creds:
                response["authentication"] = masked_creds
            response["is_dynamic_registration"] = self.credentials.get("client_information", {}).get(
                "is_dynamic_registration", True
            )

        return response

    def retrieve_client_information(self) -> OAuthClientInformation | None:
        """OAuth client information if available"""
        credentials = self.decrypt_credentials()
        if not credentials:
            return None

        # Check if we have nested client_information structure
        if "client_information" not in credentials:
            return None
        client_info_data = credentials["client_information"]
        if isinstance(client_info_data, dict):
            if "encrypted_client_secret" in client_info_data:
                client_info_data["client_secret"] = encrypter.decrypt_token(
                    self.tenant_id, client_info_data["encrypted_client_secret"]
                )
            return OAuthClientInformation.model_validate(client_info_data)
        return None

    def retrieve_tokens(self) -> OAuthTokens | None:
        """Retrieve OAuth tokens if authentication is complete.

        Returns:
            OAuthTokens if the provider has been authenticated, None otherwise.
        """
        if not self.credentials:
            return None
        credentials = self.decrypt_credentials()
        access_token = credentials.get("access_token", "")
        # Return None if access_token is empty to avoid generating invalid "Authorization: Bearer " header.
        # Note: We don't check for whitespace-only strings here because:
        # 1. OAuth servers don't return whitespace-only access tokens in practice
        # 2. Even if they did, the server would return 401, triggering the OAuth flow correctly
        if not access_token:
            return None
        return OAuthTokens(
            access_token=access_token,
            token_type=credentials.get("token_type", DEFAULT_TOKEN_TYPE),
            expires_in=int(credentials.get("expires_in", str(DEFAULT_EXPIRES_IN)) or DEFAULT_EXPIRES_IN),
            refresh_token=credentials.get("refresh_token", ""),
        )

    def masked_server_url(self) -> str:
        """Masked server URL for display"""
        parsed = urlparse(self.decrypt_server_url())
        if parsed.path and parsed.path != "/":
            masked = parsed._replace(path="/******")
            return masked.geturl()
        return parsed.geturl()

    def _mask_value(self, value: str) -> str:
        """Mask a sensitive value for display"""
        if len(value) > MIN_UNMASK_LENGTH:
            return value[:2] + MASK_CHAR * (len(value) - 4) + value[-2:]
        else:
            return MASK_CHAR * len(value)

    def masked_headers(self) -> dict[str, str]:
        """Masked headers for display"""
        return {key: self._mask_value(value) for key, value in self.decrypt_headers().items()}

    def masked_credentials(self) -> dict[str, str]:
        """Masked credentials for display"""
        credentials = self.decrypt_credentials()
        if not credentials:
            return {}

        masked = {}

        if "client_information" not in credentials or not isinstance(credentials["client_information"], dict):
            return {}
        client_info = credentials["client_information"]
        # Mask sensitive fields from nested structure
        if client_info.get("client_id"):
            masked["client_id"] = self._mask_value(client_info["client_id"])
        if client_info.get("encrypted_client_secret"):
            masked["client_secret"] = self._mask_value(
                encrypter.decrypt_token(self.tenant_id, client_info["encrypted_client_secret"])
            )
        if client_info.get("client_secret"):
            masked["client_secret"] = self._mask_value(client_info["client_secret"])
        return masked

    def decrypt_server_url(self) -> str:
        """Decrypt server URL"""
        return encrypter.decrypt_token(self.tenant_id, self.server_url)

    def _decrypt_dict(self, data: dict[str, Any]) -> dict[str, Any]:
        """Generic method to decrypt dictionary fields"""
        from core.tools.utils.encryption import create_provider_encrypter

        if not data:
            return {}

        # Only decrypt fields that are actually encrypted
        # For nested structures, client_information is not encrypted as a whole
        encrypted_fields = []
        for key, value in data.items():
            # Skip nested objects - they are not encrypted
            if isinstance(value, dict):
                continue
            # Only process string values that might be encrypted
            if isinstance(value, str) and value:
                encrypted_fields.append(key)

        if not encrypted_fields:
            return data

        # Create dynamic config only for encrypted fields
        config = [BasicProviderConfig(type=BasicProviderConfig.Type.SECRET_INPUT, name=key) for key in encrypted_fields]

        encrypter_instance, _ = create_provider_encrypter(
            tenant_id=self.tenant_id,
            config=config,
            cache=NoOpProviderCredentialCache(),
        )

        # Decrypt only the encrypted fields
        decrypted_data = encrypter_instance.decrypt({k: data[k] for k in encrypted_fields})

        # Merge decrypted data with original data (preserving non-encrypted fields)
        result = data.copy()
        result.update(decrypted_data)

        return result

    def decrypt_headers(self) -> dict[str, Any]:
        """Decrypt headers"""
        return self._decrypt_dict(self.headers)

    def decrypt_credentials(self) -> dict[str, Any]:
        """Decrypt credentials"""
        return self._decrypt_dict(self.credentials)

    def decrypt_authentication(self) -> dict[str, Any]:
        """Decrypt authentication"""
        # Option 1: if headers is provided, use it and don't need to get token
        headers = self.decrypt_headers()

        # Option 2: Add OAuth token if authed and no headers provided
        if not self.headers and self.authed:
            token = self.retrieve_tokens()
            if token:
                headers["Authorization"] = f"{token.token_type.capitalize()} {token.access_token}"
        return headers