"""Lambda token_secret_rotation function module."""

import os

from aws_lambda_powertools import Logger
from aws_lambda_powertools.utilities.data_classes import (
    SecretsManagerEvent,
    event_source,
)
from aws_lambda_powertools.utilities.parameters import SecretsProvider
from aws_lambda_powertools.utilities.typing import LambdaContext

import sentry_sdk
from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration

from jwtauth.utils import jwt_auth_from_environment
import boto3
from mypy_boto3_secretsmanager.client import SecretsManagerClient

from src.auth0 import generate_auth_token
from src.models import DUMMY_VAL, M2MToken, ClientCredentials, create_dummy_m2m_token
from config import APPLICATION_NAME, ENVIRONMENT

logger = Logger(service=f"{ENVIRONMENT}-{APPLICATION_NAME}")
secrets_provider = SecretsProvider()


STEP_CREATE_SECRET = "createSecret"
STEP_SET_SECRET = "setSecret"
STEP_TEST_SECRET = "testSecret"
STEP_FINISH_SECRET = "finishSecret"


try:
    sentry_dsn = str(
        secrets_provider.get(f"{ENVIRONMENT}/{APPLICATION_NAME}/SENTRY_DSN")
    )
    sentry_sdk.init(
        dsn=sentry_dsn,
        traces_sample_rate=1.0,
        integrations=[AwsLambdaIntegration()],
    )
except Exception:
    logger.exception("Could not initialize Sentry SDK")


class Auth0M2MRotationError(Exception):
    pass


# We are ignoring the error: Untyped decorator makes function "handler" untyped  [misc]
# AWS Powertools is not 100% compatible with strict typing
# We might have it with PowerTools v3, but for now it's still v2
@event_source(data_class=SecretsManagerEvent)  # type: ignore
@logger.inject_lambda_context(log_event=True)
def handler(event: SecretsManagerEvent, context: LambdaContext) -> None:
    """Secrets Manager Rotation Lambda.

    Args:
        event (dict): Lambda dictionary of event parameters. Keys must include:
            - SecretId: The secret ARN or identifier
            - ClientRequestToken: The ClientRequestToken of the secret version
            - Step: The rotation step, one of:
                * createSecret
                * setSecret
                * testSecret
                * finishSecret

        context (LambdaContext): The Lambda runtime information

    Raises:
        ResourceNotFoundException: If the secret with
            the specified arn and stage does not exist

        ValueError: If the secret is not properly configured for rotation

        KeyError: If the event parameters do not contain the expected keys

    """
    arn = event.secret_id
    token = event.client_request_token
    step = event.step

    logger.set_correlation_id(f"{event.secret_id}_{event.client_request_token}")
    sentry_sdk.set_tags({"secret_id": arn, "step": step})

    # Setup the client
    service_client = boto3.client(
        "secretsmanager",
        endpoint_url=os.environ.get(
            "SECRETS_MANAGER_ENDPOINT",
            "https://secretsmanager.us-east-1.amazonaws.com",
        ),
    )

    # Make sure the version is staged correctly
    metadata = service_client.describe_secret(SecretId=arn)
    if not metadata["RotationEnabled"]:
        logger.error(
            "Secret %s is not enabled for rotation" % arn,
            extra={
                "arn": arn,
                "token": token,
                "step": step,
            },
        )
        raise ValueError("Secret %s is not enabled for rotation" % arn)
    versions = metadata["VersionIdsToStages"]
    if token not in versions:
        logger.error(
            "Secret version %s has no stage for rotation of secret %s." % (token, arn),
            extra={
                "arn": arn,
                "token": token,
                "step": step,
            },
        )
        raise ValueError(
            "Secret version %s has no stage for rotation of secret %s." % (token, arn)
        )
    if "AWSCURRENT" in versions[token]:
        logger.info(
            "Secret version %s already set as AWSCURRENT for secret %s." % (token, arn),
            extra={
                "arn": arn,
                "token": token,
                "step": step,
            },
        )
        return
    elif "AWSPENDING" not in versions[token]:
        logger.error(
            "Secret version %s not set as AWSPENDING for rotation of secret %s."
            % (token, arn),
            extra={
                "arn": arn,
                "token": token,
                "step": step,
            },
        )
        raise ValueError(
            "Secret version %s not set as AWSPENDING for rotation of secret %s."
            % (token, arn)
        )

    if step == STEP_CREATE_SECRET:
        create_secret(service_client, arn, token, context)

    elif step == STEP_SET_SECRET:
        set_secret(service_client, arn, token)

    elif step == STEP_TEST_SECRET:
        test_secret(service_client, arn, token)

    elif step == STEP_FINISH_SECRET:
        finish_secret(service_client, arn, token)

    else:
        raise ValueError("Invalid step parameter %s for secret %s" % (step, arn))


