# -*- coding: utf-8 -*-
"""Inference core for PixelUP.

Model loading is handled by spandrel, which detects the architecture from the
checkpoint itself. Everything below is the part spandrel deliberately leaves to
the caller: image <-> tensor conversion, tiled inference, and alpha handling.
"""

from __future__ import annotations

import gc
import math
import os
import threading
from dataclasses import dataclass, field
from functools import lru_cache

import numpy as np
import torch
from PIL import Image
from spandrel import ImageModelDescriptor, ModelLoader, ModelTiling

WEIGHTS_DIR = os.environ.get("PIXELUP_WEIGHTS", os.path.join(os.path.abspath("."), "weights"))

# Tiles are square, in input pixels. Anything larger than this is split.
DEFAULT_TILE = 512
TILE_OVERLAP = 32
MIN_TILE = 64


@dataclass(frozen=True)
class ModelSpec:
    label: str
    url: str
    # Second URL for models shipping a paired denoise checkpoint (general v3).
    wdn_url: str | None = None
    tags: tuple[str, ...] = field(default_factory=tuple)


MODELS: dict[str, ModelSpec] = {
    "realesr-general-x4v3": ModelSpec(
        label="General v3 · recommended",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
        wdn_url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
        tags=("fast",),
    ),
    "RealESRGAN_x4plus": ModelSpec(
        label="x4 Plus · photos",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
    ),
    "RealESRNet_x4plus": ModelSpec(
        label="x4 Net · smoother, fewer artifacts",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth",
    ),
    "RealESRGAN_x4plus_anime_6B": ModelSpec(
        label="x4 Anime · illustrations, line art",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
    ),
    "RealESRGAN_x2plus": ModelSpec(
        label="x2 Plus · lighter, 2x native",
        url="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
    ),
}

GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth"

# Guards weight downloads so two concurrent requests can't write the same file.
_download_lock = threading.Lock()


# --------------------------------------------------------------------------- #
# device / dtype
# --------------------------------------------------------------------------- #

def get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def pick_dtype(model: ImageModelDescriptor, device: torch.device) -> torch.dtype:
    """fp16 on CUDA when the architecture supports it; fp32 everywhere else.

    MPS is deliberately kept at fp32 — half precision on Metal still produces
    NaNs in some conv paths, and the speedup there is small.
    """
    if device.type == "cuda" and model.supports_half:
        return torch.float16
    return torch.float32


# --------------------------------------------------------------------------- #
# weights
# --------------------------------------------------------------------------- #

def download_weight(url: str) -> str:
    """Downloads url into WEIGHTS_DIR if absent, returns the local path."""
    os.makedirs(WEIGHTS_DIR, exist_ok=True)
    filename = os.path.basename(url.split("?")[0])
    path = os.path.join(WEIGHTS_DIR, filename)
    if os.path.isfile(path):
        return path
    with _download_lock:
        # Re-check: another thread may have finished while we waited.
        if not os.path.isfile(path):
            torch.hub.download_url_to_file(url, path, progress=True)
    return path


def supports_denoise(model_name: str) -> bool:
    """Only the general v3 model ships the paired checkpoint DNI needs."""
    return MODELS[model_name].wdn_url is not None if model_name in MODELS else False


def _blend_state_dicts(a: dict, b: dict, weight: float) -> dict:
    """Deep network interpolation: a * weight + b * (1 - weight).

    This is how Real-ESRGAN exposed denoise strength — blending the standard
    checkpoint with its with-denoise twin rather than switching between them.
    """
    return {
        key: value * weight + b[key] * (1 - weight) if key in b else value
        for key, value in a.items()
    }


@lru_cache(maxsize=3)
def load_model(model_name: str, denoise: float = 1.0) -> ImageModelDescriptor:
    """Loads and caches a model on the active device.

    Cached because the previous implementation rebuilt the network and re-read
    the checkpoint from disk on every single request. `denoise` is only honoured
    for models with a paired wdn checkpoint; it is part of the cache key.
    """
    if model_name not in MODELS:
        raise ValueError(f"Unknown model: {model_name}")

    spec = MODELS[model_name]
    loader = ModelLoader()
    path = download_weight(spec.url)

    if spec.wdn_url is not None and denoise != 1.0:
        wdn_path = download_weight(spec.wdn_url)
        blended = _blend_state_dicts(
            loader.load_state_dict_from_file(path),
            loader.load_state_dict_from_file(wdn_path),
            denoise,
        )
        model = loader.load_from_state_dict(blended)
    else:
        model = loader.load_from_file(path)

    if not isinstance(model, ImageModelDescriptor):
        raise ValueError(f"{model_name} is not an image-to-image model")

    device = get_device()
    model.to(device, pick_dtype(model, device)).eval()
    return model


@lru_cache(maxsize=1)
def load_face_model() -> ImageModelDescriptor:
    """Loads GFPGAN v1.4 through spandrel (no gfpgan/basicsr packages)."""
    model = ModelLoader().load_from_file(download_weight(GFPGAN_URL))
    if not isinstance(model, ImageModelDescriptor):
        raise ValueError("GFPGAN checkpoint did not load as an image model")
    device = get_device()
    model.to(device, pick_dtype(model, device)).eval()
    return model


def clear_model_cache() -> None:
    load_model.cache_clear()
    load_face_model.cache_clear()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# --------------------------------------------------------------------------- #
# tensor conversion
# --------------------------------------------------------------------------- #

