# app.py
# AI Video Enhancer 4K - Gradio app for Hugging Face Spaces
# With batched GPU processing to avoid ZeroGPU timeout

import os
import shutil
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Tuple, List

import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image
import cv2
from huggingface_hub import hf_hub_download

# Config
TEMP_DIR = Path(tempfile.gettempdir()) / "hf_video_enhancer"
TEMP_DIR.mkdir(parents=True, exist_ok=True)

# Cache model path globally so we don't re-download
_model_cache = {}


def run_cmd(cmd):
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if p.returncode != 0:
        raise RuntimeError(f"Command failed: {p.stderr.decode()}")
    return p.stdout.decode()


def probe_video(video_path: str) -> Tuple[float, int, int, float]:
    cmd = [
        "ffprobe", "-v", "error",
        "-select_streams", "v:0",
        "-show_entries", "stream=width,height,duration,r_frame_rate",
        "-of", "default=noprint_wrappers=1:nokey=0",
        video_path
    ]
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out = p.stdout.decode()
    width = height = 0
    duration = 0.0
    fps = 30.0
    
    for line in out.splitlines():
        if line.startswith("width="):
            width = int(line.split("=")[1])
        elif line.startswith("height="):
            height = int(line.split("=")[1])
        elif line.startswith("duration="):
            try:
                duration = float(line.split("=")[1])
            except:
                pass
        elif line.startswith("r_frame_rate="):
            try:
                fps_str = line.split("=")[1]
                if "/" in fps_str:
                    num, den = fps_str.split("/")
                    fps = float(num) / float(den)
                else:
                    fps = float(fps_str)
            except:
                pass
    
    return duration, width, height, fps


def extract_frames(video_path: str, frames_dir: Path):
    frames_dir.mkdir(parents=True, exist_ok=True)
    run_cmd([
        "ffmpeg", "-y", "-i", video_path,
        "-vsync", "0",
        str(frames_dir / "%06d.png")
    ])


def reassemble_video(frames_dir: Path, audio_src: str, out_path: str, fps: float = 30.0):
    tmp_video = str(frames_dir.parent / "tmp_video.mp4")
    run_cmd([
        "ffmpeg", "-y", "-framerate", str(fps),
        "-i", str(frames_dir / "%06d.png"),
        "-c:v", "libx264", "-preset", "veryfast", "-pix_fmt", "yuv420p",
        "-crf", "18", tmp_video
    ])

    p = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "a", "-show_entries", 
         "stream=codec_type", "-of", "default=noprint_wrappers=1", audio_src],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    
    if p.stdout.decode().strip():
        run_cmd([
            "ffmpeg", "-y", "-i", tmp_video, "-i", audio_src,
            "-c:v", "copy", "-c:a", "aac",
            "-map", "0:v:0", "-map", "1:a:0", out_path
        ])
        os.remove(tmp_video)
    else:
        shutil.move(tmp_video, out_path)


def simple_upscale(img: np.ndarray, scale: int) -> np.ndarray:
    """Simple bicubic upscaling using OpenCV"""
    h, w = img.shape[:2]
    return cv2.resize(img, (w * scale, h * scale), interpolation=cv2.INTER_CUBIC)


def get_model_path(scale: int) -> str:
    """Download model weights (cached)"""
    global _model_cache
    if scale not in _model_cache:
        if scale == 2:
            _model_cache[scale] = hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x2.pth")
        else:
            _model_cache[scale] = hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth")
    return _model_cache[scale]


@spaces.GPU(duration=120)
def enhance_batch(frame_paths: List[str], scale: int = 4) -> int:
    """
    Enhance a SMALL BATCH of frames using Real-ESRGAN.
    Called multiple times for different batches to avoid timeout.
    """
    from spandrel import ImageModelDescriptor, ModelLoader
    
    if not frame_paths:
        return 0
    
    model_path = get_model_path(scale)
    model = ModelLoader().load_from_file(model_path)
    assert isinstance(model, ImageModelDescriptor)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device).eval()
    
    processed = 0
    for frame_path in frame_paths:
        img = cv2.imread(frame_path)
        if img is None:
            continue
            
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).float().div(255.0)
        tensor = tensor.unsqueeze(0).to(device)
        
        with torch.no_grad():
            output = model(tensor)
        
        output = output.squeeze(0).cpu().clamp(0, 1).mul(255).byte()
        output = output.permute(1, 2, 0).numpy()
        output_bgr = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
        
        cv2.imwrite(frame_path, output_bgr)
        processed += 1
    
    return processed


