"""Module that provides generation of upload data."""

import hashlib
import uuid
from datetime import datetime, timezone
from typing import Any

import sentry_sdk
from botocore.exceptions import ClientError
from sqlalchemy.exc import SQLAlchemyError

from assets import config
from assets.connectors import s3 as s3_connector
from assets.constants import (
    api,
    asset_status as asset_status_constants,
    asset_types,
    asset_upload as asset_upload_constants,
    asset_upload_types as asset_upload_types_constants,
    product as product_constants,
    s3,
)
from assets.exceptions import (
    AssetUploadError,
    InvalidAssetUploadType,
    InvalidProductStatus,
    S3MultipartUploadFailed,
    TrackDoesNotBelongToProduct,
    WrongAssetStatus,
)
from assets.models import (
    asset_status as asset_status_model,
    asset_upload as asset_upload_model,
    asset_upload_type as asset_upload_type_model,
    ows_product as ows_product_model,
    ows_track as ows_track_model,
    token,
)


def _generate_filename() -> str:
    unique_name = uuid.uuid4()
    underscored_name = str(unique_name).replace("-", "_")
    return underscored_name


def _get_additional_fields_for_asset_upload(
    product_id: int, asset_upload_type: str
) -> dict[str, Any]:
    """Get additional fields for asset_upload.

    Args:
        product_id (int): Product id
        asset_upload_type (str): Asset upload type name

    Returns:
        dict: upc and is_correction flag.
    """
    product = ows_product_model.get_product_by_id(product_id)

    release_status = product["status"]
    context_type = product["context_type"]
    upc = product["upc"]

    if release_status == product_constants.PRODUCT_STATUS_LABEL_PROCESSING:
        return {"upc": upc, "is_correction": False}

    if release_status == product_constants.PRODUCT_STATUS_IN_CONTENT:
        if context_type == "digital":
            # TODO CDAM-3960: Atmos uploads temporarily bypass the correction/approval flow.
            if asset_upload_type == asset_upload_types_constants.ATMOS:
                return {"upc": upc, "is_correction": False}
            return {"upc": upc, "is_correction": True}
        if context_type == "physical":
            # We allow replacing artwork for in_content physical products without approval.
            return {"upc": upc, "is_correction": False}

    raise InvalidProductStatus("Product is not in a state that allows asset upload")


def create_asset_upload(
    *,
    product_id: int,
    track_unique_id: int,
    original_filename: str,
    user_id: str,
    asset_upload_type: str
    | None = None,  # TODO: CDAM-3825: Remove default once asset_upload_type is required in the schema
) -> str:
    """Create an asset upload.

    An asset_upload record will be created with asset, product, and track info,
    and an asset_status record will be created with status "uploading".

    An AWS S3 multipart upload will be initiated and the upload_id will be saved.

    Args:
        product_id (int): Product id
        track_unique_id (int): Track unique id
        asset_upload_type (str | None): Asset upload type
        original_filename (str): Original filename
        user_id (str): User id

    Returns:
        str: Unique filename for the asset upload.
    """
    # TODO: CDAM-3825: Remove this derivation once asset_upload_type is required
    if not asset_upload_type:
        asset_upload_type = (
            asset_upload_types_constants.STEREO
            if track_unique_id
            else asset_upload_types_constants.STATIC_ARTWORK
        )

    # Verify asset_upload_type + track_unique_id combination is valid
    is_image_type = asset_upload_type in asset_upload_types_constants.IMAGE_UPLOAD_TYPES
    if is_image_type and track_unique_id:
        raise InvalidAssetUploadType(
            f"Image upload type '{asset_upload_type}' cannot have a track_unique_id"
        )
    if not is_image_type and not track_unique_id:
        raise InvalidAssetUploadType(
            f"Audio upload type '{asset_upload_type}' requires a track_unique_id"
        )

    asset_upload_type_id = asset_upload_type_model.resolve_asset_upload_type_id(
        asset_upload_type
    )

    # Verify product exists and can accept uploads, and get additional fields
    additional_fields_for_asset_upload = _get_additional_fields_for_asset_upload(
        product_id, asset_upload_type
    )

    # Verify track exists and is on product
    if track_unique_id:
        track_response = ows_track_model.get_track_by_id(track_unique_id)
        if track_response["product_id"] != product_id:
            raise TrackDoesNotBelongToProduct("Track does not belong to the product")

    filename = _generate_filename()

    # Uploads are stored as arbitrary data. We validate them to see what
    # the actual content is.
    s3_multipart_upload = s3_connector.get_s3_client().create_multipart_upload(
        Bucket=config.RAW_ASSETS_BUCKET_NAME,
        Key=filename,
        ContentType="application/octet-stream",
    )

    asset_upload_model.create_asset_upload_with_context(
        user_id=user_id,
        asset_upload_type_id=asset_upload_type_id,
        token=s3_multipart_upload["UploadId"],
        filename=filename,
        product_id=product_id,
        upc=additional_fields_for_asset_upload["upc"],
        track_unique_id=track_unique_id,
        is_correction=additional_fields_for_asset_upload["is_correction"],
        original_filename=original_filename,
    )

    return filename


