"""Core upscaling engine for 4K Upscaler Pro."""

import os
import tempfile

import cv2
import numpy as np
from PIL import Image

try:
    from basicsr.archs.rrdbnet_arch import RRDBNet
    from realesrgan import RealESRGANer

    REALESRGAN_AVAILABLE = True
except ImportError:
    REALESRGAN_AVAILABLE = False


# ── Resolution targets ──────────────────────────────────────────────
RESOLUTIONS = {
    "4K (3840×2160)": (3840, 2160),
    "2K (2560×1440)": (2560, 1440),
    "1080p (1920×1080)": (1920, 1080),
    "2x Original": None,
}


def get_dimensions(orig_w, orig_h, target_str):
    """Calculate target dimensions based on original size and target resolution."""
    dims = RESOLUTIONS.get(target_str)
    if dims is None:
        return orig_w * 2, orig_h * 2
    max_w, max_h = dims
    scale = min(max_w / orig_w, max_h / orig_h)
    return int(orig_w * scale), int(orig_h * scale)


# ── Model loader ────────────────────────────────────────────────────
_cached_upscaler = None


def load_upscaler(scale=4):
    """Load Real-ESRGAN model (cached after first call)."""
    global _cached_upscaler
    if _cached_upscaler is not None:
        return _cached_upscaler
    if not REALESRGAN_AVAILABLE:
        return None
    model = RRDBNet(
        num_in_ch=3, num_out_ch=3, num_feat=64,
        num_block=23, num_grow_ch=32, scale=scale,
    )
    upscaler = RealESRGANer(
        scale=scale,
        model_path=(
            f"https://github.com/xinntao/Real-ESRGAN/releases/"
            f"download/v0.1.0/RealESRGAN_x{scale}plus.pth"
        ),
        model=model,
        tile=512,
        tile_pad=10,
        pre_pad=0,
        half=False,
    )
    _cached_upscaler = upscaler
    return upscaler


def opencv_upscale(img_np, target_w, target_h):
    """High-quality Lanczos resize."""
    return cv2.resize(img_np, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)


# ── Image upscaling ────────────────────────────────────────────────
def upscale_image(image_path, method, target_str, progress=None):
    """Upscale a single image. Returns (output_path, info_text, download_path)."""
    if image_path is None:
        return None, "❌ No image provided.", None

    if progress:
        progress(0.1, desc="📂 Loading image...")

    image = Image.open(image_path).convert("RGB")
    orig_w, orig_h = image.size
    img_np = np.array(image)
    target_w, target_h = get_dimensions(orig_w, orig_h, target_str)

    if progress:
        progress(0.3, desc=f"🔍 Scaling {orig_w}×{orig_h} → {target_w}×{target_h}...")

    used_method = method

    if method == "Real-ESRGAN (AI)" and REALESRGAN_AVAILABLE:
        try:
            upscaler = load_upscaler(scale=4)
            if progress:
                progress(0.5, desc="🤖 Running Real-ESRGAN AI model...")
            if upscaler:
                output, _ = upscaler.enhance(img_np, outscale=target_w / orig_w)
                output = cv2.resize(
                    output, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4
                )
            else:
                output = opencv_upscale(img_np, target_w, target_h)
                used_method = "Lanczos (fallback)"
        except Exception as e:
            output = opencv_upscale(img_np, target_w, target_h)
            used_method = f"Lanczos (fallback: {str(e)[:50]})"
    else:
        if progress:
            progress(0.5, desc="⚡ Applying Lanczos4 interpolation...")
        output = opencv_upscale(img_np, target_w, target_h)
        used_method = "Lanczos (Fast)"

    if progress:
        progress(0.8, desc="💾 Saving result...")

    result_img = Image.fromarray(output)
    out_path = os.path.join(tempfile.gettempdir(), "upscaled_output.png")
    result_img.save(out_path, format="PNG", optimize=True)

    info = (
        f"✅ Upscaled: {orig_w}×{orig_h} → {target_w}×{target_h}\n"
        f"📐 Method: {used_method}\n"
        f"📄 Format: PNG | Size: {os.path.getsize(out_path) / 1024 / 1024:.1f} MB"
    )

    if progress:
        progress(1.0, desc="✨ Done!")

    return out_path, info, out_path


# ── Video upscaling ─────────────────────────────────────────────────
def upscale_video(video_path, method, target_str, progress=None):
    """Upscale a video file. Returns (output_path, info_text)."""
    if not video_path:
        return None, "❌ No video provided."

    if progress:
        progress(0.05, desc="📂 Opening video...")

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, "❌ Could not open video."

    try:
        orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        target_w, target_h = get_dimensions(orig_w, orig_h, target_str)

        out_path = os.path.join(tempfile.gettempdir(), "upscaled_video.mp4")
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        writer = cv2.VideoWriter(out_path, fourcc, fps, (target_w, target_h))

        if not writer.isOpened():
            return None, "❌ Could not create output video writer."

        try:
            upscaler = None
            if method == "Real-ESRGAN (AI)" and REALESRGAN_AVAILABLE:
                try:
                    upscaler = load_upscaler(scale=4)
                except Exception:
                    upscaler = None

            frame_idx = 0
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

                if upscaler:
                    try:
                        out_frame, _ = upscaler.enhance(
                            frame_rgb, outscale=target_w / orig_w
                        )
                        out_frame = cv2.resize(
                            out_frame, (target_w, target_h),
                            interpolation=cv2.INTER_LANCZOS4,
                        )
                    except Exception:
                        out_frame = opencv_upscale(frame_rgb, target_w, target_h)
                else:
                    out_frame = opencv_upscale(frame_rgb, target_w, target_h)

                writer.write(cv2.cvtColor(out_frame, cv2.COLOR_RGB2BGR))
                frame_idx += 1

                if total_frames > 0 and progress:
                    pct = 0.1 + 0.85 * (frame_idx / total_frames)
                    progress(
                        pct,
                        desc=f"🎬 Processing frame {frame_idx}/{total_frames}",
                    )

            info = (
                f"✅ Video upscaled: {orig_w}×{orig_h} → {target_w}×{target_h}\n"
                f"🎞 {total_frames} frames @ {fps:.1f} fps\n"
                f"📄 Size: {os.path.getsize(out_path) / 1024 / 1024:.1f} MB"
            )
        finally:
            writer.release()
    finally:
        cap.release()

    if progress:
        progress(1.0, desc="✨ Done!")

    return out_path, info