import gradio as gr
from PIL import Image
import os
import logging
from moviepy.video.io.VideoFileClip import VideoFileClip
import numpy as np
import tempfile

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.FileHandler("image_converter.log"), logging.StreamHandler()],
)
logger = logging.getLogger(__name__)


def cleanup_temp_files():
    """Clean up temporary files"""
    temp_files = ["input_image.png", "converted_image.webp"]
    for file in temp_files:
        if os.path.exists(file):
            os.remove(file)


def get_image_stats(image_path, pil_image):
    """Get image file size (in KB) and dimensions"""
    file_size_mb = os.path.getsize(image_path) / (1024)
    width, height = pil_image.size
    return f"Size: {file_size_mb:.0f}KB, Dimensions: {width}x{height}"


def convert_png_to_webp(image, quality=80):
    if image is None:
        logger.error("No image provided")
        return None, None, "", ""

    try:
        cleanup_temp_files()  # Clean up before conversion

        # Save input image temporarily to get file size
        input_path = "input_image.png"
        image.save(input_path, "PNG")
        input_stats = get_image_stats(input_path, image)

        logger.info(
            f"Starting conversion - Image size: {image.size}, Mode: {image.mode}"
        )

        # Convert image to WebP format
        output_path = "converted_image.webp"
        image = image.convert("RGBA")  # Preserve transparency
        image.save(output_path, "WEBP", quality=quality)

        # Get output stats
        output_image = Image.open(output_path)
        output_stats = get_image_stats(output_path, output_image)

        logger.info(f"Successfully converted image to WebP with quality {quality}")
        return output_path, Image.open(output_path), input_stats, output_stats

    except Exception as e:
        logger.error(f"Error converting image: {str(e)}")
        cleanup_temp_files()
        return None, None, "", ""


def convert_video_to_webp(video_path, fps=10, quality=80):
    if video_path is None:
        logger.error("No video provided")
        return None, None, ""

    try:
        # Load video and extract frames
        video = VideoFileClip(video_path)
        duration = video.duration

        logger.info(f"Starting conversion - Duration: {duration:.1f}s, FPS: {fps}")
        # Create temporary directory for frames
        with tempfile.TemporaryDirectory() as temp_dir:
            frames = []
            for frame in video.iter_frames(fps=fps):
                # Convert frame to PIL Image
                pil_frame = Image.fromarray(frame)
                frames.append(pil_frame)

            # Save as animated WebP
            output_path = "converted_animation.webp"
            frames[0].save(
                output_path,
                save_all=True,
                append_images=frames[1:],
                duration=int(1000 / fps),
                quality=quality,
                method=6,
            )

            # Get stats
            stats = f"Duration: {duration:.1f}s, Frames: {len(frames)}, FPS: {fps}"

        video.close()
        return output_path, Image.open(output_path), stats

    except Exception as e:
        logger.error(f"Error converting video: {str(e)}")
        return None, None, ""


# Replace the Gradio Interface with Blocks
with gr.Blocks(title="Media Converter", theme=gr.themes.Soft()) as demo:
    with gr.Tabs():
        with gr.Tab("PNG to WebP"):
            gr.Markdown(
                "## PNG to WebP Converter\nUpload a PNG image and convert it to WebP format with adjustable quality."
            )

            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(type="pil", label="Upload PNG Image")
                    quality = gr.Slider(10, 100, value=80, step=5, label="Quality")

                with gr.Column(scale=1):
                    output_file = gr.File(label="Download WebP Image")
                    output_preview = gr.Image(label="Preview")
                    input_stats = gr.Textbox(label="Input Image Details")
                    output_stats = gr.Textbox(label="Output Image Details")

            # Set up event listeners for automatic conversion
            input_image.change(
                fn=convert_png_to_webp,
                inputs=[input_image, quality],
                outputs=[output_file, output_preview, input_stats, output_stats],
            )
            quality.change(
                fn=convert_png_to_webp,
                inputs=[input_image, quality],
                outputs=[output_file, output_preview, input_stats, output_stats],
            )
        with gr.Tab("Video to WebP"):
            with gr.Row():
                with gr.Column(scale=1):
                    input_video = gr.Video(label="Upload MOV Video")
                    with gr.Row():
                        fps = gr.Slider(1, 30, value=10, step=1, label="FPS")
                        quality = gr.Slider(10, 100, value=80, step=5, label="Quality")

                with gr.Column(scale=1):
                    output_file = gr.File(label="Download WebP Animation")
                    output_preview = gr.Image(label="Preview")
                    video_stats = gr.Textbox(label="Video Details")

            # Set up event listeners for video conversion
            input_video.change(
                fn=convert_video_to_webp,
                inputs=[input_video, fps, quality],
                outputs=[output_file, output_preview, video_stats],
            )
        with gr.Tab("README"):
            with gr.Column(scale=1):
                gr.Markdown(
                    """### ⓘ Note:
                            The converter works both locally and in remote.<br> 
                            For more details (source codes, clone repo, etc.), visit [README](https://huggingface.co/spaces/prasanthntu/media-converter/blob/main/README.md)""",
                    label="Note",
                    show_label=True,
                    container=True,
                )


if __name__ == "__main__":
    demo.launch()