def get_presigned_urls_for_asset_upload(
    *,
    part_numbers: list[int],
    filename: str,
    user_id: str,
) -> dict[str, Any]:
    """Get presigned urls for asset upload.

    Args:
        part_numbers (list): List of upload part numbers generate presigned urls for.
        filename (str): Name of the file to upload.
        user_id (str): User id.

    Returns:
        dict: Contains information for uploading.
    """
    get_asset_upload_response = asset_upload_model.get_asset_upload(
        filename,
        user_id=user_id,
        api_version=api.API_VERSION_V2,
    )

    get_last_asset_status_response = asset_status_model.get_last_asset_status(
        get_asset_upload_response["id"]
    )

    if (
        get_last_asset_status_response["status"]
        != asset_status_constants.STATUS_UPLOADING
    ):
        raise WrongAssetStatus(
            "asset_upload not in state that allows generation of presigned urls for uploads."
        )

    return {
        "part_number_to_presigned_url": {
            part_number: s3_connector.get_s3_client(
                use_accelerate_endpoint=True
            ).generate_presigned_url(
                ClientMethod="upload_part",
                Params={
                    "Bucket": config.RAW_ASSETS_BUCKET_NAME,
                    "Key": filename,
                    "UploadId": get_asset_upload_response["token"],
                    "PartNumber": part_number,
                },
                ExpiresIn=config.ASSET_UPLOAD_PRESIGNED_URL_EXPIRES_IN_SECONDS,
            )
            for part_number in part_numbers
        },
    }


def complete_multipart_upload(
    *,
    filename: str,
    parts: list[dict[str, Any]],
    user_id: str,
) -> None:
    """
    Complete a multipart upload.

    Args:
        filename (str): Name of the file to upload.
        parts (list): Upload part info.
        user_id (str): User id.
    """
    get_asset_upload_response = asset_upload_model.get_asset_upload(
        filename,
        user_id=user_id,
        api_version=api.API_VERSION_V2,
    )

    get_last_asset_status_response = asset_status_model.get_last_asset_status(
        get_asset_upload_response["id"]
    )

    if (
        get_last_asset_status_response["status"]
        != asset_status_constants.STATUS_UPLOADING
    ):
        raise WrongAssetStatus(
            "asset_upload not in state that allows completion of upload."
        )

    try:
        s3_connector.get_s3_client().complete_multipart_upload(
            Bucket=config.RAW_ASSETS_BUCKET_NAME,
            Key=filename,
            UploadId=get_asset_upload_response["token"],
            MultipartUpload={
                "Parts": [
                    {
                        "PartNumber": part["part_number"],
                        "ETag": part["etag"],
                    }
                    for part in parts
                ],
            },
        )
    except ClientError as client_error:
        client_error_code = client_error.response["Error"]["Code"]
        if client_error_code in (
            "EntityTooSmall",
            "InvalidPart",
            "InvalidPartOrder",
            "NoSuchUpload",
        ):
            raise S3MultipartUploadFailed(str(client_error)) from client_error

        raise

    try:
        asset_status_model.create_asset_status(
            asset_upload_id=get_asset_upload_response["id"],
            status=asset_status_constants.STATUS_UPLOAD_COMPLETE,
            status_time=datetime.now(timezone.utc),
        )
    except SQLAlchemyError as e:
        sentry_sdk.capture_exception(e)