def create_secret(
    service_client: SecretsManagerClient, arn: str, token: str, context: LambdaContext
) -> None:
    """Create the secret.

    This method first checks for the existence of a secret with AWSPENDING stage
    and the passed in token. If one does not exist, it will generate a
    new secret and put it with as AWSPENDING stage with the passed in token.

    This method requires a valid AWSCURRENT stage secret in order to get
    the ClientCredentials SecretsManager ARN to generate the new secret.

    Args:
        service_client (client): The secrets manager service client

        arn (string): The secret ARN or other identifier

        token (string): The ClientRequestToken associated
            with the secret version

    Raises:
        InvalidRequestException: If it is unable to find the
            token as an existing versionId

        ResourceNotFoundException: If the secret with the
            specified arn and stage does not exist

    """
    try:
        service_client.get_secret_value(
            SecretId=arn,
            VersionId=token,
            VersionStage="AWSPENDING",
        )
        logger.info(
            "createSecret: Successfully retrieved secret for %s." % arn,
            extra={
                "arn": arn,
                "token": token,
                "step": STEP_CREATE_SECRET,
            },
        )
    except (
        service_client.exceptions.InvalidRequestException,
        service_client.exceptions.ResourceNotFoundException,
    ):
        _create_secret(service_client, arn, token, context)

        logger.info(
            "createSecret: Successfully put secret for ARN %s to AWSPENDING with version %s."
            % (arn, token),
            extra={
                "arn": arn,
                "token": token,
                "step": STEP_CREATE_SECRET,
            },
        )


def set_secret(
    service_client: SecretsManagerClient,
    arn: str,
    token: str,
) -> None:
    """Set the secret.

    AWS Docs:
    This method should set the AWSPENDING secret in the service that
    the secret belongs to. For example, if the secret is a database
    credential, this method should take the value of the AWSPENDING
    secret and set the user's password to this value in the database.

    In our case, the "database" is Auth0, and we do not need to
    set this password on Auth0. Instead, this step is used to validate
    client_credentials_secret_arn is same between
    AWSPENDING and AWSCURRENT.
    Args:
        service_client (client): The secrets manager service client

        arn (string): The secret ARN or other identifier

        token (string): The ClientRequestToken associated
            with the secret version

    """
    pending_secret_value = service_client.get_secret_value(
        SecretId=arn,
        VersionId=token,
        VersionStage="AWSPENDING",
    )["SecretString"]
    logger.info(
        "setSecret: Successfully retrieved AWSPENDING secret for %s, %s."
        % (arn, token),
        extra={
            "arn": arn,
            "token": token,
            "step": STEP_SET_SECRET,
        },
    )

    current_secret_value = service_client.get_secret_value(
        SecretId=arn,
        VersionStage="AWSCURRENT",
    )["SecretString"]
    logger.info(
        "setSecret: Successfully retrieved AWSCURRENT secret for %s." % arn,
        extra={
            "arn": arn,
            "token": token,
            "step": STEP_SET_SECRET,
        },
    )

    pending_secret = M2MToken.model_validate_json(pending_secret_value)
    current_secret = M2MToken.model_validate_json(current_secret_value)
    assert (
        pending_secret.client_credentials_secret_arn
        == current_secret.client_credentials_secret_arn
    ), "Client Credentials Secret ARN must not change between stages."
    logger.info(
        "setSecret: Successfully set %s" % arn,
        extra={
            "arn": arn,
            "token": token,
            "step": STEP_SET_SECRET,
        },
    )


