import gradio as gr
import os
import cv2
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
import shutil

# Thư mục tạm thời để lưu trữ file tải lên
UPLOAD_FOLDER = "uploads"
OUTPUT_FOLDER = "outputs"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# Khởi tạo mô hình Real-ESRGAN
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
upsampler = RealESRGANer(
    scale=4,
    model_path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
    model=model,
    tile=0,
    tile_pad=10,
    pre_pad=0,
    half=True
)

def enhance_video(input_video):
    """
    Hàm xử lý video đầu vào và trả về video đã tăng cường.
    """
    # Lưu file tải lên vào thư mục uploads
    input_path = os.path.join(UPLOAD_FOLDER, "input.mp4")
    output_path = os.path.join(OUTPUT_FOLDER, "enhanced_output.mp4")
    input_video.save(input_path)

    # Tách video thành khung hình
    frames_folder = os.path.join(UPLOAD_FOLDER, "frames")
    os.makedirs(frames_folder, exist_ok=True)
    extract_frames(input_path, frames_folder)

    # Xử lý từng khung hình với Real-ESRGAN
    enhanced_frames_folder = os.path.join(OUTPUT_FOLDER, "enhanced_frames")
    os.makedirs(enhanced_frames_folder, exist_ok=True)
    enhance_frames(frames_folder, enhanced_frames_folder)

    # Ghép các khung hình đã tăng cường thành video
    combine_frames_to_video(enhanced_frames_folder, output_path)

    # Xóa các thư mục tạm thời
    shutil.rmtree(frames_folder)
    shutil.rmtree(enhanced_frames_folder)

    return output_path

def extract_frames(video_path, output_folder):
    """
    Tách video thành các khung hình.
    """
    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_path = os.path.join(output_folder, f"frame_{frame_count:04d}.png")
        cv2.imwrite(frame_path, frame)
        frame_count += 1
    cap.release()

def enhance_frames(input_folder, output_folder):
    """
    Tăng cường từng khung hình bằng Real-ESRGAN.
    """
    for frame_name in os.listdir(input_folder):
        frame_path = os.path.join(input_folder, frame_name)
        img = cv2.imread(frame_path, cv2.IMREAD_UNCHANGED)
        enhanced_img, _ = upsampler.enhance(img, outscale=4)
        output_path = os.path.join(output_folder, frame_name)
        cv2.imwrite(output_path, enhanced_img)

def combine_frames_to_video(frames_folder, output_path):
    """
    Ghép các khung hình thành video.
    """
    frame_rate = 30  # Số khung hình/giây
    first_frame = os.path.join(frames_folder, os.listdir(frames_folder)[0])
    height, width, _ = cv2.imread(first_frame).shape

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    video_writer = cv2.VideoWriter(output_path, fourcc, frame_rate, (width, height))

    for frame_name in sorted(os.listdir(frames_folder)):
        frame_path = os.path.join(frames_folder, frame_name)
        frame = cv2.imread(frame_path)
        video_writer.write(frame)

    video_writer.release()

# Tạo giao diện Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Video Enhancement with Real-ESRGAN")
    gr.Markdown("Upload a video and get an enhanced version with improved quality.")
    
    with gr.Row():
        input_video = gr.File(label="Upload Video")
        output_video = gr.Video(label="Enhanced Video")
    
    submit_button = gr.Button("Enhance Video")
    submit_button.click(enhance_video, inputs=input_video, outputs=output_video)

# Khởi chạy ứng dụng
demo.launch()