"""Connector for assuming role to another AWS Account's Secrets Manager."""
import logging
from datetime import UTC, datetime
from uuid import uuid4

import boto3
from botocore.exceptions import ClientError
from mypy_boto3_secretsmanager.client import SecretsManagerClient
from mypy_boto3_sts.client import STSClient
from pydantic import (
    BaseModel,
    Field,
)

logger = logging.getLogger(__name__)


class AssumeRoleCredentials(BaseModel):
    """AssumeRoleCredentials is a representation of the STS Credentials"""

    aws_access_key_id: str | None = Field(None, alias="AccessKeyId")
    aws_secret_access_key: str | None = Field(None, alias="SecretAccessKey")
    aws_session_token: str | None = Field(None, alias="SessionToken")
    expiration: datetime | None = Field(None, alias="Expiration")


class SecretScheduledForDeletionError(Exception):
    pass


def is_expired(expiration_datetime: datetime) -> bool:
    """Utility function to compare a time to current to determine if it is expired"""
    current = datetime.now(UTC)
    if expiration_datetime < current:
        return True
    return False


class AssumeRoleSecretsManager:
    """Assume Role to Secrets Manager."""

    def __init__(
        self,
        environment: str,
        sts_client: STSClient,
        aws_account_id: str,
        aws_assume_role_name: str,
        duration_seconds: int = 3600,
    ):
        """Initialize AssumeRoleSecretsManager."""
        self._environment = environment
        self._sts_client = sts_client
        self._aws_account_id = aws_account_id
        self._aws_assume_role_name = aws_assume_role_name
        self._duration_seconds = duration_seconds
        self._expiration = None
        self._check_expiration_and_setup_secrets_manager_client()

    def _build_assume_role_arn(self) -> str:
        """Build arn for the role being assumed."""
        return f"arn:aws:iam::{self._aws_account_id}:role/{self._environment}-{self._aws_assume_role_name}"  # noqa: E501

    def _get_assume_role_credentials(self) -> AssumeRoleCredentials:
        """Invoke assume role and return the credentials."""
        result = self._sts_client.assume_role(
            RoleArn=self._build_assume_role_arn(),
            RoleSessionName=str(uuid4()),
            DurationSeconds=self._duration_seconds,
        )
        return AssumeRoleCredentials.model_validate(result["Credentials"])

    def _setup_secrets_manager_client(
        self,
        credentials: AssumeRoleCredentials,
    ) -> None:
        """Create the secrets manager client from assume role credentials."""
        credentials_dict = credentials.model_dump()
        expiration = credentials_dict.pop("expiration")

        self._secrets_manager_client: SecretsManagerClient = boto3.client(
            "secretsmanager",
            **credentials_dict,
        )
        self._expiration = expiration

    def _is_secrets_manager_client_expired(self) -> bool:
        """Compare current time and expiration to determine if the client is expired."""
        if not self._expiration:
            return True
        return is_expired(self._expiration)

    def _check_expiration_and_setup_secrets_manager_client(self) -> None:
        """Handle an expired secrets manager client."""
        if self._is_secrets_manager_client_expired():
            credentials = self._get_assume_role_credentials()
            self._setup_secrets_manager_client(credentials)

    def does_secret_exist(self, secret_name: str) -> bool:
        """Return if secret exists."""
        self._check_expiration_and_setup_secrets_manager_client()
        try:
            self._secrets_manager_client.describe_secret(SecretId=secret_name)
        except ClientError as client_error:
            if client_error.response["Error"]["Code"] == "ResourceNotFoundException":
                logger.warning(
                    "Secret does not exist",
                    extra={
                        "secret_name": secret_name,
                        "aws_account_id": self._aws_account_id,
                    },
                )
            elif client_error.response["Error"]["Code"] == "AccessDeniedException":
                logger.error(
                    client_error.response["Error"]["Message"],
                    extra={
                        "secret_name": secret_name,
                        "aws_account_id": self._aws_account_id,
                    },
                )
            return False
        except Exception:
            logger.exception("Could not get secret")
            return False
        return True

    def get_secret_string(self, secret_name: str) -> str:
        """Get the secret."""
        self._check_expiration_and_setup_secrets_manager_client()
        try:
            result = self._secrets_manager_client.get_secret_value(SecretId=secret_name)
        except ClientError as client_error:
            if "marked for deletion" in client_error.response["Error"]["Message"]:
                raise SecretScheduledForDeletionError(
                    f"Secret name: {secret_name} is scheduled for deletion"
                ) from client_error
            raise client_error
        return result["SecretString"]

    def save_secret_string(self, secret_name: str, secret_value: str) -> bool:
        """Write/update the secret."""
        self._check_expiration_and_setup_secrets_manager_client()

        try:
            response = self._secrets_manager_client.put_secret_value(
                SecretId=secret_name,
                SecretString=secret_value,
            )

            return "VersionId" in response
        except Exception:
            logger.exception(
                "Unable to update secret",
                extra={
                    "secret_name": secret_name,
                    "aws_account_id": self._aws_account_id,
                },
            )
            return False