def to_tensor(array: np.ndarray, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """HWC uint8 (or HW for single channel) -> (1, C, H, W) float in [0, 1]."""
    if array.ndim == 2:
        array = array[:, :, None]
    tensor = torch.from_numpy(np.ascontiguousarray(array.transpose(2, 0, 1)))
    return tensor.unsqueeze(0).to(device=device, dtype=dtype).div_(255.0)


def to_array(tensor: torch.Tensor) -> np.ndarray:
    """(1, C, H, W) float in [0, 1] -> HWC uint8.

    Every operation here is out-of-place: spandrel runs models under
    `torch.inference_mode`, and in-place updates to the tensors it returns raise
    outside that context.
    """
    out = tensor.detach().float().clamp(0, 1).squeeze(0)
    out = (out * 255.0).round().to(dtype=torch.uint8, device="cpu")
    return out.permute(1, 2, 0).numpy()


# --------------------------------------------------------------------------- #
# tiled inference
# --------------------------------------------------------------------------- #

def _run_whole(model: ImageModelDescriptor, image: torch.Tensor) -> torch.Tensor:
    with torch.no_grad():
        return model(image)


def run_tiled(
    model: ImageModelDescriptor,
    image: torch.Tensor,
    tile: int = DEFAULT_TILE,
    overlap: int = TILE_OVERLAP,
    progress_cb=None,
) -> torch.Tensor:
    """Runs the model over the image in overlapping tiles.

    Tiles overlap and are cropped back to their non-overlapping core, so seams
    land inside the padded region rather than on the output boundary. Falls back
    to progressively smaller tiles when the device runs out of memory.
    """
    _, _, h, w = image.shape
    scale = model.scale

    if model.tiling == ModelTiling.INTERNAL or (tile <= 0) or (h <= tile and w <= tile):
        return _run_whole(model, image)

    out = torch.empty(
        (1, model.output_channels, h * scale, w * scale),
        dtype=image.dtype,
        device=image.device,
    )

    tiles_y = math.ceil(h / tile)
    tiles_x = math.ceil(w / tile)
    total = tiles_y * tiles_x
    done = 0

    for ty in range(tiles_y):
        for tx in range(tiles_x):
            # Core region this tile is responsible for.
            y0, x0 = ty * tile, tx * tile
            y1, x1 = min(y0 + tile, h), min(x0 + tile, w)

            # Padded region actually fed to the model.
            py0, px0 = max(y0 - overlap, 0), max(x0 - overlap, 0)
            py1, px1 = min(y1 + overlap, h), min(x1 + overlap, w)

            patch = image[:, :, py0:py1, px0:px1]
            result = _run_whole(model, patch)

            # Offset of the core inside the padded result, in output pixels.
            oy, ox = (y0 - py0) * scale, (x0 - px0) * scale
            ch, cw = (y1 - y0) * scale, (x1 - x0) * scale
            out[:, :, y0 * scale:y1 * scale, x0 * scale:x1 * scale] = result[
                :, :, oy:oy + ch, ox:ox + cw
            ]

            done += 1
            if progress_cb is not None:
                progress_cb(done / total)

    return out


def run_model(
    model: ImageModelDescriptor,
    array: np.ndarray,
    tile: int = DEFAULT_TILE,
    progress_cb=None,
) -> np.ndarray:
    """Runs a model over an HWC uint8 array, retrying with smaller tiles on OOM.

    Architectures spandrel marks as tiling-DISCOURAGED (they lean on global image
    context, so seams show) start whole-image and only fall back to tiles if
    memory actually runs out.
    """
    device, dtype = model.device, model.dtype
    image = to_tensor(array, device, dtype)

    current = 0 if model.tiling == ModelTiling.DISCOURAGED else tile

    while True:
        try:
            return to_array(run_tiled(model, image, tile=current, progress_cb=progress_cb))
        except (torch.cuda.OutOfMemoryError, RuntimeError) as err:
            if not _is_oom(err):
                raise
            if current <= 0:
                current = tile if tile > 0 else DEFAULT_TILE
            elif current > MIN_TILE:
                current = max(current // 2, MIN_TILE)
            else:
                raise
            gc.collect()
            if device.type == "cuda":
                torch.cuda.empty_cache()


def _is_oom(err: Exception) -> bool:
    if isinstance(err, torch.cuda.OutOfMemoryError):
        return True
    message = str(err).lower()
    return "out of memory" in message or "can't allocate" in message


# --------------------------------------------------------------------------- #
# alpha
# --------------------------------------------------------------------------- #

def split_alpha(img: Image.Image) -> tuple[np.ndarray, np.ndarray | None]:
    """Returns (rgb, alpha) as uint8 arrays; alpha is None for opaque images."""
    if img.mode == "RGBA":
        array = np.array(img)
        return array[:, :, :3], array[:, :, 3]
    if img.mode in ("LA", "PA") or (img.mode == "P" and "transparency" in img.info):
        array = np.array(img.convert("RGBA"))
        return array[:, :, :3], array[:, :, 3]
    return np.array(img.convert("RGB")), None


def upscale_alpha(model: ImageModelDescriptor, alpha: np.ndarray, tile: int) -> np.ndarray:
    """Upscales the alpha channel through the model to keep hard cutout edges.

    Bicubic would soften logo and cutout edges, which is the visible failure
    case, so the channel is replicated to RGB and run through the same network.
    """
    rgb = np.repeat(alpha[:, :, None], 3, axis=2)
    return run_model(model, rgb, tile=tile)[:, :, 0]