"""Video upload application on modal."""

import concurrent
import math
import json
import logging
import os
import time
from collections import namedtuple
from typing import Literal

import boto3
import ffmpeg
import modal
import requests
from datadog import initialize, statsd

from suno_utils.worker.video_utils.video_sprite import process_video_sprite
from suno_utils.utils.image_extraction import ImageExtractionError
from suno_utils.worker.schema import QueueItem
from suno_utils.worker.settings import s3_client
from suno_utils.worker.utils import retry_decorator
from suno_utils.worker.modal_base import get_modal_base_image

logger = logging.getLogger(__name__)
logging.basicConfig()
logger.setLevel(logging.INFO)

############## CHANGE THESE ##############
DEPLOYMENT_TYPE = "dev"

##########################################
DEPLOYMENT_TYPES = {"dev", "prod"}
assert DEPLOYMENT_TYPE in DEPLOYMENT_TYPES
# orchestrator talks to studio api
APP_NAME = f"upload-video-{DEPLOYMENT_TYPE}"
assert APP_NAME.endswith(DEPLOYMENT_TYPE)
MOUNT_PATH = "/suno/models"
UPLOADS_S3_BUCKET = "suno-data-uploads"
RAW_UPLOADS_S3_BUCKET = "suno-uploads"
retry_s3_download = retry_decorator(3, wait_seconds=20)(s3_client.download_file)
retry_s3_upload = retry_decorator(3, wait_seconds=20)(s3_client.upload_file)
Dimensions = namedtuple("Dimensions", ["width", "height"])
PreviewFormat = Literal["gif", "mp4"]
DEFAULT_PREVIEW_FORMAT: PreviewFormat = "mp4"
PREVIEW_DURATIONS = {"gif": 2, "mp4": 10}

# for maximum video upload duration
MAX_SCENES_UPLOAD_DURATION = 31
MAX_COVER_VIDEO_UPLOAD_DURATION = 11
MAX_HOOK_VIDEO_UPLOAD_DURATION = 121
# arbitrary number, just for better experience
MIN_UPLOAD_DURATION = 1

aws_secret = modal.Secret.from_name("studio-aws")
SECRETS = [
    aws_secret,
    modal.Secret.from_dict(
        {
            "DD_SITE": "datadoghq.com",
            "DD_ENV": DEPLOYMENT_TYPE,
            "DD_SERVICE": DEPLOYMENT_TYPE,
            "DD_LOGS_ENABLED": "false",
            "DD_TRACE_ENABLED": "true" if DEPLOYMENT_TYPE == "dev" else "false",
        },
    ),
    modal.Secret.from_name("datadog-metrics"),
    modal.Secret.from_name("api-callback-token"),
    modal.Secret.from_name("hive-secret"),
]


def download_model_wrapper_b() -> None:
    pass


base_image = get_modal_base_image().add_local_python_source("suno_utils", copy=False)

app = modal.App(APP_NAME, image=base_image)


