import gradio as gr
from moviepy.editor import VideoFileClip, CompositeVideoClip
import cv2
import numpy as np
import tempfile
import shutil
import moviepy.video.fx.all as vfx
import os
from moviepy.editor import ImageSequenceClip

def remove_green_background(overlay_video_bytes):
    overlay_temp = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    overlay_temp.write(overlay_video_bytes)
    overlay_temp.close()
    overlay_video_path = overlay_temp.name
    
    cap = cv2.VideoCapture(overlay_video_path)
    frame_width = int(cap.get(3))
    frame_height = int(cap.get(4))
    fps = cap.get(cv2.CAP_PROP_FPS)

    no_bg_temp = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    out = cv2.VideoWriter(no_bg_temp.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height), True)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        lower_green = np.array([36, 25, 25])
        upper_green = np.array([86, 255,255])
        
        mask = cv2.inRange(hsv, lower_green, upper_green)
        mask_inv = cv2.bitwise_not(mask)
        
        bg_removed = cv2.bitwise_and(frame, frame, mask=mask_inv)
        
        out.write(bg_removed)

    cap.release()
    out.release()
    
    return no_bg_temp.name

position_map = {
    "Top Left": ('left', 'top'),
    "Top Center": ('center', 'top'),
    "Top Right": ('right', 'top'),
    "Middle Left": ('left', 'center'),
    "Middle Center": ('center', 'center'),
    "Middle Right": ('right', 'center'),
    "Bottom Left": ('left', 'bottom'),
    "Bottom Center": ('center', 'bottom'),
    "Bottom Right": ('right', 'bottom'),
}

def remove_green_background_and_save_images(overlay_video_path):
    cap = cv2.VideoCapture(overlay_video_path)
    frame_idx = 0
    temp_dir = tempfile.mkdtemp()

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        lower_green = np.array([36, 25, 25])
        upper_green = np.array([86, 255,255])
        
        mask = cv2.inRange(hsv, lower_green, upper_green)
        mask_inv = cv2.bitwise_not(mask)
        
        bgra = cv2.cvtColor(frame, cv2.COLOR_BGR2BGRA)
        bgra[:, :, 3] = mask_inv
        
        cv2.imwrite(os.path.join(temp_dir, f"frame_{frame_idx:05d}.png"), bgra)
        frame_idx += 1

    cap.release()
    return temp_dir


def combine_videos(background_video_bytes, overlay_video_bytes, position, scale):
    background_temp = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    background_temp.write(background_video_bytes)
    background_temp.close()
    background_video_path = background_temp.name

    overlay_temp = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    overlay_temp.write(overlay_video_bytes)
    overlay_temp.close()
    overlay_video_path = overlay_temp.name

    temp_image_sequence_dir = remove_green_background_and_save_images(overlay_video_path)

    overlay_clip = ImageSequenceClip(temp_image_sequence_dir, fps=24).resize(scale / 10)

    background_clip = VideoFileClip(background_video_path)
    
    pos = position_map[position]
    final_clip = CompositeVideoClip([background_clip, overlay_clip.set_position(pos)], size=background_clip.size)
    
    output_video_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
    final_clip.write_videofile(output_video_path, codec='libx264', audio_codec='aac')

    shutil.rmtree(temp_image_sequence_dir)
    os.unlink(background_temp.name)
    os.unlink(overlay_temp.name)

    return output_video_path

css = """
footer { 
    visibility: hidden; 
}
"""

iface = gr.Interface(css=css,
    fn=combine_videos,
    inputs=[
        gr.File(type="binary", label="Background Video"),
        gr.File(type="binary", label="Overlay Video"),
        gr.Dropdown(["Top Left", "Top Center", "Top Right", "Middle Left", "Middle Center", "Middle Right", "Bottom Left", "Bottom Center", "Bottom Right"], label="Position"),
        gr.Slider(minimum=1, maximum=100, value=5, label="Scale (x0.1)"),
    ],
    outputs=gr.Video(label="Combined Video")
)

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