import os
import subprocess
import tempfile

import numpy as np
import torch
from PIL import Image, ImageEnhance

from .palettes import PALETTE_NAMES, get_palette_rgb

# The compiled engine lives next to this file, inside the custom node folder, so it
# works standalone on any Windows machine. To use your own build, just replace this
# .exe in the node folder.
_HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_EXE = os.path.join(_HERE, "spritefusion-pixel-snapper.exe")


def tensor_to_pil(img_tensor):
    arr = (img_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
    return Image.fromarray(arr, mode="RGB")


def pil_to_tensor(img):
    arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
    return torch.from_numpy(arr)


def remap_to_palette(img, rgb_palette, dither):
    """Reduce the image to the exact colors of a fixed palette."""
    pal_img = Image.new("P", (1, 1))
    flat = []
    for r, g, b in rgb_palette:
        flat.extend([r, g, b])
    flat.extend([0, 0, 0] * (256 - len(rgb_palette)))
    pal_img.putpalette(flat)
    dither_mode = Image.FLOYDSTEINBERG if dither else Image.NONE
    return img.convert("RGB").quantize(
        palette=pal_img, colors=len(rgb_palette), dither=dither_mode
    ).convert("RGB")


def correct_aspect(img, src_size):
    """Fix the 'flattening' from the engine: the edge-walker can emit a slightly
    different number of cells on X vs Y, distorting the aspect ratio. We trim the
    over-represented axis (nearest) so the output aspect matches the source."""
    ow, oh = img.size
    sw, sh = src_size
    if min(ow, oh, sw, sh) <= 0:
        return img
    src_aspect = sw / sh
    out_aspect = ow / oh
    if abs(out_aspect - src_aspect) < 1e-3:
        return img
    if out_aspect > src_aspect:
        new_w, new_h = max(1, round(oh * src_aspect)), oh   # too wide -> trim width
    else:
        new_w, new_h = ow, max(1, round(ow / src_aspect))   # too tall -> trim height
    if (new_w, new_h) == (ow, oh):
        return img
    return img.resize((new_w, new_h), Image.NEAREST)


def run_snapper(exe_path, pil_in, k):
    """Save the image, run the spritefusion-pixel-snapper binary and read the result."""
    if not (isinstance(exe_path, str) and os.path.isfile(exe_path)):
        raise FileNotFoundError(
            f"Pixel Snapper engine not found. Place 'spritefusion-pixel-snapper.exe' "
            f"in the node folder: {_HERE}"
        )
    tmpdir = tempfile.mkdtemp(prefix="pixelsnap_")
    in_path = os.path.join(tmpdir, "in.png")
    out_path = os.path.join(tmpdir, "out.png")
    pil_in.save(in_path)
    result = subprocess.run(
        [exe_path, in_path, out_path, str(int(k))],
        capture_output=True, text=True,
    )
    if result.returncode != 0 or not os.path.isfile(out_path):
        raise RuntimeError(
            f"Pixel Snapper failed (code {result.returncode}).\n"
            f"stdout: {result.stdout}\nstderr: {result.stderr}"
        )
    out_img = Image.open(out_path).convert("RGB").copy()
    return out_img


class PixelSnapperNode:
    """Calls the compiled engine 'spritefusion-pixel-snapper.exe' which detects the
    native pixel art grid and reduces the image to k colors. Supports color
    adjustments (pre) and remapping to a fixed video game palette (post)."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "colors_k": ("INT", {"default": 16, "min": 2, "max": 256, "step": 1}),
                "saturation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
                "contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
                "brightness": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
                "fix_aspect": ("BOOLEAN", {"default": True}),
                "keep_output_size": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "palette": ("GAME_PALETTE",),
                "palette_dither": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "run"
    CATEGORY = "image/pixel_art"

    def run(self, image, colors_k, saturation, contrast, brightness,
            fix_aspect, keep_output_size, palette=None, palette_dither=False):
        # If a palette is connected, the k passed to the engine is the palette size.
        k = len(palette) if palette is not None else colors_k

        out_batch = []
        target_size = None
        for i in range(image.shape[0]):
            img = tensor_to_pil(image[i])
            orig_size = img.size

            # Color adjustments BEFORE the snap (they affect the engine's k-means).
            if brightness != 1.0:
                img = ImageEnhance.Brightness(img).enhance(brightness)
            if contrast != 1.0:
                img = ImageEnhance.Contrast(img).enhance(contrast)
            if saturation != 1.0:
                img = ImageEnhance.Color(img).enhance(saturation)

            result = run_snapper(DEFAULT_EXE, img, k)

            # Fix the engine's aspect-ratio flattening (edge-walker cell mismatch).
            if fix_aspect:
                result = correct_aspect(result, orig_size)

            # Remap to a fixed palette AFTER the snap (exact video game colors).
            if palette is not None:
                result = remap_to_palette(result, palette, palette_dither)

            if keep_output_size:
                result = result.resize(orig_size, Image.NEAREST)

            # Ensure consistent size across the batch (torch.stack requires it).
            if target_size is None:
                target_size = result.size
            elif result.size != target_size:
                result = result.resize(target_size, Image.NEAREST)

            out_batch.append(pil_to_tensor(result))

        return (torch.stack(out_batch),)


class VideoGamePaletteNode:
    """Provides a classic video game color palette for the Pixel Snapper."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "palette_name": (PALETTE_NAMES, {"default": PALETTE_NAMES[0]}),
            },
        }

    RETURN_TYPES = ("GAME_PALETTE",)
    RETURN_NAMES = ("palette",)
    FUNCTION = "run"
    CATEGORY = "image/pixel_art"

    def run(self, palette_name):
        return (get_palette_rgb(palette_name),)