# TODO experiment with the parameters
@app.cls(
    cpu=1,
    secrets=SECRETS,
    timeout=240,
    min_containers=1,
    scaledown_window=600,
    retries=modal.Retries(
        max_retries=1,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    cloud="aws",
    region="us-east",
    buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1,
)
@modal.concurrent(max_inputs=4)
class VideoUploadStub:
    """Video upload stub and functions."""

    def __init__(self):
        """Set up VideoUploadStub."""
        self.modal_f_image_description = modal.Cls.lookup(
            f"image-extraction-{DEPLOYMENT_TYPE}",
            "ImageExtractionStub",
        )().generate_image_description

        options = {"statsd_host": "127.0.0.1", "statsd_port": 8125}
        initialize(**options)

    @modal.method()
    def trigger_mediaconvert_job(self, queue_item_json: str):
        """Loop video for audio duration, upload to mediaconvert source bucket"""
        item = QueueItem(**json.loads(queue_item_json))
        audio_s3_id = item.metadata.get("audio_s3_id", None)
        video_key = item.metadata.get("video_key", None)
        video_duration = item.metadata.get("video_duration", None)
        audio_duration = item.metadata.get("audio_duration", None)

        # download video_key from s3
        temp_input_path = f"/tmp/tmp_{video_key.split('/')[-1]}"
        retry_s3_download(UPLOADS_S3_BUCKET, f"studio/uploads/{video_key}.mp4", temp_input_path)
        print(f"Downloaded video from {video_key}. ID {item.id}")

        # Calculate how many times we need to loop the video
        if not all([video_duration, audio_duration]):
            print(f"Missing required duration parameters. ID {item.id}")
            return {"success": False, "error": "Missing duration parameters"}

        # Calculate loop count (rounded up)
        loop_count = math.ceil(audio_duration / video_duration)

        # Create output path for the looped video
        looped_video_path = f"/tmp/looped_{video_key.split('/')[-1]}.mp4"

        try:
            # Create a file with repeated inputs
            input_list_file = "/tmp/input_list.txt"
            with open(input_list_file, "w") as f:
                for _ in range(loop_count):
                    f.write(f"file '{temp_input_path}'\n")

            # Use ffmpeg concat demuxer to loop the video
            (
                ffmpeg.input(input_list_file, format="concat", safe=0)
                .output(looped_video_path, c="copy", t=str(audio_duration))
                .run(capture_stdout=True, capture_stderr=True)
            )

            # Upload to mediaconvert source bucket
            first_two = item.id[:2]
            second_two = item.id[2:4]
            dest_key = f"{first_two}/{second_two}/{item.id}.mp4"

            # Upload to the mediaconvert source bucket
            s3_client = boto3.client("s3")
            s3_client.upload_file(
                looped_video_path,
                "suno-media-sour",  # Mediaconvert source bucket
                dest_key,
                ExtraArgs={
                    "Metadata": {
                        "audio-file-bucket": "suno-data-uploads",
                        "audio-file-key": f"studio/uploads/{audio_s3_id}.mp3",
                        "environment": DEPLOYMENT_TYPE,
                        "callback-url": item.callback_url,
                    }
                },
            )

            print(f"Successfully uploaded looped video to mediaconvert source bucket. ID {item.id}")
            return {"success": True, "dest_key": dest_key}

        except ffmpeg.Error as e:
            print(
                f"FFmpeg error while looping video: {e.stderr.decode() if e.stderr else str(e)}. ID {item.id}"
            )
            return {"success": False, "error": "FFmpeg error"}
        except Exception as e:
            print(f"Error processing video loop: {str(e)}. ID {item.id}")
            return {"success": False, "error": str(e)}

    @modal.method()
    def process_uploaded_video(self, queue_item_json: str):
        """Transcode, moderate and upload video to S3, then generate the description."""
        t0 = time.time()
        item = QueueItem(**json.loads(queue_item_json))
        upload_id = item.id
        video_cover_clip_id = item.metadata.get("clip_id", None)
        is_clip_video_cover = video_cover_clip_id is not None or item.metadata.get(
            "is_video_cover", False
        )
        require_video_sprite = item.metadata.get("require_video_sprite", False)
        fail_task_on_moderation = item.metadata.get("fail_task_on_moderation", True)
        video_upload_type = item.metadata.get("video_upload_type", None)
        title = item.title
        if not title:
            title = "Uploaded Video"
        # this is the key on s3 that maps to the uploaded video file
        upload_key = item.metadata.get("upload_key", "")
        print(f"upload_key: {upload_key} \n Upload id {upload_id}")

        # Define the S3 path for the uploaded file
        try:
            s3_client.get_object(Bucket=RAW_UPLOADS_S3_BUCKET, Key=upload_key)
        except:
            print(
                f"File {upload_id} doesn't exist on s3. Did the upload finish? Upload id {upload_id}\n"
            )
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 0,
                    "error_type": "upload_failure_not_found",
                    "error_message": "Can't fetch the uploaded video.",
                },
                blocking=True,
            )
            return {}

        # Download the file from S3 to a temporary location
        temp_input_path = f"/tmp/tmp_{upload_key.split('/')[-1]}"
        retry_s3_download(RAW_UPLOADS_S3_BUCKET, upload_key, temp_input_path)
        print(f"Downloaded {upload_key} in {time.time() - t0:.3f}s  Upload id {upload_id}\n")
        result = {
            "success": False,
            "duration": -1,
            "s3_id": None,
            "image_s3_ids": [],
            "video_dimensions": None,
        }

        # Read video metadata, fetch duration and dimensions
        try:
            probe = ffmpeg.probe(temp_input_path)

            video_dimensions = self._get_video_dimensions(probe)
            result["video_dimensions"] = video_dimensions

            result["duration"] = float(probe["format"]["duration"])
        except Exception as e:
            print(f"Error getting duration: {e} Upload id {upload_id}\n")
            return result

        max_duration = self._get_max_duration(video_upload_type, is_clip_video_cover)
        if result["duration"] > max_duration or result["duration"] < MIN_UPLOAD_DURATION:
            print(
                f"Video duration {result['duration']} is not within the allowed range. Upload id {upload_id}\n"
            )
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 0,
                    "error_type": "upload_failure_video_length",
                    "error_message": f"Video too {'long' if result['duration'] > max_duration else 'short'}",
                },
                blocking=True,
            )
            return result

        duration = int(result["duration"])
        video_dimensions = result["video_dimensions"]
        thumbnail_dimensions = self.get_thumbnail_dimensions(video_dimensions)
        # get duration bucket by 10s
        duration_bucket = (
            (duration + 9) // 10 * 10
        )  # round up to nearest 10 - buckets will be 10, 20, 30

        with concurrent.futures.ThreadPoolExecutor() as executor:
            # Submit the functions to the executor
            future_image_lyrics = None
            future_process_snapshot = None
            future_process_preview = None
            task = "scene"
            if video_upload_type == "video_hook":
                task = "hook"
                future_process_snapshot = executor.submit(
                    self._process_cover_image_from_video, upload_id, temp_input_path, video_dimensions
                )
            elif is_clip_video_cover:
                task = "cover"
                preview_format = item.metadata.get("preview_format", None) or DEFAULT_PREVIEW_FORMAT
                future_process_snapshot = executor.submit(
                    self._process_cover_image_from_video, upload_id, temp_input_path, video_dimensions
                )
                future_process_preview = executor.submit(
                    self._process_cover_preview_from_video,
                    upload_id,
                    temp_input_path,
                    thumbnail_dimensions,
                    duration,
                    preview_format,
                )
            else:  # scenes
                future_image_lyrics = executor.submit(
                    self._process_image_to_text,
                    upload_id,
                    temp_input_path,
                    duration,
                    duration_bucket,
                    thumbnail_dimensions,
                )

            maintain_audio = video_upload_type == "video_hook"
            future_transcode = executor.submit(
                self._process_video_transcode,
                upload_id,
                upload_key,
                temp_input_path,
                duration_bucket,
                maintain_audio,
            )
            future_nsfw_check = executor.submit(
                self._detect_nsfw, RAW_UPLOADS_S3_BUCKET, upload_key, upload_id, duration_bucket
            )

            if require_video_sprite:
                future_video_sprites = executor.submit(
                    process_video_sprite,
                    upload_id,
                    temp_input_path,
                    duration,
                    duration_bucket,
                    video_dimensions,
                    retry_s3_upload,
                )

        with statsd.timed(
            "video_processing",
            tags=[
                f"env:{DEPLOYMENT_TYPE}",
                f"video_upload_type:{video_upload_type}",
                "sub_task:transcode",
            ],
        ):
            # Wait for the results
            transcode_result = future_transcode.result()

        with statsd.timed(
            "video_processing",
            tags=[
                f"env:{DEPLOYMENT_TYPE}",
                f"video_upload_type:{video_upload_type}",
                "sub_task:moderation",
            ],
        ):
            # Handle NSFW check with error handling
            try:
                nsfw_result = future_nsfw_check.result()
            except Exception as e:
                logger.error(f"NSFW check failed for upload {upload_id}: {e}")
                # Default to safe (not NSFW) if check fails
                nsfw_result = False

        with statsd.timed(
            "video_processing",
            tags=[f"env:{DEPLOYMENT_TYPE}", f"video_upload_type:{video_upload_type}", "sub_task:sprite"],
        ):
            video_sprite_result = future_video_sprites.result() if require_video_sprite else None

        image_lyrics_result = None
        snapshot_process_result = None
        image_s3_id = None
        gif_s3_id = None
        video_preview_s3_id = None
        description = None

        # wait for image-to-text - for scenes
        if future_image_lyrics is not None:
            with statsd.timed(
                "video_processing",
                tags=[
                    f"env:{DEPLOYMENT_TYPE}",
                    f"video_upload_type:{video_upload_type}",
                    "sub_task:image_lyrics",
                ],
            ):
                image_lyrics_result = future_image_lyrics.result()
                image_s3_id = image_lyrics_result["image_s3_id"]
                description = image_lyrics_result["lyrics"]
        # wait for processing snapshot / cover image - for clip video cover
        else:
            with statsd.timed(
                "video_processing",
                tags=[
                    f"env:{DEPLOYMENT_TYPE}",
                    f"video_upload_type:{video_upload_type}",
                    "sub_task:snapshot",
                ],
            ):
                if future_process_snapshot is not None:
                    snapshot_process_result = future_process_snapshot.result()
                    image_s3_id = snapshot_process_result["image_s3_id"]
                if future_process_preview is not None:
                    preview_process_result = future_process_preview.result()
                    gif_s3_id = preview_process_result.get("gif_s3_id")
                    video_preview_s3_id = preview_process_result.get("mp4_s3_id")

        if video_sprite_result and video_sprite_result.success:
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 1,
                    "status": "passed_video_sprite_generation",
                },
                blocking=True,
            )
        elif video_sprite_result and not video_sprite_result.success:
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 0,
                    "error_type": "upload_failure_sprite",
                    "error_message": video_sprite_result.error_message,
                },
                blocking=True,
            )
            return {}  # stop processing

        if nsfw_result:
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 0,
                    "error_type": "upload_failure_nsfw",
                    "error_message": "Video is not safe for work.",
                },
                blocking=True,
            )
            if fail_task_on_moderation:
                return {}
        else:
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 1,
                    "status": "passed_nsfw_check",
                },
                blocking=True,
            )

        if transcode_result["success"]:
            # send a notification for success if get here
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 1,
                    "status": "passed_video_processing",
                },
                blocking=True,
            )
        else:
            if not transcode_result["s3_id"]:
                item.notify_progress(
                    {
                        "id": item.id,
                        "request_id": upload_id,
                        "type": "upload_video",
                        "ok": 0,
                        "error_type": "upload_failure_decode_video",
                        "error_message": "Can't parse uploaded video. Source is corrupted.",
                    },
                    blocking=True,
                )
                return {}

        if image_lyrics_result is not None and not image_lyrics_result["success"]:
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 0,
                    "error_type": "upload_failure_extract_lyrics",
                    "error_message": "Can't extract lyrics from video.",
                },
                blocking=True,
            )
            return {}

        if snapshot_process_result is not None and not snapshot_process_result["success"]:
            item.notify_progress(
                {
                    "id": item.id,
                    "request_id": upload_id,
                    "type": "upload_video",
                    "ok": 0,
                    "error_type": "upload_failure_snapshot",
                    "error_message": "Can't extract cover snapshot from video.",
                },
                blocking=True,
            )

        print(f"Dimensions are {video_dimensions} Upload id {upload_id}\n")
        # send back the information
        item.notify_progress(
            {
                "id": item.id,
                "request_id": upload_id,
                "type": "upload_video",
                "s3_id": transcode_result["s3_id"],
                "image_s3": image_s3_id,
                "gif_s3": gif_s3_id,
                "mp4_preview_s3": video_preview_s3_id,
                "title": title,
                "description": description,
                "duration": result["duration"],
                "is_description_lyrics": True,
                "video_cover_clip_id": video_cover_clip_id,
                "video_width": video_dimensions.width if video_dimensions else None,
                "video_height": video_dimensions.height if video_dimensions else None,
                "video_sprite": video_sprite_result.sprite_schema.to_dict()
                if video_sprite_result
                else None,
                "is_nsfw": nsfw_result,
            },
            blocking=True,
        )
        statsd.distribution(
            "video_processing.total_time",
            time.time() - t0,
            tags=[
                f"duration_bucket:{duration_bucket}",
                f"task:{task}",
                f"video_upload_type:{video_upload_type}",
            ],
        )
        print(f"Done with {upload_id}, {time.time() - t0:.3f}s \n Upload id {upload_id}")
        if image_lyrics_result:
            print(
                f"Generated lyrics for {upload_id}, {image_lyrics_result['lyrics']} Upload id {upload_id}\n"
            )
        else:
            print(f"Skipped generating lyrics for clip video cover. Upload id {upload_id}\n")
        return {}

    def _get_max_duration(self, video_upload_type: str, is_clip_video_cover: bool) -> int:
        if video_upload_type == "video_hook":
            return MAX_HOOK_VIDEO_UPLOAD_DURATION

        # reuse the logic for clip video cover for now, we can change this logic once we has a separate util function to collect all different config params
        if is_clip_video_cover:
            return MAX_COVER_VIDEO_UPLOAD_DURATION

        return MAX_SCENES_UPLOAD_DURATION

    def _get_video_dimensions(self, probe) -> Dimensions | None:
        """
        Gets the video dimensions from the output of an ffprobe
        """
        video_stream = next(
            (stream for stream in probe["streams"] if stream["codec_type"] == "video"), None
        )
        if video_stream:
            width = int(video_stream["width"])
            height = int(video_stream["height"])

            # sometimes width and height are swapped because the video is "rotated" - correct this before continuing
            rotation = next(
                (
                    side_data.get("rotation", 0)
                    for side_data in video_stream.get("side_data_list", [])
                    if side_data.get("side_data_type") == "Display Matrix"
                ),
                0,
            )
            if abs(rotation) in [90, 270]:
                width, height = height, width

            return Dimensions(width, height)

        return None

    def _process_video_transcode(
        self,
        upload_id: str,
        upload_key: str,
        temp_input_path: str,
        duration_bucket: int,
        maintain_audio: bool,
    ):
        """Transcode and upload video to S3"""
        print("\ninside process transcode\n")
        t0 = time.time()
        result = {"success": False, "s3_id": None}

        s3_id = f"video_upload_{upload_id}"
        s3_new_path = f"studio/uploads/{s3_id}.mp4"

        file_extension = temp_input_path.split(".")[-1].lower()
        if file_extension not in ["mp4", "mov", "avi", "mkv", "wmv", "webm"]:
            print(f"Unsupported file extension: {file_extension} Upload id {upload_id}\n")
            return result

        new_video_file_path = f"/tmp/{upload_id}.mp4"
        if file_extension == "mp4":
            # Remove audio and ensure reasonable bitrate (unless it's a hook video)
            try:
                if not maintain_audio:
                    # Remove audio completely
                    (
                        ffmpeg.input(temp_input_path)
                        .output(
                            new_video_file_path,
                            vcodec="libx264",
                            video_bitrate="2250k",  # Target bitrate of 2.25 Mbps
                            maxrate="2700k",  # Maximum bitrate (1.2× target)
                            bufsize="5400k",  # Buffer size (2× target)
                            preset="ultrafast",
                            an=None,  # Remove audio completely
                        )
                        .run(capture_stdout=True, capture_stderr=True)
                    )
                else:
                    # keep audio
                    (
                        ffmpeg.input(temp_input_path)
                        .output(
                            new_video_file_path,
                            vcodec="libx264",
                            video_bitrate="2250k",  # Target bitrate of 2.25 Mbps
                            maxrate="2700k",  # Maximum bitrate (1.2× target)
                            bufsize="5400k",  # Buffer size (2× target)
                            preset="ultrafast",
                            # audio settings
                            acodec="aac",
                            audio_bitrate="192k",
                            ar="44100",
                        )
                        .run(capture_stdout=True, capture_stderr=True)
                    )
            except ffmpeg.Error as e:
                print(
                    f"Error occurred during mp4 processing: {e.stderr.decode()} Upload id {upload_id}\n"
                )
                return result

        else:
            # Convert non mp4 file into mp4
            try:
                (
                    ffmpeg.input(temp_input_path)
                    .output(
                        new_video_file_path,
                        format="mp4",
                        vcodec="libx264",
                        acodec="aac",
                        audio_bitrate="192k",
                        preset="ultrafast",
                        ar="44100",
                    )
                    .run(capture_stdout=True, capture_stderr=True)
                )
                print(f"Conversion successful: {new_video_file_path}")
            except ffmpeg.Error as e:
                print(f"Error occurred during transcoding: {e.stderr.decode()} Upload id {upload_id}\n")
                return result
        try:
            retry_s3_upload(
                new_video_file_path,
                UPLOADS_S3_BUCKET,
                s3_new_path,
                ExtraArgs={
                    "ContentType": "video/mp4",
                },
            )
            # Log new s3 path for Modal logs.
            print(f"upload_id: {upload_id}, s3_new_path: {s3_new_path}")
            result["s3_id"] = s3_id
        except Exception as e:
            print(f"upload video: {upload_id} failed with error: {e}")
            return result

        result["success"] = True
        end_time = time.time()
        print(f"Done with transcode {end_time - t0:.3f}s Upload id {upload_id}\n")
        statsd.distribution(
            "video_processing.transcode", end_time - t0, tags=[f"duration_bucket:{duration_bucket}"]
        )
        return result

    def _process_cover_preview_from_video(
        self,
        upload_id: str,
        temp_input_path: str,
        image_dimensions: Dimensions | None,
        duration: int,
        format: PreviewFormat,
    ):
        """Generates and uploads a GIF preview for an uploaded video."""
        t0 = time.time()
        result = {"success": False, "gif_s3_id": "", "mp4_s3_id": ""}
        preview_s3_id = f"video_upload_{upload_id}_preview.{format}"
        preview_s3_without_ext = f"video_upload_{upload_id}_preview"
        output_path = f"/tmp/{preview_s3_id}"

        # Use the shorter of PREVIEW_DURATION or actual video duration
        preview_duration = min(PREVIEW_DURATIONS[format], duration)

        def _ffmpeg_gif():
            """Scale video and limit to 3fps"""
            (
                ffmpeg.input(temp_input_path, t=preview_duration)
                .filter("fps", fps=3)
                .filter(
                    "scale",
                    w=image_dimensions.width if image_dimensions else -1,
                    h=image_dimensions.height if image_dimensions else -1,
                )
                .output(output_path)
                .run(capture_stdout=True, capture_stderr=True)
            )

        def _ffmpeg_mp4():
            """Scale video, rip out audio"""
            (
                ffmpeg.input(temp_input_path, t=preview_duration)
                .filter(
                    "scale",
                    w=image_dimensions.width if image_dimensions else -1,
                    h=image_dimensions.height if image_dimensions else -1,
                )
                .filter("pad", "ceil(iw/2)*2", "ceil(ih/2)*2", 0, 0)
                .output(output_path, crf=23, an=None)
                .run(capture_stdout=True, capture_stderr=True)
            )

        try:
            # Generate preview using ffmpeg
            if format == "gif":
                _ffmpeg_gif()
            else:
                _ffmpeg_mp4()

            # Upload preview to S3
            s3_output_path = f"studio/uploads/{preview_s3_id}"
            content_types_by_format = {"gif": "image/gif", "mp4": "video/mp4"}
            retry_s3_upload(
                output_path,
                UPLOADS_S3_BUCKET,
                s3_output_path,
                ExtraArgs={"ContentType": f"{content_types_by_format[format]}"},
            )
            result["success"] = True
            result[f"{format}_s3_id"] = preview_s3_without_ext

        except Exception as e:
            print(
                f"Preview generation for video {upload_id} with format {format} failed with error: {e}, {e.stderr.decode() if hasattr(e, 'stderr') and e.stderr is not None else ''}"
            )

        print(
            f"Done with preview generation {time.time() - t0:.3f}s with format {format} Upload id {upload_id}\n"
        )
        return result

    def _process_cover_image_from_video(
        self,
        upload_id: str,
        temp_input_path: str,
        image_dimensions: Dimensions | None,
        t: int = 0,
    ):
        """Generates and uploads the cover image snapshot (thumbnail + large) for an uploaded video."""
        result = {"success": False, "image_s3_id": ""}
        t, image_s3_id, image_file, thumbnail_s3_id, thumbnail_file = self.process_time_slice(
            t, upload_id, temp_input_path, image_dimensions
        )
        if thumbnail_s3_id is None:
            return result

        # Upload thumbnail and large image to s3
        try:
            for s3_id, file in [(image_s3_id, image_file), (thumbnail_s3_id, thumbnail_file)]:
                s3_image_path = f"studio/uploads/{s3_id}.jpeg"
                retry_s3_upload(file, UPLOADS_S3_BUCKET, s3_image_path)
            result["image_s3_id"] = thumbnail_s3_id
        except Exception as e:
            print(f"upload video: {upload_id} failed with error: {e} Upload id {upload_id}\n")
            return result

        else:
            result["success"] = True
            result["image_s3_id"] = thumbnail_s3_id
        return result

    def get_thumbnail_dimensions(
        self, dimensions, short_side_length=360, long_side_length=-1
    ) -> Dimensions:
        width = dimensions.width
        height = dimensions.height

        return (
            Dimensions(short_side_length, long_side_length)
            if width <= height
            else Dimensions(long_side_length, short_side_length)
        )

    def process_time_slice(self, t, upload_id, temp_input_path, dimensions):
        """
        Returns an image s3 id and reference to local image file to upload for slice {t} of the video
        Also, for the first slice, returns an s3 id and reference to local image file for a thumbnail
        """
        # use base with prefix image_ for thumbnail and image_large for full size
        image_s3_base = f"video_upload_{upload_id}_snapshot_{t}s"

        image_s3_id = f"image_large_{image_s3_base}"
        image_file = f"/tmp/{image_s3_id}.jpeg"
        thumbnail_s3_id = None
        thumbnail_file = None

        try:
            ffmpeg.input(temp_input_path, ss=t).output(image_file, preset="ultrafast", vframes=1).run(
                capture_stdout=True, capture_stderr=True
            )

            # for the first slice, generate a thumbnail
            if t == 0 and dimensions:
                thumbnail_s3_id = f"image_{image_s3_base}"
                thumbnail_file = f"/tmp/{thumbnail_s3_id}.jpeg"
                # thumbnail_width, thumbnail_height = self.get_thumbnail_dimensions(dimensions)
                thumbnail_width, thumbnail_height = dimensions.width, dimensions.height
                ffmpeg.input(image_file).filter("scale", thumbnail_width, thumbnail_height).output(
                    thumbnail_file
                ).run(capture_stdout=True, capture_stderr=True)

        except ffmpeg.Error as e:
            print(
                f"ffmpeg error generating snapshot or thumbnail for {t}s: {e.stderr.decode()} Upload id {upload_id}\n"
            )
            return t, None, None, None, None
        except Exception as e:
            print(
                f"non-ffmpeg error generating snapshot or thumbnail for {t}s: {str(e)} Upload id {upload_id}\n"
            )
            return t, None, None, None, None

        return t, image_s3_id, image_file, thumbnail_s3_id, thumbnail_file

    def _process_image_to_text(
        self,
        upload_id: str,
        temp_input_path: str,
        duration: int,
        duration_bucket: int,
        image_dimensions: Dimensions | None,
    ):
        """Generate lyrics based on a list of images"""
        t0 = time.time()

        image_s3_ids = []
        image_files = []
        result = {"success": False, "lyrics": ""}

        start_time = time.time()
        # start processing slices
        try:
            with concurrent.futures.ThreadPoolExecutor() as executor:
                futures = [
                    executor.submit(
                        self.process_time_slice, t, upload_id, temp_input_path, image_dimensions
                    )
                    for t in range(0, duration, 6)
                ]
                results = {t: None for t in range(0, duration, 6)}
                thumbnail_result = None
                for future in concurrent.futures.as_completed(futures):
                    t, image_s3_id, image_file, thumbnail_s3_id, thumbnail_file = future.result()
                    if image_s3_id and image_file:
                        results[t] = (image_s3_id, image_file)
                    if thumbnail_s3_id and thumbnail_file:
                        thumbnail_result = (thumbnail_s3_id, thumbnail_file)
                for t in range(0, duration, 6):
                    if results[t]:
                        image_s3_id, image_file = results[t]
                        image_s3_ids.append(image_s3_id)
                        image_files.append(image_file)
                if thumbnail_result:
                    image_s3_ids.append(thumbnail_result[0])
                    image_files.append(thumbnail_result[1])
        except ffmpeg.Error as e:
            print(f"Error generating snapshots: {e.stderr.decode()} Upload id {upload_id}\n")
            return result

        end_image_extraction_time = time.time()

        print(f"Done with image extraction {time.time() - t0:.3f}s Upload id {upload_id}\n")
        try:
            for image_s3_id, image_file in zip(image_s3_ids, image_files):
                s3_image_path = f"studio/uploads/{image_s3_id}.jpeg"
                retry_s3_upload(image_file, UPLOADS_S3_BUCKET, s3_image_path)
            result["image_s3_id"] = image_s3_id
        except Exception as e:
            print(f"upload video: {upload_id} failed with error: {e} Upload id {upload_id}\n")
            return result

        end_snapshot_upload_time = time.time()

        print(f"Done with image upload {time.time() - t0:.3f}s Upload id {upload_id}\n")
        try:
            image_lyrics = self.modal_f_image_description.remote(
                QueueItem(
                    id=upload_id,
                    metadata={
                        "image_to_song_s3_ids": image_s3_ids,
                        "user_context": "",
                    },
                ).model_dump_json(),
                image_moderation_threshold=0,
            )
        # network errors in image extraction pipeline
        except Exception as e:
            print(f"Error in image extraction pipeline: {e} Upload id {upload_id}\n")
            return result

        end_image_lyrics_time = time.time()
        statsd.distribution(
            "video_processing.snapshot_creation",
            end_image_extraction_time - start_time,
            tags=[f"duration_bucket:{duration_bucket}"],
        )
        statsd.distribution(
            "video_processing.snapshot_upload",
            end_snapshot_upload_time - end_image_extraction_time,
            tags=[f"duration_bucket:{duration_bucket}"],
        )
        statsd.distribution(
            "video_processing.generate_image_descriptions",
            end_image_lyrics_time - end_snapshot_upload_time,
            tags=[f"duration_bucket:{duration_bucket}"],
        )
        statsd.distribution(
            "video_processing.image_to_text",
            end_image_lyrics_time - start_time,
            tags=[f"duration_bucket:{duration_bucket}"],
        )

        # system or user error
        if isinstance(image_lyrics, ImageExtractionError):
            print(f"Error in image extraction pipeline: {image_lyrics} Upload id {upload_id}\n")
            return None

        result["lyrics"] = image_lyrics
        result["success"] = True
        print(f"Done with image to text {time.time() - t0:.3f}s Upload id {upload_id}\n")
        return result

    def _detect_nsfw(
        self, bucket_name: str, upload_key: str, upload_id: str, duration_bucket: int
    ) -> bool:
        """Detect if the video is nsfw."""
        print(f"Starting NSFW check Upload id {upload_id}\n")
        video_moderation_threshold: float = 0.9
        s3_client = boto3.client("s3")
        presigned_url = s3_client.generate_presigned_url(
            "get_object",
            Params={"Bucket": bucket_name, "Key": upload_key},
            ExpiresIn=300,  # 5 minutes
        )

        headers = {"Authorization": f"Token {os.environ['HIVE_EAST_API_KEY']}"}
        data = {"url": presigned_url}

        start_time = time.time()
        model_response = requests.post(
            "https://api-va1.thehive.ai/api/v2/task/sync", headers=headers, data=data
        )

        # Record the latency
        end_time = time.time()
        total_time = end_time - start_time
        statsd.distribution(
            "video_processing.hive_moderation", total_time, tags=[f"duration_bucket:{duration_bucket}"]
        )

        # Parse the Hive API response
        response_dict = model_response.json()
        mod_classes = (
            response_dict.get("status", [{}])[0]
            .get("response", {})
            .get("output", [{}])[0]
            .get("classes", [])
        )
        print(
            f"Got NSFW check response from Hive, status is {response_dict.get('status', [{}])[0].get('status', None)} Upload id {upload_id}\n"
        )
        # Handle unexpected Hive API output
        if len(mod_classes) == 0:
            raise ValueError(f"No moderation classes found for video {upload_key}")
        # Check the Hive API `general_nsfw` class likelihood score
        general_nsfw_class = list(filter(lambda cls: cls["class"] == "general_nsfw", mod_classes))
        general_nsfw_score = general_nsfw_class[0].get("score", None) if general_nsfw_class else None
        # Handle unexpected Hive API output
        if general_nsfw_score is None:
            raise ValueError(f"No moderation score found for video {upload_key}")
        if general_nsfw_score > video_moderation_threshold:
            return True
        # If the likelihood is below the threshold, it means the video is safe for work
        else:
            return False


