import os
import boto3
import uuid
import mysql.connector
import json
from secrets_manager.python_ext import PythonSecretsManager

ENVIRONMENT = "prod"
SERVICE_NAME = 'distro-scripts-create-asset-finals'
RAW_ASSETS_S3_BUCKET = f"{ENVIRONMENT}-orcd-raw-assets"
ASSETS_TRANSCODING_V2_STATE_MACHINE_ARN = f"arn:aws:states:us-east-1:437795906767:stateMachine:{ENVIRONMENT}-assets-transcoding-v2"

secrets_manager_client = PythonSecretsManager(
    environment=ENVIRONMENT,
    service_name=SERVICE_NAME,
)

mysql_connection = mysql.connector.connect(
    host=os.environ.get("OWS_ASSETS_DB_HOST"),
    user=os.environ.get("OWS_ASSETS_DB_USER"),
    password=secrets_manager_client.get_cred("OWS_ASSETS_DB_PASSWORD"),
    database="ows_assets",
)
s3_client = boto3.client("s3")
sfn_client = boto3.client("stepfunctions")


def generate_unique_filename():
    """Generate a unique filename."""
    return str(uuid.uuid4()).replace("-", "_")


def get_asset_upload(asset_upload_id):
    """Get an asset_upload record by asset_upload_id."""
    with mysql_connection.cursor(dictionary=True) as cursor:
        cursor.execute(
            """
            SELECT *
            FROM asset_upload
            WHERE id = %(asset_upload_id)s
            """,
            {"asset_upload_id": asset_upload_id},
        )
        return cursor.fetchone()


def get_asset_uploads_by_filename(filename):
    """Get asset_upload records by filename."""
    with mysql_connection.cursor(dictionary=True) as cursor:
        cursor.execute(
            """
            SELECT *
            FROM asset_upload
            WHERE filename = %(filename)s
                AND api_version = 2
            """,
            {"filename": filename},
        )
        return cursor.fetchall()


def update_asset_upload_filename(asset_upload_id, new_filename):
    """Update asset_upload filename by asset_upload_id."""
    with mysql_connection.cursor(dictionary=True) as cursor:
        cursor.execute(
            """
            UPDATE asset_upload
            SET filename = %(new_filename)s
            WHERE id = %(asset_upload_id)s
            """,
            {
                "asset_upload_id": asset_upload_id,
                "new_filename": new_filename,
            },
        )
        mysql_connection.commit()


def insert_upload_complete_asset_status(asset_upload_id):
    """Insert upload_complete asset_status record for a given asset_upload_id."""
    with mysql_connection.cursor(dictionary=True) as cursor:
        cursor.execute(
            """
            INSERT INTO asset_status (asset_upload_id, status, message, status_time)
            VALUES (%(asset_upload_id)s, 'upload_complete', '{}', NOW())
            """,
            {"asset_upload_id": asset_upload_id},
        )
        mysql_connection.commit()


def delete_asset_final_records_for_asset_upload(asset_upload_id):
    """Delete records from the asset_final table for a given asset_upload_id."""
    with mysql_connection.cursor(dictionary=True) as cursor:
        cursor.execute(
            "DELETE FROM asset_final WHERE asset_upload_id = %(asset_upload_id)s",
            {"asset_upload_id": asset_upload_id},
        )
        mysql_connection.commit()


def trigger_assets_transcoding_v2_state_machine(filename):
    """Trigger the ENV-assets-transcoding-v2 state machine."""
    sfn_client.start_execution(
        stateMachineArn=ASSETS_TRANSCODING_V2_STATE_MACHINE_ARN,
        name=f"create-asset-finals_{uuid.uuid4()}",
        input=json.dumps(
            {
                "detail-type": "Object Created",
                "detail": {
                    "bucket": {"name": RAW_ASSETS_S3_BUCKET},
                    "object": {"key": filename},
                }
            }
        ),
    )


def copy_s3_object_with_metadata(
    *,
    source_filename,
    destination_filename,
    metadata,
):
    """Copy an S3 object to a new key with updated metadata, preserving the original content type."""
    original_object = s3_client.head_object(
        Bucket=RAW_ASSETS_S3_BUCKET, Key=source_filename
    )
    content_type = original_object["ContentType"]

    s3_client.copy_object(
        CopySource={"Bucket": RAW_ASSETS_S3_BUCKET, "Key": source_filename},
        Bucket=RAW_ASSETS_S3_BUCKET,
        Key=destination_filename,
        ContentType=content_type,
        Metadata=metadata,
        MetadataDirective="REPLACE",
    )


