"""
AI Image & Video Upscaler – Gradio UI
Powered by Real-ESRGAN
"""

import os
import tempfile

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

from upscaler import Upscaler

# Try to import the interactive slider; fall back to plain images
try:
    from gradio_imageslider import ImageSlider

    HAS_SLIDER = True
except ImportError:
    HAS_SLIDER = False

engine = Upscaler()


# ── helpers ────────────────────────────────────────────────────────────
def _pil_from_bgr(arr: np.ndarray) -> Image.Image:
    return Image.fromarray(cv2.cvtColor(arr, cv2.COLOR_BGR2RGB))


def _parse_inputs(scale_str: str, model_str: str):
    scale = int(scale_str.replace("x", ""))
    mtype = "anime" if "Anime" in model_str else "photo"
    return scale, mtype


# ── image handler ──────────────────────────────────────────────────────
def upscale_image(img_path, scale_str, model_str, progress=gr.Progress()):
    if img_path is None:
        gr.Warning("Upload an image first.")
        return None, None, None, ""

    scale, mtype = _parse_inputs(scale_str, model_str)

    progress(0.10, desc="Loading model…")
    progress(0.25, desc="Upscaling – this may take a moment…")
    output_bgr = engine.upscale_image(img_path, scale=scale, model_type=mtype)
    progress(0.90, desc="Preparing preview…")

    original_bgr = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
    # Flatten alpha channel so cvtColor always gets a 3-channel BGR image
    if original_bgr.ndim == 3 and original_bgr.shape[2] == 4:
        original_bgr = cv2.cvtColor(original_bgr, cv2.COLOR_BGRA2BGR)
    oh, ow = original_bgr.shape[:2]
    nh, nw = output_bgr.shape[:2]

    # save full-res result for download
    out_file = os.path.join(tempfile.gettempdir(), "upscaled_output.png")
    cv2.imwrite(out_file, output_bgr)

    upscaled_pil = _pil_from_bgr(output_bgr)
    original_pil = _pil_from_bgr(original_bgr).resize((nw, nh), Image.LANCZOS)

    progress(1.0, desc="Done!")
    status = f"✅  {ow}×{oh}  →  {nw}×{nh}  ({scale}x {mtype})"

    if HAS_SLIDER:
        return upscaled_pil, (original_pil, upscaled_pil), out_file, status
    return upscaled_pil, None, out_file, status


# ── video handler ──────────────────────────────────────────────────────
def upscale_video(video_path, scale_str, model_str, progress=gr.Progress()):
    if video_path is None:
        gr.Warning("Upload a video first.")
        return None, ""

    scale, mtype = _parse_inputs(scale_str, model_str)
    out_file = os.path.join(tempfile.gettempdir(), "upscaled_video.mp4")

    def on_progress(pct):
        progress(pct, desc=f"Upscaling frames… {int(pct * 100)}%")

    progress(0.05, desc="Loading model…")
    engine.upscale_video(
        video_path, out_file, scale=scale, model_type=mtype, progress_cb=on_progress
    )
    progress(1.0, desc="Done!")
    return out_file, "✅  Video upscaled successfully!"


# ── UI ─────────────────────────────────────────────────────────────────
custom_css = """
#title { text-align: center; }
.gr-button-primary { min-height: 48px !important; font-size: 1.1em !important; }
"""

with gr.Blocks(title="AI Upscaler") as demo:
    gr.Markdown(
        "# 🔍 AI Image & Video Upscaler\n"
        "*Powered by Real-ESRGAN  •  100 % Local & Private*",
        elem_id="title",
    )

    with gr.Tabs():
        # ─── Image tab ─────────────────────────────────────────────
        with gr.Tab("🖼️  Image"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=1, min_width=300):
                    img_input = gr.Image(
                        label="Upload Image",
                        type="filepath",
                        sources=["upload", "clipboard"],
                        height=300,
                    )
                    img_scale = gr.Radio(
                        ["2x", "4x"],
                        label="Upscale Factor",
                        value="4x",
                        info="Higher = bigger output & longer processing",
                    )
                    img_model = gr.Dropdown(
                        ["General (Photo)", "Anime / Illustration"],
                        label="AI Model",
                        value="General (Photo)",
                        info="Pick the best match for your content",
                    )
                    img_btn = gr.Button("🚀 Upscale Image", variant="primary", size="lg")
                    img_status = gr.Textbox(label="Status", interactive=False, lines=1)

                with gr.Column(scale=2, min_width=400):
                    img_output = gr.Image(
                        label="Upscaled Result",
                        height=400,
                    )
                    img_download = gr.File(label="Download Full Resolution")

            # before / after comparison
            gr.Markdown("### 🔄 Before / After Comparison")
            if HAS_SLIDER:
                img_slider = ImageSlider(
                    label="Drag the slider to compare",
                    type="pil",
                )
            else:
                gr.Markdown(
                    "*Install `gradio_imageslider` for an interactive comparison slider.*"
                )
                img_slider = gr.Image(label="Comparison (install gradio_imageslider for slider)", visible=False)

            img_btn.click(
                fn=upscale_image,
                inputs=[img_input, img_scale, img_model],
                outputs=[img_output, img_slider, img_download, img_status],
            )

        # ─── Video tab ─────────────────────────────────────────────
        with gr.Tab("🎬  Video"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=1, min_width=300):
                    vid_input = gr.Video(label="Upload Video", height=300)
                    vid_scale = gr.Radio(
                        ["2x", "4x"],
                        label="Upscale Factor",
                        value="2x",
                        info="2x recommended for video (much faster)",
                    )
                    vid_model = gr.Dropdown(
                        ["General (Photo)", "Anime / Illustration"],
                        label="AI Model",
                        value="General (Photo)",
                    )
                    vid_btn = gr.Button("🚀 Upscale Video", variant="primary", size="lg")
                    vid_status = gr.Textbox(label="Status", interactive=False, lines=1)

                with gr.Column(scale=2, min_width=400):
                    vid_output = gr.Video(label="Upscaled Video", height=400)

            vid_btn.click(
                fn=upscale_video,
                inputs=[vid_input, vid_scale, vid_model],
                outputs=[vid_output, vid_status],
            )

    gr.Markdown(
        "---\n"
        "Models are downloaded automatically on first use (~65 MB each).  \n"
        "GPU (CUDA) is used when available; otherwise falls back to CPU."
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)