def _test_process_uploaded_video(model: VideoUploadStub):
    """
    Make sure you've got AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY set or these checks will fail
    """
    thumbnail_key = "studio/uploads/image_video_upload_testprocessuploadedvideo_snapshot_0s.jpeg"
    large_image_key = "studio/uploads/image_large_video_upload_testprocessuploadedvideo_snapshot_0s.jpeg"
    sprite_key = "studio/uploads/video_upload_sprite_testprocessuploadedvideo.jpeg"

    queue_item = (
        QueueItem(
            id="testprocessuploadedvideo",
            metadata={
                "upload_key": "test/d78819dd-4735-475c-9a28-8ae8f1cd98ca.mp4",
                "require_video_sprite": True,
            },
        )
    ).json()
    # make sure we clean up the s3 resources before running the test
    try:
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=thumbnail_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=large_image_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=sprite_key)
    except Exception as e:
        pass

    model.process_uploaded_video.remote(queue_item)

    # check for thumbnail and large image in the suno-data-uploads bucket
    try:
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=thumbnail_key)
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=large_image_key)
        # Check whether video sprite exists
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=sprite_key)
    except Exception as e:
        print(f"Headobject failed with error {e}")
        assert False

    print("Found snapshot and thumbnail on s3. Cleaning them up now.")

    # clean up resources before next test run
    try:
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=thumbnail_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=large_image_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=sprite_key)
    except Exception as e:
        print(
            f"Delete object failed with error {e} - make sure you clean up s3 resources or this test won't work :)"
        )

    return


