# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
#

"""
AWS utilities for working with AWS services.

This module provides functions for:
- AWS authentication and session management
- S3 URL parsing
- API URL handling
"""

from urllib.parse import urljoin, urlparse, urlsplit

import boto3

from src.exceptions import MalformedS3Url


def get_aws_credentials(session=None, ask_mfa=False, token=None, serial_number=None):
    """
    Get AWS credentials with optional MFA support.

    Args:
        session: Optional existing boto3 session
        ask_mfa: Whether to prompt for MFA
        token: MFA token code
        serial_number: MFA device serial number

    Returns:
        dict: AWS credentials
    """
    if not session:
        session = boto3.Session()

    if ask_mfa and (token is None or serial_number is None):
        raise ValueError("MFA token and serial number are required when ask_mfa is True")

    if ask_mfa:
        sts = session.client("sts")
        response = sts.get_session_token(DurationSeconds=3600, SerialNumber=serial_number, TokenCode=token)
        credentials = response["Credentials"]
        return {
            "aws_access_key_id": credentials["AccessKeyId"],
            "aws_secret_access_key": credentials["SecretAccessKey"],
            "aws_session_token": credentials["SessionToken"],
        }
    else:
        creds = session.get_credentials()
        # If frozen credentials (like from a config file)
        if hasattr(creds, "access_key") and hasattr(creds, "secret_key"):
            credentials = {"aws_access_key_id": creds.access_key, "aws_secret_access_key": creds.secret_key}
            if hasattr(creds, "token") and creds.token:
                credentials["aws_session_token"] = creds.token
            return credentials
        # Regular credentials
        else:
            frozen = creds.get_frozen_credentials()
            credentials = {"aws_access_key_id": frozen.access_key, "aws_secret_access_key": frozen.secret_key}
            if frozen.token:
                credentials["aws_session_token"] = frozen.token
            return credentials


def get_aws_session(credentials=None, region=None, profile=None):
    """
    Create a boto3 session from AWS credentials.

    Args:
        credentials: Optional dict of AWS credentials
        region: AWS region
        profile: AWS profile name

    Returns:
        boto3.Session: Boto3 session object
    """
    session_args = {}
    if profile:
        session_args["profile_name"] = profile
    if region:
        session_args["region_name"] = region
    if credentials:
        session_args.update(credentials)
    return boto3.Session(**session_args)


def parse_s3_url(s3_url: str) -> tuple[str, str]:
    """
    Parse an S3 URL, returning a tuple containing the bucket name and key.

    Args:
        s3_url: s3 URL, like s3://some_bucket/some_path/and/maybe/a.file

    Returns:
        tuple: bucket name (some_bucket) and key (some_path/and/maybe/a.file)

    Raises:
        MalformedS3Url: If the URL is not a valid S3 URL
    """
    parsed_config_location = urlparse(s3_url)
    if parsed_config_location.scheme == "s3":
        bucket = parsed_config_location.netloc
        key = parsed_config_location.path.lstrip("/")
    else:
        raise MalformedS3Url(f"S3 url '{s3_url}' is malformed")

    return bucket, key


def join_api_url(endpoint: str, path: str) -> str:
    """
    Join API endpoint and path ensuring proper formatting.

    Args:
        endpoint: Base API endpoint URL
        path: API path to append

    Returns:
        str: Joined URL
    """
    return urljoin(endpoint.rstrip("/") + "/", path.lstrip("/"))


def clean_url(url: str) -> str | None:
    """
    Clean a URL by removing duplicate slashes and ensuring proper format.

    Args:
        url: URL to clean

    Returns:
        str: Cleaned URL, or None if URL is invalid
    """
    if not url:
        return None
    parts = urlsplit(url)
    if not all([parts.scheme, parts.netloc]):
        return None
    return f"{parts.scheme}://{parts.netloc}/{parts.path.replace('//', '/').rstrip('/').lstrip('/')}"
