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

import asyncio
import json
import unicodedata
import uuid
from typing import Any

from fastapi import HTTPException
from pydantic import UUID4

from product_staging import config
from product_staging.api import datasources
from product_staging.api.schemas.assets_upload import (
    CompleteAssetUploadRequest,
    CompleteAssetUploadResponse,
)
from product_staging.api.schemas.bulk_session import (
    AssetStatus,
    UpdateBulkSessionRequest,
)
from product_staging.api.schemas.bulk_session_asset_file import (
    FileStatus,
    UpdateBulkSessionAssetRequest,
)
from product_staging.connectors import db
from product_staging.constants.error import ERROR_MESSAGE_NO_BULK_SESSION
from product_staging.logic import bulk_session as bulk_session_logic
from product_staging.logic.utils import s3, utils
from product_staging.models import bulk_session as bulk_session_model
from product_staging.models import (
    bulk_session_asset_file as bulk_session_asset_file_model,
)


@db.db_session_wrap
async def create_bulk_assets_upload_and_get_presigned_urls(
    assets: list,
    bulk_session_id: UUID4,
    identity_uuid: UUID4,
    cloud_transfer_job_id: str | None = None,
    session=None,
):
    """Creates an assets bulk upload in S3, saves related assets to the
    database and generate presigned urls.
    """
    bulk_session = await bulk_session_logic.get_bulk_session(
        bulk_session_id, identity_uuid=identity_uuid, enrich=False, session=session
    )
    if not bulk_session:
        raise HTTPException(
            status_code=404,
            detail=ERROR_MESSAGE_NO_BULK_SESSION,
        )

    await bulk_session_model.update_bulk_session(
        bulk_session_id,
        identity_uuid,
        UpdateBulkSessionRequest(
            asset_status=AssetStatus.transferring
            if cloud_transfer_job_id
            else AssetStatus.incomplete
        ),
    )

    metadata_json_file = await s3.get_file(
        await bulk_session.validated_metadata_json_file()
    )
    metadata_json = json.loads((await metadata_json_file["Body"].read()))

    asset_dict = {}
    for asset in assets:
        asset_dict[unicodedata.normalize("NFC", asset.original_filename)] = {
            "parts": asset.parts
        }

    for product in metadata_json:
        artwork = product["product_info"]["product"]["artwork"]
        if artwork:
            filename = unicodedata.normalize("NFC", artwork["filename"])
            if filename in asset_dict:
                artwork["bucket"] = config.OWS_PRODUCT_STAGING_S3_BUCKET
                artwork["key"] = f"assets/{str(uuid.uuid4())}"
                asset_dict[filename]["key"] = artwork["key"]
                asset_dict[filename]["type"] = "image"

        for track in product["product_info"]["tracks"]:
            audio = track["asset"]
            if audio:
                filename = unicodedata.normalize("NFC", audio["filename"])
                if filename in asset_dict:
                    audio["bucket"] = config.OWS_PRODUCT_STAGING_S3_BUCKET
                    audio["key"] = f"assets/{str(uuid.uuid4())}"
                    asset_dict[filename]["key"] = audio["key"]
                    asset_dict[filename]["type"] = "audio"

    filtered_asset_dict = {
        name: info for name, info in asset_dict.items() if "key" in info
    }
    unidentified_assets = list(asset_dict.keys() - filtered_asset_dict.keys())

    s3_client = datasources.get_s3_client()

    multipart_upload_results = await asyncio.gather(
        *[
            s3_client.create_multipart_upload(
                Bucket=config.OWS_PRODUCT_STAGING_S3_BUCKET,
                Key=info["key"],
                Metadata={
                    "bulk_session_id": str(bulk_session_id),
                    "identity_uuid": str(identity_uuid),
                },
            )
            for _, info in filtered_asset_dict.items()
        ]
    )

    save_db_results = await bulk_session_asset_file_model.save_multiple(
        [
            bulk_session_asset_file_model.init_bulk_session_asset_file(
                bulk_session_id=bulk_session_id,
                upload_token=multipart_upload_results[i]["UploadId"],
                original_filename=filename,
                s3_filename=multipart_upload_results[i]["Key"],
                identity_uuid=identity_uuid,
                file_type=info["type"],
                cloud_transfer_job_id=cloud_transfer_job_id,
            )
            for i, (filename, info) in enumerate(filtered_asset_dict.items())
        ]
    )

    await s3.write_file_stream(
        await bulk_session.validated_metadata_json_file(),
        utils.list_to_stream(metadata_json),
        metadata={
            "bulk_session_id": str(bulk_session_id),
            "identity_uuid": str(identity_uuid),
        },
    )

    presigned_urls_results = await asyncio.gather(
        *[
            s3.get_multipart_upload_presigned_urls(
                parts=info["parts"],
                s3_filename=save_db_results[i]["s3_filename"],
                upload_token=save_db_results[i]["upload_token"],
            )
            for i, info in enumerate(filtered_asset_dict.values())
        ]
    )

    return {
        "bulk_session_id": str(bulk_session_id),
        "assets": [
            {
                **save_db_results[i],
                "presigned_urls": presigned_urls_results[i],
            }
            for i in range(len(filtered_asset_dict))
        ],
        "unidentified_assets": unidentified_assets,
    }


async def complete_asset_upload(
    request: CompleteAssetUploadRequest,
    s3_filename: str,
    file: dict[str, Any],
    bulk_session: dict[str, Any],
) -> CompleteAssetUploadResponse:
    """Complete asset upload request."""
    await s3.complete_multipart_upload(
        s3_filename=s3_filename,
        upload_token=file["upload_token"],
        parts=request.parts,
    )

    await bulk_session_asset_file_model.update_bulk_session_asset_file(
        s3_filename=s3_filename,
        update_req=UpdateBulkSessionAssetRequest(file_status=FileStatus.validating),
    )

    return CompleteAssetUploadResponse.model_validate(
        {
            "bulk_session_id": bulk_session["bulk_session_id"],
            "bulk_session_asset_file_id": file["bulk_session_asset_file_id"],
        }
    )