import base64
import logging
from typing import Any, cast

import aioboto3
import botocore

from labels.model.sources import AwsCredentials, AwsRole
from labels.utils.exceptions import AwsCredentialsNotFoundError

LOGGER = logging.getLogger(__name__)


async def get_credentials(credentials: AwsRole) -> AwsCredentials | None:
    try:
        async with aioboto3.Session().client(service_name="sts") as sts_client:
            response = await sts_client.assume_role(
                ExternalId=credentials.external_id,
                RoleArn=credentials.role,
                RoleSessionName="FluidAttacksRoleVerification",
            )
            aws_credentials = AwsCredentials(
                access_key_id=response["Credentials"]["AccessKeyId"],
                secret_access_key=response["Credentials"]["SecretAccessKey"],
                session_token=response["Credentials"]["SessionToken"],
            )
    except botocore.exceptions.ClientError:  # type: ignore[misc]
        LOGGER.exception("We found problems on your AWS credentials")
        raise
    else:
        return aws_credentials


async def obtain_auth_data(  # type: ignore[explicit-any]
    credentials: AwsCredentials,
    region: str = "us-east-1",
) -> dict[str, Any]:
    try:
        session = aioboto3.Session(
            aws_access_key_id=credentials.access_key_id,
            aws_secret_access_key=credentials.secret_access_key,
            aws_session_token=credentials.session_token,
        )
        async with session.client("ecr", region) as client:
            authorization_data = await client.get_authorization_token()
            return cast("dict[str, Any]", authorization_data)  # type: ignore[explicit-any, misc]
    except botocore.exceptions.ClientError:  # type: ignore[misc]
        LOGGER.exception("We found problems on your AWS credentials")
        raise


async def get_token(credentials: AwsCredentials, region: str = "us-east-1") -> str:
    authorization_data = await obtain_auth_data(credentials=credentials, region=region)  # type: ignore[misc]
    auth_resp = authorization_data["authorizationData"][0]  # type: ignore[misc]
    token = auth_resp["authorizationToken"]  # type: ignore[misc]
    _, password = base64.b64decode(token).decode("utf-8").split(":")  # type: ignore[misc]
    return password


async def ecr_connection(role: AwsRole, region: str = "us-east-1") -> str:
    credentials = await get_credentials(role)

    if not credentials:
        raise AwsCredentialsNotFoundError

    token = await get_token(credentials, region=region)
    return f"AWS:{token}"