RESAMPLE = {
    "nearest": Image.NEAREST,
    "box": Image.BOX,
    "bilinear": Image.BILINEAR,
    "hamming": Image.HAMMING,
    "bicubic": Image.BICUBIC,
    "lanczos": Image.LANCZOS,
}


class PixelArtResizerNode:
    """Resizes pixel art. To scale up without blurring use 'nearest'; to shrink a
    large image down to pixel art resolution use 'box' or 'nearest'."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "mode": (["scale_factor", "target_size"], {"default": "scale_factor"}),
                "scale_factor": ("FLOAT", {"default": 2.0, "min": 0.05, "max": 64.0, "step": 0.05}),
                "width": ("INT", {"default": 128, "min": 1, "max": 8192, "step": 1}),
                "height": ("INT", {"default": 128, "min": 1, "max": 8192, "step": 1}),
                "interpolation": (list(RESAMPLE.keys()), {"default": "nearest"}),
                "keep_aspect": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "run"
    CATEGORY = "image/pixel_art"

    def run(self, image, mode, scale_factor, width, height, interpolation, keep_aspect):
        filt = RESAMPLE[interpolation]
        out_batch = []
        target = None
        for i in range(image.shape[0]):
            img = tensor_to_pil(image[i])
            w, h = img.size

            if mode == "scale_factor":
                nw = max(1, round(w * scale_factor))
                nh = max(1, round(h * scale_factor))
            else:
                if keep_aspect:
                    ratio = min(width / w, height / h)
                    nw = max(1, round(w * ratio))
                    nh = max(1, round(h * ratio))
                else:
                    nw, nh = int(width), int(height)

            result = img.resize((nw, nh), filt)

            # Batch needs a consistent size for torch.stack.
            if target is None:
                target = result.size
            elif result.size != target:
                result = result.resize(target, filt)

            out_batch.append(pil_to_tensor(result))

        return (torch.stack(out_batch),)


class PixelArtPreviewNode:
    """Pixel art preview with crisp rendering (nearest-neighbor) and an interactive
    magnifier. Hover over the image to zoom in. Output node."""

    def __init__(self):
        self._prefix = "pixelart_preview"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "magnifier": ("BOOLEAN", {"default": True}),
                "zoom": ("INT", {"default": 6, "min": 2, "max": 32, "step": 1}),
                "lens_size": ("INT", {"default": 140, "min": 40, "max": 400, "step": 10}),
            },
            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
        }

    RETURN_TYPES = ()
    FUNCTION = "preview"
    OUTPUT_NODE = True
    CATEGORY = "image/pixel_art"

    def preview(self, images, magnifier, zoom, lens_size, prompt=None, extra_pnginfo=None):
        import folder_paths

        output_dir = folder_paths.get_temp_directory()
        full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
            self._prefix, output_dir, images[0].shape[1], images[0].shape[0]
        )
        results = []
        for image in images:
            arr = (255.0 * image.cpu().numpy()).clip(0, 255).astype(np.uint8)
            img = Image.fromarray(arr)
            file = f"{filename}_{counter:05}_.png"
            img.save(os.path.join(full_output_folder, file), compress_level=1)
            results.append({"filename": file, "subfolder": subfolder, "type": "temp"})
            counter += 1

        return {
            "ui": {
                "images": results,
                "magnifier": [magnifier],
                "zoom": [zoom],
                "lens_size": [lens_size],
            }
        }


NODE_CLASS_MAPPINGS = {
    "PixelSnapper": PixelSnapperNode,
    "VideoGamePalette": VideoGamePaletteNode,
    "PixelArtResizer": PixelArtResizerNode,
    "PixelArtPreview": PixelArtPreviewNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PixelSnapper": "Pixel Snapper (Engine)",
    "VideoGamePalette": "Video Game Palette",
    "PixelArtResizer": "Pixel Art Resizer",
    "PixelArtPreview": "Pixel Art Preview",
}