def process_video(video_file, scale: int = 4, progress=gr.Progress()) -> Tuple[str, str]:
    """Main video processing - handles file I/O outside GPU function"""
    if video_file is None:
        return "⚠️ Please upload a video file.", None
    
    ts = int(time.time() * 1000)
    base_dir = TEMP_DIR / f"job_{ts}"
    base_dir.mkdir(parents=True, exist_ok=True)
    in_path = base_dir / "input_video"

    try:
        shutil.copy(video_file, in_path)
    except Exception as e:
        return f"Error: {e}", None

    try:
        duration, w, h, fps = probe_video(str(in_path))
    except Exception as e:
        shutil.rmtree(base_dir, ignore_errors=True)
        return f"Error probing video: {e}", None

    if duration <= 0:
        shutil.rmtree(base_dir, ignore_errors=True)
        return "Could not determine video duration.", None

    # Limit for ZeroGPU - process max ~15 seconds of video
    max_frames = int(fps * 15)
    
    progress(0.05, f"Video: {w}x{h}, {duration:.1f}s")

    frames_dir = base_dir / "frames"
    try:
        extract_frames(str(in_path), frames_dir)
    except Exception as e:
        shutil.rmtree(base_dir, ignore_errors=True)
        return f"Failed extracting frames: {e}", None

    frame_files = sorted(frames_dir.glob("*.png"))
    num_frames = len(frame_files)
    
    # Limit frames if too many
    if num_frames > max_frames:
        progress(0.1, f"Limiting to {max_frames} frames...")
        for f in frame_files[max_frames:]:
            f.unlink()
        frame_files = frame_files[:max_frames]
        num_frames = max_frames

    progress(0.15, f"Enhancing {num_frames} frames...")

    # Pre-download model (outside GPU call)
    try:
        get_model_path(scale)
    except Exception as e:
        print(f"Model download failed: {e}")

    # Process in SMALL BATCHES (10 frames per GPU call to avoid timeout)
    batch_size = 10
    total_enhanced = 0
    use_fallback = False
    
    for batch_start in range(0, num_frames, batch_size):
        batch_end = min(batch_start + batch_size, num_frames)
        batch_paths = [str(f) for f in frame_files[batch_start:batch_end]]
        
        if not use_fallback:
            try:
                enhanced = enhance_batch(batch_paths, scale)
                total_enhanced += enhanced
                print(f"Batch {batch_start}-{batch_end}: enhanced {enhanced} frames")
            except Exception as e:
                print(f"GPU batch failed: {e}, switching to fallback...")
                use_fallback = True
        
        if use_fallback:
            # Fallback to CPU upscaling for this batch
            for fp in batch_paths:
                img = cv2.imread(fp)
                if img is not None:
                    upscaled = simple_upscale(img, scale)
                    cv2.imwrite(fp, upscaled)
                    total_enhanced += 1
        
        # Update progress
        pct = 0.15 + 0.75 * (batch_end / num_frames)
        progress(pct, f"Processed {batch_end}/{num_frames} frames")

    progress(0.92, "Creating video...")
    
    out_video = base_dir / "enhanced_output.mp4"
    try:
        reassemble_video(frames_dir, str(in_path), str(out_video), fps)
    except Exception as e:
        shutil.rmtree(base_dir, ignore_errors=True)
        return f"Failed reassembling: {e}", None

    shutil.rmtree(frames_dir, ignore_errors=True)

    try:
        _, out_w, out_h, _ = probe_video(str(out_video))
        method = "bicubic" if use_fallback else "Real-ESRGAN"
        progress(1.0, "Done!")
        return f"✅ Done! {w}x{h} → {out_w}x{out_h} ({method})", str(out_video)
    except:
        return "✅ Done!", str(out_video)


# Gradio UI
with gr.Blocks(title="AI Video Enhancer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎬 AI Video Enhancer")
    gr.Markdown("Upscale videos using Real-ESRGAN AI enhancement.")
    
    # LOGIN BUTTON - This allows ZeroGPU to recognize your Pro account
    gr.LoginButton()
    
    with gr.Row():
        with gr.Column(scale=2):
            video_in = gr.File(label="Upload video", file_types=[".mp4", ".avi", ".mov", ".mkv", ".webm"])
            scale_choice = gr.Radio(choices=[2, 4], value=4, label="Upscale Factor")
            btn = gr.Button("🚀 Enhance", variant="primary")
            status = gr.Textbox(label="Status", interactive=False)
        with gr.Column(scale=1):
            out_video = gr.Video(label="Result")

    gr.Markdown("**Note:** Limited to ~15 seconds for ZeroGPU. Longer videos will be truncated.")

    btn.click(fn=process_video, inputs=[video_in, scale_choice], outputs=[status, out_video])

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