def test_secret(
    service_client: SecretsManagerClient,
    arn: str,
    token: str,
) -> None:
    """Test the secret.

    AWS Docs: This method should validate that the AWSPENDING secret works in the service
    that the secret belongs to. For example, if the secret
    is a database credential, this method should validate that the user
    can login with the password in AWSPENDING and that the user has
    all of the expected permissions against the database.

    In our case, we want to test:
    * the token in the AWSPENDING secret can be decoded/is_valid
    * the expires_at in the AWSPENDING secret is in the future

    Args:
        service_client (client): The secrets manager service client
        arn (string): The secret ARN or other identifier
        token (string): The ClientRequestToken associated
            with the secret version
    """
    # Get the AWSPENDING secret
    pending_secret_value = service_client.get_secret_value(
        SecretId=arn,
        VersionId=token,
        VersionStage="AWSPENDING",
    )["SecretString"]
    logger.info(
        "testSecret: Successfully retrieved AWSPENDING secret for %s, %s."
        % (arn, token),
        extra={
            "arn": arn,
            "token": token,
            "step": STEP_TEST_SECRET,
        },
    )
    pending_secret = M2MToken.model_validate_json(pending_secret_value)

    if pending_secret.token == DUMMY_VAL:
        logger.info(
            "testSecret: AWSPENDING secret is still dummy for %s, %s; skipping expires_at check"
            % (arn, token),
            extra={
                "arn": arn,
                "token": token,
                "step": STEP_TEST_SECRET,
            },
        )
        return

    # Decode the token - this will raise an exception
    # if it cannot be decoded/isn't valid
    auth = jwt_auth_from_environment(environment=ENVIRONMENT)
    auth.get_token(pending_secret.token)
    logger.info(
        "testSecret: Successfully tested AWSPENDING secret for %s, %s." % (arn, token),
        extra={
            "arn": arn,
            "token": token,
            "step": STEP_TEST_SECRET,
        },
    )

    # Only check the pending secret expiration if not DUMMY_VAL
    assert not pending_secret.is_expired(), f"testSecret: AWSPENDING secret expires_at {pending_secret.expires_at} - should not be expired"

    logger.info(
        "testSecret: Successfully confirmed expiration of AWSPENDING secret for %s, %s."
        % (arn, token),
        extra={
            "arn": arn,
            "token": token,
            "step": STEP_TEST_SECRET,
        },
    )


def finish_secret(
    service_client: SecretsManagerClient,
    arn: str,
    token: str,
) -> None:
    """Finish the secret.

    This method finalizes the rotation process by marking the secret
    version passed in as the AWSCURRENT secret.

    Args:
        service_client (client): The secrets manager service client

        arn (string): The secret ARN or other identifier

        token (string): The ClientRequestToken associated with
            the secret version

    Raises:
        ResourceNotFoundException: If the secret with the specified
            arn does not exist

    """
    # First describe the secret to get the current version
    metadata = service_client.describe_secret(SecretId=arn)
    current_version = None
    for version in metadata["VersionIdsToStages"]:
        if "AWSCURRENT" in metadata["VersionIdsToStages"][version]:
            if version == token:
                # The correct version is already marked as current, return
                logger.info(
                    "finishSecret: Version %s already marked as AWSCURRENT for %s"
                    % (version, arn),
                    extra={
                        "arn": arn,
                        "token": token,
                        "step": STEP_FINISH_SECRET,
                    },
                )
                return
            current_version = version
            break

    assert current_version, "current_version must not be None"

    # Finalize by staging the secret version current
    service_client.update_secret_version_stage(
        SecretId=arn,
        VersionStage="AWSCURRENT",
        MoveToVersionId=token,
        RemoveFromVersionId=current_version,
    )
    logger.info(
        "finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s."
        % (token, arn),
        extra={
            "arn": arn,
            "token": token,
            "step": STEP_FINISH_SECRET,
        },
    )


def _create_secret(
    service_client: SecretsManagerClient, arn: str, token: str, context: LambdaContext
) -> None:
    """Create the new secret value and include details for handled errors."""
    try:
        current_secret_value = service_client.get_secret_value(
            SecretId=arn,
            VersionStage="AWSCURRENT",
        )["SecretString"]
        current_secret = M2MToken.model_validate_json(current_secret_value)
        client_credentials_secret_arn = current_secret.client_credentials_secret_arn

        client_credentials_value = service_client.get_secret_value(
            SecretId=client_credentials_secret_arn,
            VersionStage="AWSCURRENT",
        )["SecretString"]
        if client_credentials_value == DUMMY_VAL:
            logger.warning(
                "createSecret: Client credentials %s have not been configured, setting dummy M2MToken."
                % (client_credentials_secret_arn),
                extra={
                    "arn": arn,
                    "token": token,
                    "step": STEP_CREATE_SECRET,
                    "client_credentials_secret_arn": client_credentials_secret_arn,
                },
            )
            pending_secret = create_dummy_m2m_token(client_credentials_secret_arn)
        else:
            client_credentials = ClientCredentials.model_validate_json(
                client_credentials_value,
            )

            # Generate the new secret
            oauth_token = generate_auth_token(client_credentials)
            pending_secret = oauth_token.convert_to_m2m_token(
                client_credentials_secret_arn
            )

        service_client.put_secret_value(
            SecretId=arn,
            ClientRequestToken=token,
            SecretString=pending_secret.model_dump_json(),
            VersionStages=["AWSPENDING"],
        )
    except Exception as ex:
        # re-raise the error with additional error info.
        raise Auth0M2MRotationError(
            f"create_secret error. request_id:{context.aws_request_id}. arn:{arn}"
        ) from ex