"""
Secret Manager Handler

Handles retrieving secrets from different secret management systems.
"""
import base64
import os
from typing import Any, Optional

import litellm
from litellm._logging import print_verbose
from litellm.types.secret_managers.main import KeyManagementSystem


def _is_base64(s):
    """Check if a string is valid base64."""
    import binascii
    try:
        return base64.b64encode(base64.b64decode(s)).decode() == s
    except binascii.Error:
        return False


def get_secret_from_manager( # noqa: PLR0915
    client: Any,
    key_manager: str,
    secret_name: str,
    key_management_settings: Optional[Any] = None,
) -> Optional[str]:
    """
    Get a secret from the configured secret manager.
    
    Args:
        client: The secret manager client instance
        key_manager: The type of key manager (e.g., "azure_key_vault", "google_kms", etc.)
        secret_name: The name/path of the secret to retrieve
        key_management_settings: Optional settings for the key management system
        
    Returns:
        The secret value as a string, or None if not found
        
    Raises:
        ValueError: If the secret cannot be retrieved or required parameters are missing
        Exception: For other errors during secret retrieval
    """
    secret = None
    
    if (
        key_manager == KeyManagementSystem.AZURE_KEY_VAULT.value
        or type(client).__module__ + "." + type(client).__name__
        == "azure.keyvault.secrets._client.SecretClient"
    ):  # support Azure Secret Client - from azure.keyvault.secrets import SecretClient
        secret = client.get_secret(secret_name).value
        
    elif (
        key_manager == KeyManagementSystem.GOOGLE_KMS.value
        or client.__class__.__name__ == "KeyManagementServiceClient"
    ):
        encrypted_secret: Any = os.getenv(secret_name)
        if encrypted_secret is None:
            raise ValueError("Google KMS requires the encrypted secret to be in the environment!")
        b64_flag = _is_base64(encrypted_secret)
        if b64_flag is True:  # if passed in as encoded b64 string
            encrypted_secret = base64.b64decode(encrypted_secret)
            ciphertext = encrypted_secret
        else:
            raise ValueError(
                "Google KMS requires the encrypted secret to be encoded in base64"
            )  # fix for this vulnerability https://huntr.com/bounties/ae623c2f-b64b-4245-9ed4-f13a0a5824ce
        response = client.decrypt(
            request={
                "name": litellm._google_kms_resource_name,
                "ciphertext": ciphertext,
            }
        )
        secret = response.plaintext.decode("utf-8")  # assumes the original value was encoded with utf-8
        
    elif key_manager == KeyManagementSystem.AWS_KMS.value:
        """
        Only check the tokens which start with 'aws_kms/'. This prevents latency impact caused by checking all keys.
        """
        encrypted_value = os.getenv(secret_name, None)
        if encrypted_value is None:
            raise Exception("AWS KMS - Encrypted Value of Key={} is None".format(secret_name))
        # Decode the base64 encoded ciphertext
        ciphertext_blob = base64.b64decode(encrypted_value)

        # Set up the parameters for the decrypt call
        params = {"CiphertextBlob": ciphertext_blob}
        # Perform the decryption
        response = client.decrypt(**params)

        # Extract and decode the plaintext
        plaintext = response["Plaintext"]
        secret = plaintext.decode("utf-8")
        if isinstance(secret, str):
            secret = secret.strip()
            
    elif key_manager == KeyManagementSystem.AWS_SECRET_MANAGER.value:
        from litellm.secret_managers.aws_secret_manager_v2 import (
            AWSSecretsManagerV2,
        )

        if isinstance(client, AWSSecretsManagerV2):
            primary_secret_name = None
            if key_management_settings is not None:
                primary_secret_name = key_management_settings.primary_secret_name
                
            secret = client.sync_read_secret(
                secret_name=secret_name,
                primary_secret_name=primary_secret_name,
            )
            print_verbose(f"get_secret_value_response: {secret}")
            
    elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value:
        try:
            secret = client.get_secret_from_google_secret_manager(secret_name)
            print_verbose(f"secret from google secret manager:  {secret}")
            if secret is None:
                raise ValueError(f"No secret found in Google Secret Manager for {secret_name}")
        except Exception as e:
            print_verbose(f"An error occurred - {str(e)}")
            raise e
            
    elif key_manager == KeyManagementSystem.HASHICORP_VAULT.value:
        try:
            secret = client.sync_read_secret(secret_name=secret_name)
            if secret is None:
                raise ValueError(f"No secret found in Hashicorp Secret Manager for {secret_name}")
        except Exception as e:
            print_verbose(f"An error occurred - {str(e)}")
            raise e
            
    elif key_manager == KeyManagementSystem.CYBERARK.value:
        try:
            secret = client.sync_read_secret(secret_name=secret_name)
            if secret is None:
                raise ValueError(f"No secret found in CyberArk Secret Manager for {secret_name}")
        except Exception as e:
            print_verbose(f"An error occurred - {str(e)}")
            raise e
            
    elif key_manager == KeyManagementSystem.CUSTOM.value:
        # Check if client is a CustomSecretManager instance
        from litellm.integrations.custom_secret_manager import CustomSecretManager
        
        if isinstance(client, CustomSecretManager):
            secret = client.sync_read_secret(
                secret_name=secret_name,
                optional_params=key_management_settings.model_dump() if key_management_settings else None,
            )
            if secret is None:
                raise ValueError(f"No secret found in Custom Secret Manager for {secret_name}")
        else:
            raise ValueError(
                f"Custom secret manager client must be an instance of CustomSecretManager, got {type(client).__name__}"
            )
        
    elif key_manager == "local":
        secret = os.getenv(secret_name)
        
    else:  # assume the default is infisicial client
        secret = client.get_secret(secret_name).secret_value
        
    return secret