def _test_process_uploaded_clip_video_cover(model: VideoUploadStub):
    """
    Make sure you've got AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY set or these checks will fail
    """
    queue_item = (
        QueueItem(
            id="testprocessuploadedclipvideocover",
            metadata={
                "upload_key": "test/cc86aba8-c118-4f30-875b-6d94d8ae15ef.mp4",
                "is_video_cover": True,
            },
        )
    ).json()

    model.process_uploaded_video.remote(queue_item)

    thumbnail_key = (
        "studio/uploads/image_video_upload_testprocessuploadedclipvideocover_snapshot_0s.jpeg"
    )
    large_image_key = (
        "studio/uploads/image_large_video_upload_testprocessuploadedclipvideocover_snapshot_0s.jpeg"
    )
    gif_key = "studio/uploads/video_upload_testprocessuploadedclipvideocover.gif"
    mp4_key = "studio/uploads/video_upload_testprocessuploadedclipvideocover_preview.mp4"

    # check for thumbnail and large image in the suno-data-uploads bucket
    try:
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=thumbnail_key)
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=large_image_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=gif_key)
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=mp4_key)
    except Exception as e:
        print(f"Headobject failed with error {e}")
        assert False

    print("Found snapshot and thumbnail on s3. Cleaning them up now.")

    # clean up resources before next test run
    try:
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=thumbnail_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=gif_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=large_image_key)

    except Exception as e:
        print(
            f"Delete object failed with error {e} - make sure you clean up s3 resources or this test won't work :)"
        )

    return


