import gradio as gr
import cv2
import numpy as np
from PIL import Image
import tempfile
import os

def enhance_video(input_video, target_fps, upscale_factor, denoise_strength, sharpen_strength):
    if input_video is None:
        return None, "Silakan upload video terlebih dahulu"
    
    try:
        cap = cv2.VideoCapture(input_video)
        original_fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        new_width = int(width * upscale_factor)
        new_height = int(height * upscale_factor)
        
        temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        output_path = temp_output.name
        temp_output.close()
        
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, target_fps, (new_width, new_height))
        
        frame_interpolation_ratio = target_fps / original_fps
        
        frames_buffer = []
        frame_idx = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            enhanced_frame = frame.copy()
            
            if denoise_strength > 0:
                enhanced_frame = cv2.fastNlMeansDenoisingColored(
                    enhanced_frame, None, 
                    h=denoise_strength, 
                    hColor=denoise_strength, 
                    templateWindowSize=7, 
                    searchWindowSize=21
                )
            
            if sharpen_strength > 0:
                kernel = np.array([[-1,-1,-1],
                                   [-1, 9 + sharpen_strength,-1],
                                   [-1,-1,-1]])
                enhanced_frame = cv2.filter2D(enhanced_frame, -1, kernel)
            
            enhanced_frame = cv2.resize(enhanced_frame, (new_width, new_height), 
                                       interpolation=cv2.INTER_LANCZOS4)
            
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
            lab = cv2.cvtColor(enhanced_frame, cv2.COLOR_BGR2LAB)
            lab[:,:,0] = clahe.apply(lab[:,:,0])
            enhanced_frame = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
            
            frames_buffer.append(enhanced_frame)
            
            if len(frames_buffer) >= 2 and frame_interpolation_ratio > 1:
                num_interpolated = int(frame_interpolation_ratio) - 1
                for i in range(num_interpolated):
                    alpha = (i + 1) / (num_interpolated + 1)
                    interpolated = cv2.addWeighted(frames_buffer[-2], 1-alpha, 
                                                   frames_buffer[-1], alpha, 0)
                    out.write(interpolated)
            
            out.write(enhanced_frame)
            
            if len(frames_buffer) > 2:
                frames_buffer.pop(0)
            
            frame_idx += 1
        
        cap.release()
        out.release()
        
        status_msg = f"✅ Video berhasil diproses!\n"
        status_msg += f"📊 Resolusi: {width}x{height} → {new_width}x{new_height}\n"
        status_msg += f"🎬 FPS: {original_fps:.1f} → {target_fps}\n"
        status_msg += f"📦 Total frames: {frame_count} → {frame_idx * int(frame_interpolation_ratio)}"
        
        return output_path, status_msg
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def create_demo():
    with gr.Blocks(title="Ultra Video Enhancer 4D Pro") as demo:
        gr.Markdown("""
        # 🎬 Ultra Video Enhancer 4D Pro
        ### Tingkatkan kualitas video Anda dengan AI Enhancement
        
        **Fitur:**
        - 🎯 Upscaling hingga 4x resolusi
        - 🚀 Frame interpolation hingga 240 FPS
        - ✨ AI Denoising & Sharpening
        - 🎨 Adaptive contrast enhancement
        - 💎 Super resolution dengan Lanczos4
        """)
        
        with gr.Row():
            with gr.Column():
                input_video = gr.Video(label="📥 Upload Video")
                
                with gr.Accordion("⚙️ Pengaturan Enhancement", open=True):
                    target_fps = gr.Slider(
                        minimum=24, maximum=240, value=60, step=1,
                        label="🎬 Target FPS",
                        info="Frame rate output (24-240 FPS)"
                    )
                    
                    upscale_factor = gr.Slider(
                        minimum=1.0, maximum=4.0, value=2.0, step=0.5,
                        label="📐 Upscale Factor",
                        info="Tingkatkan resolusi (1x-4x)"
                    )
                    
                    denoise_strength = gr.Slider(
                        minimum=0, maximum=20, value=5, step=1,
                        label="🧹 Denoise Strength",
                        info="Kurangi noise (0=off, 20=max)"
                    )
                    
                    sharpen_strength = gr.Slider(
                        minimum=0, maximum=5, value=2, step=1,
                        label="✨ Sharpen Strength",
                        info="Pertajam detail (0=off, 5=max)"
                    )
                
                enhance_btn = gr.Button("🚀 Enhance Video", variant="primary", size="lg")
            
            with gr.Column():
                output_video = gr.Video(label="📤 Enhanced Video")
                status_text = gr.Textbox(label="📊 Status", lines=5)
        
        gr.Markdown("""
        ### 💡 Tips:
        - Video pendek (<30 detik) akan diproses lebih cepat
        - Upscale 2x dengan 60 FPS memberikan hasil optimal
        - Denoise berguna untuk video dengan noise/grain
        - Sharpen meningkatkan ketajaman edge dan detail
        
        ### ⚡ Preset Rekomendasi:
        - **Balanced**: FPS=60, Upscale=2x, Denoise=5, Sharpen=2
        - **Ultra Quality**: FPS=120, Upscale=3x, Denoise=8, Sharpen=3
        - **Speed**: FPS=30, Upscale=1.5x, Denoise=3, Sharpen=1
        """)
        
        enhance_btn.click(
            fn=enhance_video,
            inputs=[input_video, target_fps, upscale_factor, denoise_strength, sharpen_strength],
            outputs=[output_video, status_text]
        )
    
    return demo

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