core.py

import tempfile
from enum import Enum

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


class TargetResolution(str, Enum):
    FOUR_K = "4K (3840×2160)"
    TWO_K = "2K (2560×1440)"
    TEN_EIGHTY = "1080p (1920×1080)"
    TWO_X = "2x Original"


def get_dimensions(orig_w: int, orig_h: int, target: TargetResolution) -> tuple[int, int]:
    """Calculate target dimensions based on original size and target resolution."""
    if target == TargetResolution.FOUR_K:
        max_w, max_h = 3840, 2160
    elif target == TargetResolution.TWO_K:
        max_w, max_h = 2560, 1440
    elif target == TargetResolution.TEN_EIGHTY:
        max_w, max_h = 1920, 1080
    else:  # TWO_X
        return orig_w * 2, orig_h * 2

    scale_w = max_w / orig_w
    scale_h = max_h / orig_h
    scale = min(scale_w, scale_h)

    return int(orig_w * scale), int(orig_h * scale)


def load_upscaler(scale: int = 4) -> "RealESRGANer | None":
    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/download/v0.1.0/RealESRGAN_x{scale}plus.pth",
        model=model,
        tile=512,
        tile_pad=10,
        pre_pad=0,
        half=False,  # CPU-safe
    )
    return upscaler


def opencv_upscale(img_np: np.ndarray, target_w: int, target_h: int) -> np.ndarray:
    return cv2.resize(img_np, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)


def upscale_image(
    image: Image.Image | None, method: str, target_str: str
) -> tuple[Image.Image | None, str, str | None]:

    if image is None:
        return None, "❌ No image provided.", None

    try:
        target = TargetResolution(target_str)
    except ValueError:
        target = TargetResolution.FOUR_K

    orig_w, orig_h = image.size
    img_np = np.array(image)

    target_w, target_h = get_dimensions(orig_w, orig_h, target)

    if method == "Real-ESRGAN (AI)" and REALESRGAN_AVAILABLE:
        try:
            upscaler = load_upscaler(scale=4)
            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)
                method = "Lanczos (fallback: ESRGAN uninitialized)"
        except Exception as e:
            output = opencv_upscale(img_np, target_w, target_h)
            method = f"Lanczos (fallback: {str(e)[:60]})"
    else:
        output = opencv_upscale(img_np, target_w, target_h)
        method = "Lanczos (Fast)"

    result_img = Image.fromarray(output)
    info = f"✅ Upscaled: {orig_w}×{orig_h} → {target_w}×{target_h} | Method: {method}"

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        result_img.save(tmp.name, format="PNG")
        tmp_path = tmp.name

    return result_img, info, tmp_path


def upscale_video(
    video_path: str | None, method: str, target_str: str, progress_callback=None
) -> tuple[str | None, str]:

    if not video_path:
        return None, "❌ No video provided."

    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))

        try:
            target = TargetResolution(target_str)
        except ValueError:
            target = TargetResolution.FOUR_K

        target_w, target_h = get_dimensions(orig_w, orig_h, target)

        out_path = tempfile.mktemp(suffix=".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_callback:
                    progress_callback(
                        (frame_idx / total_frames), f"Processing frame {frame_idx}/{total_frames}"
                    )

            info = f"✅ Video upscaled: {orig_w}×{orig_h} → {target_w}×{target_h} | {total_frames} frames @ {fps:.1f}fps"
        finally:
            writer.release()
    finally:
        cap.release()

    return out_path, info