# app.py
# AI Video Enhancer 4K - Optimized single GPU call with in-memory processing

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

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
from spandrel import ImageModelDescriptor, ModelLoader

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

# Pre-download models at startup
MODEL_PATHS = {}

def ensure_model(scale: int) -> str:
    """Download model weights (CPU/network only)."""
    if scale not in MODEL_PATHS:
        if scale == 2:
            MODEL_PATHS[scale] = hf_hub_download(
                repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x2.pth"
            )
        else:
            MODEL_PATHS[scale] = hf_hub_download(
                repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth"
            )
    return MODEL_PATHS[scale]

try:
    ensure_model(2)
    ensure_model(4)
    print("Models pre-downloaded successfully.")
except Exception as e:
    print(f"Model pre-download skipped: {e}")


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, max_frames: int = None):
    frames_dir.mkdir(parents=True, exist_ok=True)
    cmd = ["ffmpeg", "-y", "-i", video_path, "-vsync", "0"]
    if max_frames:
        cmd.extend(["-vframes", str(max_frames)])
    cmd.append(str(frames_dir / "%06d.png"))
    run_cmd(cmd)


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:
    h, w = img.shape[:2]
    return cv2.resize(img, (w * scale, h * scale), interpolation=cv2.INTER_CUBIC)


def load_frames_to_memory(frames_dir: Path) -> list:
    """Load all frames into RAM (CPU work, not billed)."""
    frame_files = sorted(frames_dir.glob("*.png"))
    frames = []
    for fp in frame_files:
        img = cv2.imread(str(fp))
        if img is not None:
            frames.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    return frames


def save_frames_from_memory(frames: list, frames_dir: Path):
    """Write enhanced frames back to disk (CPU work, not billed)."""
    for idx, img_rgb in enumerate(frames):
        out_path = frames_dir / f"{idx + 1:06d}.png"
        cv2.imwrite(str(out_path), cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR))


@spaces.GPU(duration=180)
def enhance_all_frames_gpu(frames_rgb: list, model_path: str, scale: int = 4) -> list:
    """
    Enhance ALL frames in a SINGLE GPU call.
    Frames are already in memory (no disk I/O here).
    Uses torch.inference_mode() for faster inference with identical quality.
    """
    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()

    total = len(frames_rgb)
    print(f"Model loaded on {device}, processing {total} frames in memory...")

    enhanced = []
    with torch.inference_mode():
        for idx, img_rgb in enumerate(frames_rgb):
            tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).float().div(255.0)
            tensor = tensor.unsqueeze(0).to(device)

            output = model(tensor)
            output = output.squeeze(0).cpu().clamp(0, 1).mul(255).byte()
            output = output.permute(1, 2, 0).numpy()
            enhanced.append(output)

            if (idx + 1) % 10 == 0:
                print(f"Processed {idx + 1}/{total}")

    return enhanced


def process_video(video_file, scale: int = 4, progress=gr.Progress()) -> Tuple[str, str]:
    """Main video processing - single GPU call with in-memory frame processing."""
    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

    max_seconds = 10
    max_frames = int(fps * max_seconds)

    progress(0.05, f"Video: {w}x{h} @ {fps:.1f}fps, extracting up to {max_seconds}s...")

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

    num_frames = len(list(frames_dir.glob("*.png")))
    progress(0.15, f"Loading {num_frames} frames into memory...")

    # Load all frames into RAM (CPU work, no GPU needed)
    frames_rgb = load_frames_to_memory(frames_dir)
    if not frames_rgb:
        shutil.rmtree(base_dir, ignore_errors=True)
        return "No frames extracted.", None

    # Ensure model weights are cached (CPU/network only)
    progress(0.20, "Preparing model...")
    model_path = ensure_model(scale)

    progress(0.25, f"Enhancing {len(frames_rgb)} frames on GPU (single call)...")

    use_fallback = False
    try:
        # === THE SINGLE GPU CALL ===
        enhanced_frames = enhance_all_frames_gpu(frames_rgb, model_path, scale)
        print(f"Enhanced {len(enhanced_frames)} frames with Real-ESRGAN")
    except Exception as e:
        print(f"GPU enhancement failed: {e}")
        print("Using fallback bicubic upscaling...")
        use_fallback = True
        enhanced_frames = [simple_upscale(f, scale) for f in frames_rgb]

    progress(0.80, "Writing enhanced frames...")
    save_frames_from_memory(enhanced_frames, frames_dir)

    # Free memory
    del frames_rgb, enhanced_frames

    progress(0.85, "Reassembling 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}, {num_frames} frames)", 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("**NOTE:** The video processing may appear stuck at 25%. It is NOT stuck; it is working in the background, and will directly jump to 80% once it's finished!")

    # 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 (max 10 sec processed)",
                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(
        "**Limit: 10 seconds of video** (ZeroGPU quota). "
        "Log in to HuggingFace for more!"
    )

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

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