"""Model representation of temporary AWS token."""

import json
from typing import Any

from mypy_boto3_sts.type_defs import AssumeRoleResponseTypeDef

from assets import config
from assets.connectors import sts


def _to_dict(sts_response: AssumeRoleResponseTypeDef) -> dict[str, Any]:
    """Format boto3 sts client response.

    Args:
        sts_response (dict): boto3 response.
    Returns:
        dict: our convention for an AWS token.
    """
    credentials = sts_response["Credentials"]
    return {
        "token": credentials["SessionToken"],
        "aws_access_key_id": credentials["AccessKeyId"],
        "aws_secret_access_key": credentials["SecretAccessKey"],
        "expiration": str(credentials["Expiration"]),
    }


def _prepare_iam_policy_s3(bucket: str, filename: str) -> dict[str, Any]:
    """Modify the IAM policy to limit uploads to bucket and filename.

    Args:
        filename: Name of the file to put.
    Returns:
        dict: The IAM policy

    """
    resource = "arn:aws:s3:::{bucket}/{filename}.*".format(
        bucket=bucket, filename=filename
    )
    return {
        "Version": "2012-10-17",
        "Statement": [
            {"Effect": "Allow", "Action": ["s3:PutObject"], "Resource": resource}
        ],
    }


def get_s3_token(
    filename: str, duration: int = config.STS_TOKEN_DURATION
) -> dict[str, Any]:
    """Generate session token for S3 Upload.

    Args:
        filename (string): Name of the file, sans extension to be uploaded.
        duration (int): session token lifetime, in seconds.
            Default is 900 seconds - 15 minutes.
    Returns:
        dict: credentials.
    """
    sts_client = sts.get_sts_client()
    policy = _prepare_iam_policy_s3(config.RAW_ASSETS_BUCKET_NAME, filename)
    sts_response = sts_client.assume_role(
        RoleArn=config.IAM_ROLE,
        RoleSessionName=filename,
        DurationSeconds=duration,
        Policy=json.dumps(policy),
    )
    return _to_dict(sts_response)


def _prepare_iam_policy_s3_entity(
    bucket: str, paths: dict[str, Any], filename: str
) -> dict[str, Any]:
    """Modify the IAM policy to limit uploads to bucket and folder.

    Args:
        bucket (str): Name of the bucket.
        paths (dict): path of the file where file to be put on S3.
        filename (str): Name of the file.
    Returns:
        dict: The IAM policy

    """
    resources = []
    for key, path in paths.items():
        resources.append(
            "arn:aws:s3:::{bucket}/{path}/{filename}".format(
                bucket=bucket, path=path, filename=filename
            )
        )

    return {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": ["s3:PutObject", "s3:GetObject", "s3:DeleteObject"],
                "Resource": resources,
            },
            {
                "Effect": "Allow",
                "Action": ["cloudfront:CreateInvalidation"],
                "Resource": "*",
            },
        ],
    }


def get_s3_entity_token(
    path: dict[str, Any], filename: str, duration: int = config.STS_TOKEN_DURATION
) -> dict[str, Any]:
    """Generate session token for S3 Upload.

    Args:
        path (dict): path of the file where file to be put on S3.
        filename (string): Name of the file
        duration (int): session token lifetime, in seconds.
            Default is 900 seconds - 15 minutes.
    Returns:
        dict: credentials.
    """
    sts_client = sts.get_sts_client()
    policy = _prepare_iam_policy_s3_entity(
        config.ASSET_STORAGE_BUCKET_NAME, path, filename
    )
    sts_response = sts_client.assume_role(
        RoleArn=config.IAM_ROLE,
        RoleSessionName=filename,
        DurationSeconds=duration,
        Policy=json.dumps(policy),
    )
    return _to_dict(sts_response)