def get_full_asset_upload_filename(filename_without_extension):
    """Get the full asset_upload filename with extension."""
    response = s3_client.list_objects_v2(
        Bucket=RAW_ASSETS_S3_BUCKET, Prefix=filename_without_extension
    )

    if "Contents" not in response:
        raise Exception(f"No files found with prefix {filename_without_extension}")

    if len(response["Contents"]) > 1:
        raise Exception(
            f"Multiple files found with prefix {filename_without_extension}"
        )

    return response["Contents"][0]["Key"]


def create_asset_finals_for_asset_upload(
    *,
    asset_upload_id,
    s3_metadata_to_set=None,
    should_clear_s3_metadata=False,
):
    """Create asset_final assets and records for a given asset_upload_id."""
    asset_upload = get_asset_upload(asset_upload_id)
    is_filename_unique = (
        len(get_asset_uploads_by_filename(asset_upload["filename"])) == 1
    )
    existing_filename = get_full_asset_upload_filename(asset_upload["filename"])
    existing_filename_extension = os.path.splitext(existing_filename)[1]

    insert_upload_complete_asset_status(asset_upload_id)
    delete_asset_final_records_for_asset_upload(asset_upload_id)

    if is_filename_unique:
        if not s3_metadata_to_set and not should_clear_s3_metadata:
            # no metadata to set, so we can just trigger the transcoding state machine
            trigger_assets_transcoding_v2_state_machine(existing_filename)
            return

        # to update s3 object metadata, the object is copied to itself with the new metadata
        destination_filename = existing_filename
    else:
        # filename is not unique, so we need to copy the object to a new filename that is unique
        new_filename = generate_unique_filename()
        update_asset_upload_filename(asset_upload_id, new_filename)
        destination_filename = f"{new_filename}{existing_filename_extension}"

    copy_s3_object_with_metadata(
        source_filename=existing_filename,
        destination_filename=destination_filename,
        metadata=s3_metadata_to_set,
    )


def has_product_id_set_in_metadata(asset_upload_id):
    """Check if product_id is set in the S3 object metadata."""
    asset_upload = get_asset_upload(asset_upload_id)
    asset_upload_filename = get_full_asset_upload_filename(asset_upload["filename"])
    response = s3_client.head_object(Bucket=RAW_ASSETS_S3_BUCKET, Key=asset_upload_filename)
    return "product_id" in response.get("Metadata", {})


if __name__ == "__main__":
    inputs = json.loads(os.environ["CREATE_ASSET_FINALS_INPUTS"])
    asset_upload_ids = inputs["asset_upload_ids"]
    # example of s3_metadata_to_set:
    # s3_metadata_to_set = {
    #     'bypass_image_dimension_validation': '1',
    #     'bypass_image_mode_validation': '1',
    # }
    s3_metadata_to_set = inputs.get("s3_metadata_to_set", {})

    # preventing setting this metadata, particularly product_id, allows this to flow through the new acknowledge flow
    # https://github.com/theorchard/lambda-assets/blob/2a7ea87d82730dfdb95ad4fb8ea0bc6c945fde24/lambda/acknowledge/src/app.py#L73
    # https://github.com/theorchard/ows-assets/blob/c4fa3e596a2bfda0e119ac17a5d06832bee8c67b/assets/models/asset_upload.py#L423
    restricted_s3_metadata = {
        "asset_type",
        "product_id",
        "original_filename",
        "upc",
        "track_unique_id",
        "is_correction",
    }
    if any(
        key in s3_metadata_to_set
        for key in restricted_s3_metadata
    ):
        raise Exception(
            f"s3_metadata_to_set should not include: {', '.join(restricted_s3_metadata)}"
        )

    assert asset_upload_ids, "asset_upload_ids is empty"

    for asset_upload_id in asset_upload_ids:
        create_asset_finals_for_asset_upload(
            asset_upload_id=asset_upload_id,
            s3_metadata_to_set=s3_metadata_to_set,
            # clearing metadata, particularly product_id, from metadata allows this to flow through the new acknowledge flow
            should_clear_s3_metadata=has_product_id_set_in_metadata(asset_upload_id)
        )

    mysql_connection.close()