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

import json
from typing import Any
from zoneinfo import ZoneInfo

from mypy_boto3_sts.type_defs import CredentialsTypeDef

from video import config
from video.connectors import sts
from video.constants import job_io_fields
from video.constants.thumbnails import UPLOAD_CUSTOM_THUMBNAIL


def _to_dict(
    credentials: CredentialsTypeDef,
    workflow_job_id: int,
    workflow_job_type: str,
    filename: str,
) -> dict[str, Any]:
    """Standardize AWS credentials dictionary to our convention.

    Args:
        credentials (dict): AWS credentials dictionary.
    Returns:
        dict: our convention for an AWS token.
    """
    if workflow_job_type == UPLOAD_CUSTOM_THUMBNAIL:
        s3_key = "thumbnails/custom/{}/{}".format(workflow_job_id, filename)
    else:
        s3_key = "raw/{}/{}".format(workflow_job_id, filename)

    return {
        job_io_fields.TOKEN: credentials["SessionToken"],
        job_io_fields.S3_TOKEN_EXPIRATION: credentials["Expiration"]
        .astimezone(ZoneInfo("UTC"))
        .strftime(config.S3_TOKEN_EXPIRATION_FORMAT),
        job_io_fields.AWS_ACCESS_KEY: credentials["AccessKeyId"],
        job_io_fields.AWS_SECRET_ACCESS_KEY: credentials["SecretAccessKey"],
        job_io_fields.INPUT_VIDEO_S3_BUCKET: config.VIDEO_BUCKET_NAME,
        job_io_fields.INPUT_VIDEO_S3_KEY: s3_key,
    }


def policy(workflow_job_id: int, workflow_job_type: str) -> str:
    """Return the policy."""
    if workflow_job_type == UPLOAD_CUSTOM_THUMBNAIL:
        resource = "arn:aws:s3:::{}/thumbnails/custom/{}/*".format(
            config.VIDEO_BUCKET_NAME, workflow_job_id
        )
    else:
        resource = "arn:aws:s3:::{}/raw/{}/*".format(
            config.VIDEO_BUCKET_NAME, workflow_job_id
        )

    return json.dumps(
        {
            "Version": "2012-10-17",
            "Statement": [
                {"Effect": "Allow", "Action": ["s3:PutObject"], "Resource": resource}
            ],
        }
    )


def get_s3_token(
    workflow_job_id: int, workflow_job_type: str, filename: str
) -> dict[str, Any]:
    """Generate session token for S3 Upload.

    Args:
        workflow_job_id (int): Workflow id.
        workflow_job_type (str): Workflow type.
        filename (str): Filename.
    Returns:
        dict: containing credentials.
    """
    connection = sts.connect_to_sts()
    role_arn = config.UPLOAD_RAW_DIR_ROLE_ARN
    assert role_arn is not None, "UPLOAD_RAW_ROLE_ARN env var is not set"
    credentials = connection.assume_role(
        RoleArn=role_arn,
        RoleSessionName=workflow_job_type,
        DurationSeconds=config.UPLOAD_RAW_DIR_ROLE_DURATION_SECONDS,
        Policy=policy(workflow_job_id, workflow_job_type),
    )
    return _to_dict(
        credentials["Credentials"], workflow_job_id, workflow_job_type, filename
    )