"""AWS utils and helpers."""

from collections import namedtuple
from datetime import datetime
import json
import sys
from tempfile import NamedTemporaryFile
from typing import TypedDict

from airflow.providers.amazon.aws.hooks.s3 import S3Hook
import boto3
import s3fs
from slugify import slugify

from lib import config
from lib import constants
from lib.config import S3_SALES_BUCKET_NAME

S3Location = namedtuple('S3Location', ('key', 'url'))
S3Object = namedtuple('S3Object', ('key', 'bucket'))


class AwsCredentials(TypedDict):
    """Represent the credentials returned by `STS.Client.assume_role`."""

    AccessKeyId: str
    SecretAccessKey: str
    SessionToken: str
    Expiration: datetime


# Security Token Service
def get_sts_client():
    """Get boto3 STS client."""
    return boto3.client('sts')


def get_credentials_for_assumed_role(
    sts_client,
    role_arn,
    session_name
) -> AwsCredentials:
    """Assume a role and return the temporary credentials."""
    response = sts_client.assume_role(
        RoleArn=role_arn,
        RoleSessionName=session_name
    )
    return response['Credentials']


# Secrets Manager
def get_secrets_manager_client():
    """Get boto3 SM client."""
    return boto3.client('secretsmanager')


def get_secret(client, secret_name):
    """Fetch secret string from secrets manager."""
    secret_full_path = \
        f'{config.OWS_ENV}/{constants.AIRFLOW_SERVICE_NAME}/{secret_name}'
    secret_value = client.get_secret_value(SecretId=secret_full_path)
    return secret_value.get('SecretString')


# S3 Path Helpers
def write_file(url, content, s3_options=None):
    """Write a file to S3."""
    if s3_options is None:
        s3_options = {}

    fs = s3fs.S3FileSystem(s3_additional_kwargs=s3_options)
    file_size = sys.getsizeof(content)
    print(f'WRITING {file_size} TO {url}')
    with fs.open(url, 'wb') as f:
        f.write(content)


def copy_file(source_bucket, source_key, dest_key, dest_bucket=None):
    """Copy a file from one S3 location to another."""
    params = {
        'source_bucket_name': source_bucket,
        'dest_bucket_name': dest_bucket if dest_bucket else source_bucket
    }
    print(
        f'COPY FROM {source_bucket}/{source_key}'
        f'TO {params.get("dest_bucket_name")}/{dest_key}'
    )
    result = S3Hook().copy_object(
        source_key,
        dest_key,
        **params
    )
    print('COPY RESULT', result)
    return result


def delete_files(bucket: str, key: str = None, keys: list = None):
    """Delete a file or list of files from the specified bucket.

    Args:
        bucket (str): name of the S3 bucket
        key (str): the name of the object to be deleted from the bucket
        keys (list): the list of objects to be deleted from the bucket
    """
    params = {
        'bucket': bucket,
        'keys': key or keys
    }

    print(f'DELETE {bucket}/{key}')

    S3Hook().delete_objects(**params)


def get_file(account_id, bucket_name, key):
    """Download file from specified bucket/key."""
    file = NamedTemporaryFile(suffix='.xlsx', delete=False)
    s3_client = boto3.client('s3')
    print('RETRIEVING FILE', bucket_name, key)
    s3_client.download_file(
        bucket_name,
        key,
        file.name,
        ExtraArgs={
            'ExpectedBucketOwner': account_id
        }
    )
    return file.name


def file_exists(url):
    """Return boolean result for whether s3 file path exists."""
    fs = s3fs.S3FileSystem()
    result = fs.exists(url)
    print('FILE EXISTS RESULT', url, result)
    return result


def split_path(s3_path):
    """Return bucket name and key name out of full s3 path."""
    fs = s3fs.S3FileSystem()
    result = fs.split_path(s3_path)
    return S3Object(bucket=result[0], key=result[1])


def read_key(bucket_name, key):
    """Read contents of specified bucket/key."""
    print('READING FILE', bucket_name, key)
    return S3Hook().read_key(key, bucket_name=bucket_name)


