"""FikoYT Image Upscale — Real-ESRGAN / ESRGAN / SwinIR auto-download upscaler."""
import os

try:
    import torch
    import comfy.utils
    from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel, UpscaleModelLoader
except ImportError:
    torch = None
    ImageUpscaleWithModel = None
    UpscaleModelLoader = None

from ..core import ensure_model, get_models_dir


_PRESETS = {
    "RealESRGAN x4plus":    {"repo": "ai-forever/Real-ESRGAN", "file": "RealESRGAN_x4plus.pth"},
    "RealESRGAN x2plus":    {"repo": "ai-forever/Real-ESRGAN", "file": "RealESRGAN_x2plus.pth"},
    "RealESRGAN x4plus Anime": {"repo": "ai-forever/Real-ESRGAN", "file": "RealESRGAN_x4plus_anime_6B.pth"},
    "4x-UltraSharp":        {"repo": "uwg/upscaler", "file": "ESRGAN/4x-UltraSharp.pth"},
    "4x_NMKD-Siax":         {"repo": "uwg/upscaler", "file": "ESRGAN/4x_NMKD-Siax_200k.pth"},
}


class FikoYTImageUpscale:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "model_preset": (list(_PRESETS.keys()), {"default": "RealESRGAN x4plus"}),
                "auto_download": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "scale_after": ("FLOAT", {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.05}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"
    CATEGORY = "FikoYT/image"

    def upscale(self, image, model_preset, auto_download, scale_after=1.0):
        if torch is None or ImageUpscaleWithModel is None:
            raise RuntimeError("torch veya comfy_extras yüklü değil")

        meta = _PRESETS[model_preset]
        upscale_dir = get_models_dir("upscale_models")
        ckpt_path = os.path.join(upscale_dir, os.path.basename(meta["file"]))
        if not os.path.isfile(ckpt_path) and auto_download:
            ckpt_path = ensure_model(meta["repo"], meta["file"], "upscale_models")
        if not os.path.isfile(ckpt_path):
            raise FileNotFoundError(f"Upscale modeli yok: {ckpt_path}")

        loader = UpscaleModelLoader()
        applier = ImageUpscaleWithModel()
        upscale_model = loader.load_model(os.path.basename(ckpt_path))[0]
        upscaled = applier.upscale(upscale_model, image)[0]

        if abs(scale_after - 1.0) > 1e-3:
            n, h, w, c = upscaled.shape
            new_h = max(8, int(h * scale_after))
            new_w = max(8, int(w * scale_after))
            up_t = upscaled.permute(0, 3, 1, 2)
            up_t = torch.nn.functional.interpolate(up_t, size=(new_h, new_w), mode="bicubic", align_corners=False)
            upscaled = up_t.permute(0, 2, 3, 1).clamp(0, 1)

        return (upscaled,)


NODE_CLASS_MAPPINGS = {"FikoYTImageUpscale": FikoYTImageUpscale}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTImageUpscale": "FikoYT • Image Upscale"}