import streamlit as st
from moviepy.editor import VideoFileClip, VideoClip
from moviepy.video.fx.all import crop
import tempfile
import os
import numpy as np
from PIL import Image, ImageDraw
import hashlib

# Fix for PIL.Image.ANTIALIAS deprecation
if not hasattr(Image, 'ANTIALIAS'):
    Image.ANTIALIAS = Image.LANCZOS

# Function to create a circular mask
def create_circular_mask(size):
    # Create a square image with a black background
    img = Image.new('L', (size, size), 0)
    # Create a white circle in the middle
    draw = ImageDraw.Draw(img)
    draw.ellipse((0, 0, size, size), fill=255)
    # Convert the image to a numpy array
    mask = np.array(img) / 255.0  # Normalize to [0, 1] range
    return mask

def compute_file_hash(file):
    hash_md5 = hashlib.md5()
    for chunk in file:
        hash_md5.update(chunk)
    file.seek(0)  # Reset file pointer after reading
    return hash_md5.hexdigest()

# Function to process the video with audio and progress logging
@st.cache_data(hash_funcs={str: lambda _: None})  # Ignore `input_path` by setting its hash to None
def process_video(input_path, file_hash):
    # Load the video file
    input_video = VideoFileClip(input_path)
    w, h = input_video.size
    circle_size = 360
    aspect_ratio = w / h

    # Resize the video while maintaining aspect ratio
    if w > h:
        new_w = int(circle_size * aspect_ratio)
        new_h = circle_size
    else:
        new_w = circle_size
        new_h = int(circle_size / aspect_ratio)

    # Resize and crop the video to a square
    resized_video = input_video.resize((new_w, new_h))
    square_video = crop(
        resized_video,
        x_center=new_w / 2,
        y_center=new_h / 2,
        width=circle_size,
        height=circle_size
    )

    # Create a circular mask
    mask_array = create_circular_mask(circle_size)

    # Apply the circular mask to each frame
    def apply_mask(get_frame, t):
        frame = get_frame(t)
        alpha = mask_array[..., np.newaxis]  # Add axis for channels
        frame = frame * alpha + 255 * (1 - alpha)  # Assuming white background
        return frame.astype('uint8')

    circular_video = square_video.fl(apply_mask)

    # Add audio back to the video
    circular_video = circular_video.set_audio(input_video.audio)

    # Save output to a temporary file with progress logging
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    circular_video.write_videofile(
        temp_file.name,
        codec="libx264",
        audio_codec="aac",
        temp_audiofile=tempfile.NamedTemporaryFile(suffix=".m4a").name,
        remove_temp=True,
    )
    return temp_file.name

# Streamlit UI
st.title("Circular Video Cropper with Progress")
st.write("Upload a video, and the app will crop it into a circular format with visible progress.")

# Upload file
uploaded_file = st.file_uploader("Choose a video file", type=["mp4", "mov", "avi", "mkv"])

if uploaded_file:

    file_hash = compute_file_hash(uploaded_file)
    # Save uploaded file to a temporary location
    temp_input_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    temp_input_file.write(uploaded_file.read())
    temp_input_file.close()

    input_video = VideoFileClip(temp_input_file.name)
    total_duration = input_video.duration

    output_path = process_video(temp_input_file.name, file_hash)

    # Display success and download button
    st.success("Video processed successfully!")
    with open(output_path, "rb") as file:
        st.download_button(
            label="Download Circular Video",
            data=file,
            file_name="circular_video.mp4",
            mime="video/mp4"
        )

    os.remove(temp_input_file.name)