def location(*key_segments, bucket_name=S3_SALES_BUCKET_NAME):
    """Build an S3 location from path segments."""
    key = join(*key_segments)
    return S3Location(key=key, url=f's3://{bucket_name}/{key}')


def join(*segments):
    """Build a key from key segments."""
    return '/'.join(segments)


def sanitize(text, max_length=100, **kwargs):
    """Sanitize a string for use as part of a bucket name or url."""
    return slugify(text, max_length=max_length, **kwargs)


def read_json_from_s3(prefix, bucket_name):
    """Read file metadata from S3."""
    print('LISTING BUCKET', bucket_name, 'PREFIX', prefix)
    key = S3Hook().list_keys(bucket_name=bucket_name, prefix=prefix)[0]
    results = read_key(bucket_name, key)
    print('FILE RESULTS:', results)
    return json.loads(results)


def list_bucket_keys(bucket, prefix):
    """List keys in specified S3 bucket."""
    result = S3Hook().list_keys(bucket_name=bucket, prefix=prefix)
    print('LIST BUCKET KEYS', bucket, prefix, result)
    return result


# ECS Helpers
def get_ecs_client(credentials: AwsCredentials | None = None):
    """Get boto3 ECS client."""
    if credentials:
        return boto3.client(
            'ecs',
            region_name='us-east-1',
            aws_access_key_id=credentials['AccessKeyId'],
            aws_secret_access_key=credentials['SecretAccessKey'],
            aws_session_token=credentials['SessionToken'],
        )

    return boto3.client('ecs')


# EC2 Helpers
def get_ec2_client(credentials: AwsCredentials | None = None):
    """Get boto3 EC2 client."""
    if credentials:
        return boto3.client(
            'ec2',
            region_name='us-east-1',
            aws_access_key_id=credentials['AccessKeyId'],
            aws_secret_access_key=credentials['SecretAccessKey'],
            aws_session_token=credentials['SessionToken'],
        )

    return boto3.client('ec2')


# Network Configuration Helpers
def get_vpc(ec2_client) -> str:
    """Get the first available VPC.

    Args:
        ec2_client (boto3.client): An EC2 client.

    Returns:
        str: The VPC ID.
    """
    response = ec2_client.describe_vpcs()
    vpc_id = response['Vpcs'][0]['VpcId']
    return vpc_id


def get_security_group(ec2_client, vpc_id: str, service_name: str) -> str:
    """Get the security group corresponding to a service.

    Args:
        ec2_client (boto3.client): An EC2 client.
        vpc_id (str):  The VPC ID of the service.
        service_name (str): The name of the service.

    Returns:
        str: The security group ID.
    """
    security_group_name = '{}-task-security-group'.format(service_name)

    response = ec2_client.describe_security_groups(
        Filters=[
            {
                'Name': 'group-name',
                'Values': [
                    security_group_name,
                ]
            },
            {
                'Name': 'vpc-id',
                'Values': [
                    vpc_id,
                ]
            },
        ],
    )

    security_group_id = response['SecurityGroups'][0]['GroupId']
    return security_group_id


def get_subnet(ec2_client, vpc_id: str) -> str:
    """Get the subnet with the most available IP addresses inside a VPC.

    Args:
        ec2_client (boto3.client): An EC2 client.
        vpc_id (str): The VPC ID.

    Returns:
        str: The subnet ID.
    """
    response = ec2_client.describe_subnets(
        Filters=[
            {
                'Name': 'tag:Name',
                'Values': [
                    '*private*',
                ]
            },
            {
                'Name': 'tag:tier',
                'Values': [
                    'private',
                ]
            },
            {
                'Name': 'vpc-id',
                'Values': [
                    vpc_id,
                ]
            },
        ],
    )

    sorted_subnets = sorted(
        response['Subnets'],
        key=lambda k: k['AvailableIpAddressCount'],
        reverse=True
    )

    subnet_id = sorted_subnets[0]['SubnetId']
    return subnet_id