from typing import Callable
import ffmpeg
from collections import namedtuple
from dataclasses import dataclass

Dimensions = namedtuple("Dimensions", ["width", "height"])


UPLOADS_S3_BUCKET = "suno-data-uploads"


def get_video_sprite_dimensions(video_dimensions: Dimensions | None):
    if video_dimensions is None:
        return None
    # make the demension propotional and fit into 240p
    max_height = 240
    aspect_ratio = video_dimensions.width / video_dimensions.height

    if video_dimensions.height <= max_height:
        # If height is already less than or equal to 240p, keep original dimensions
        return {
            "thumbnail_width": video_dimensions.width,
            "thumbnail_height": video_dimensions.height,
        }

    # Scale down proportionally to fit height
    new_height = max_height
    new_width = int(new_height * aspect_ratio)

    return {
        "thumbnail_width": new_width,
        "thumbnail_height": new_height,
    }


@dataclass
class VideoSpriteSchema:
    s3_id: str
    thumbnail_width: int
    thumbnail_height: int
    image_width: int
    image_height: int
    interval: float
    total_frames: int

    def to_dict(self):
        return {
            "s3_id": self.s3_id,
            "thumbnail_width": self.thumbnail_width,
            "thumbnail_height": self.thumbnail_height,
            "image_width": self.image_width,
            "image_height": self.image_height,
            "interval": self.interval,
            "total_frames": self.total_frames,
        }


@dataclass
class VideoSpriteGenerationSchema:
    output_path: str | None = None
    sprite_schema: VideoSpriteSchema | None = None
    success: bool = False
    error_message: str | None = None


def process_video_sprite(
    upload_id: str,
    temp_input_path: str,
    video_duration: float,
    duration_bucket: int,
    video_dimensions: Dimensions,
    retry_s3_upload: Callable[[str, str, str, dict], None],
) -> VideoSpriteGenerationSchema:
    """Processes a video sprite for an uploaded video.
    The return is in /tmp/video_upload_sprite_<upload_id>.jpeg"""

    video_sprite_dimensions = get_video_sprite_dimensions(video_dimensions)
    if video_sprite_dimensions is None:
        return VideoSpriteGenerationSchema(
            output_path=None,
            sprite_schema=None,
            success=False,
            error_message="video_sprite_dimensions is None",
        )
    sprite_s3_id = f"video_upload_sprite_{upload_id}"
    output_path = f"/tmp/{sprite_s3_id}.jpeg"

    """Generate sprite from video"""
    """generate the sprite image from the video, using ffmpeg, the image is tiled up left to right, top to bottom, with the dimensions of the video_sprite_dimensions
        The interval of the sprite is 1 second per image.
    """
    try:
        # Calculate number of frames needed (1 frame per second)
        num_frames = int(video_duration)

        # Calculate grid dimensions
        grid_cols = min(8, num_frames)  # Max 8 columns
        grid_rows = (num_frames + grid_cols - 1) // grid_cols  # Ceiling division

        # Get thumbnail dimensions
        thumb_width = video_sprite_dimensions["thumbnail_width"]
        thumb_height = video_sprite_dimensions["thumbnail_height"]

        # Calculate total sprite dimensions
        sprite_width = thumb_width * grid_cols
        sprite_height = thumb_height * grid_rows

        # Build ffmpeg command to generate sprite
        stream = (
            ffmpeg.input(temp_input_path, t=duration_bucket)
            .filter("fps", fps=1)  # 1 frame per second
            .filter("scale", thumb_width, thumb_height)  # Scale to thumbnail size
            .filter("tile", f"{grid_cols}x{grid_rows}")  # Tile frames in grid with proper format
            .output(output_path, vframes=1, preset="ultrafast")  # Output single sprite image
        )

        # Run ffmpeg command
        stream.run(capture_stdout=True, capture_stderr=True)

        # Upload sprite to s3
        s3_output_path = f"studio/uploads/{sprite_s3_id}.jpeg"
        retry_s3_upload(output_path, UPLOADS_S3_BUCKET, s3_output_path, {"ContentType": "image/jpeg"})

        print(f"Upload sprite with upload_id: {upload_id}, s3_new_path: {s3_output_path}")
        sprite_schema = VideoSpriteSchema(
            s3_id=sprite_s3_id,
            thumbnail_width=thumb_width,
            thumbnail_height=thumb_height,
            image_width=sprite_width,
            image_height=sprite_height,
            interval=1,
            total_frames=num_frames,
        )
        return VideoSpriteGenerationSchema(
            output_path=output_path, sprite_schema=sprite_schema, success=True
        )
    except ffmpeg.Error as e:
        print(f"ffmpeg error generating sprite: {e.stderr.decode()}")
        return VideoSpriteGenerationSchema(
            output_path=None, sprite_schema=None, success=False, error_message=e.stderr.decode()
        )