def _test_process_uploaded_clip_video_cover_gif_preview(model: VideoUploadStub):
    """
    Make sure you've got AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY set or these checks will fail
    """
    queue_item = (
        QueueItem(
            id="testprocessuploadedclipvideocover",
            metadata={
                "upload_key": "test/cc86aba8-c118-4f30-875b-6d94d8ae15ef.mp4",
                "is_video_cover": True,
                "preview_format": "gif",
            },
        )
    ).json()

    model.process_uploaded_video.remote(queue_item)

    thumbnail_key = (
        "studio/uploads/image_video_upload_testprocessuploadedclipvideocover_snapshot_0s.jpeg"
    )
    large_image_key = (
        "studio/uploads/image_large_video_upload_testprocessuploadedclipvideocover_snapshot_0s.jpeg"
    )
    gif_key = "studio/uploads/video_upload_testprocessuploadedclipvideocover_preview.gif"

    # check for thumbnail and large image in the suno-data-uploads bucket
    try:
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=thumbnail_key)
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=large_image_key)
        s3_client.head_object(Bucket=UPLOADS_S3_BUCKET, Key=gif_key)
    except Exception as e:
        print(f"Headobject failed with error {e}")
        assert False

    print("Found snapshot and thumbnail on s3. Cleaning them up now.")

    # clean up resources before next test run
    try:
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=thumbnail_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=large_image_key)
        s3_client.delete_object(Bucket=UPLOADS_S3_BUCKET, Key=gif_key)
    except Exception as e:
        print(
            f"Delete object failed with error {e} - make sure you clean up s3 resources or this test won't work :)"
        )

    return


@app.local_entrypoint()
def main():
    model = VideoUploadStub()
    _test_process_uploaded_video(model)
    _test_process_uploaded_clip_video_cover(model)
    _test_process_uploaded_clip_video_cover_gif_preview(model)