def get_upload_permission(
    user_id: str, duration: int, asset_type: str
) -> dict[str, Any]:
    """Function that generates data required for asset upload.

    Args:
        user_id (str): User id
        duration (int): Number of seconds upload permissions should live.
        asset_type (str): Asset type (audio or image).
    Returns:
        dict: Contains information for uploading.
    """
    filename = _generate_filename()
    credentials = token.get_s3_token(filename, duration)
    s3_token = credentials["token"]

    # TODO: Temporarily deriving asset_upload_type
    # asset_type can only be audio or image at this point
    asset_upload_type = (
        asset_upload_types_constants.STEREO
        if asset_type == asset_upload_constants.ASSET_TYPE_AUDIO
        else asset_upload_types_constants.STATIC_ARTWORK
    )

    asset_upload_type_id = asset_upload_type_model.resolve_asset_upload_type_id(
        asset_upload_type
    )

    asset_upload = asset_upload_model.create_asset_upload(
        user_id, filename, s3_token, asset_upload_type_id, api.API_VERSION_V2
    )
    if not asset_upload:
        raise AssetUploadError("Failed to create an asset upload record")

    return {
        "bucket": config.RAW_ASSETS_BUCKET_NAME,
        "filename": filename,
        "credentials": credentials,
    }


def get_entity_upload_permission(
    duration: int, entity: str, entity_id: int
) -> dict[str, Any]:
    """Function that generates data required for asset upload.

    Args:
        duration (int): Number of seconds upload permissions should live.
        entity (string): Name of the entity/folder on S3 that stores images.
        entity_id (int): either artist_id or vendor_id.
    Returns:
        dict: Contains information for uploading.
    """
    entity_hash = hashlib.md5()
    entity_hash.update(str(entity_id).encode("utf8"))
    filename = "{filename}.{extension}".format(
        filename=str(entity_hash.hexdigest()),
        extension=asset_types.TYPE_FILE_JPG.lower(),
    )

    path = {}
    if entity == s3.VENDOR_ENTITY:
        logo_path = "{images}/{entity}/{logo_folder}".format(
            images=s3.IMAGE_ASSET_BUCKET,
            entity=entity,
            logo_folder=s3.VENDOR_LOGO_FOLDER,
        )
        icon_path = "{images}/{entity}/{icon_folder}".format(
            images=s3.IMAGE_ASSET_BUCKET,
            entity=entity,
            icon_folder=s3.VENDOR_ICON_FOLDER,
        )
        path = {"logo": logo_path, "icon": icon_path}
    elif entity == s3.ARTIST_ENTITY:
        photos_path = "{images}/{entity}/{photos_folder}".format(
            images=s3.IMAGE_ASSET_BUCKET,
            entity=entity,
            photos_folder=s3.ARTIST_PHOTOS_FOLDER,
        )
        thumb_path = "{images}/{entity}/{thumb_folder}".format(
            images=s3.IMAGE_ASSET_BUCKET,
            entity=entity,
            thumb_folder=s3.ARTIST_THUMB_FOLDER,
        )
        web_path = "{images}/{entity}/{web_folder}".format(
            images=s3.IMAGE_ASSET_BUCKET, entity=entity, web_folder=s3.ARTIST_WEB_FOLDER
        )
        path = {"photos": photos_path, "thumb": thumb_path, "web": web_path}

    credentials = token.get_s3_entity_token(path, filename, duration)

    return {
        "bucket": config.ASSET_STORAGE_BUCKET_NAME,
        "path": path,
        "filename": filename,
        "credentials": credentials,
    }