# -----------------------------------------------------------------------
#
# !\file CNP_random_resolution.py
# !\brief Generate random image resolutions based on a base unit and various modes.
# !\author Colin J.D. Stewart
#
# Copyright (c) 2026 Colin J.D. Stewart. All rights reserved.
#
# -----------------------------------------------------------------------

import random

class CNPRandomResolutionNode:
    """
    Generate image width/height from multiples of a base unit.
    Supports fixed, fully random, and ratio-based random modes.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mode": (["fixed", "random", "random_ratio"], {
                    "default": "random_ratio"
                }),
                "base_unit": ("INT", {
                    "default": 512,
                    "min": 64,
                    "max": 4096,
                    "step": 64
                }),
                "min_steps": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 32,
                    "step": 1
                }),
                "max_steps": ("INT", {
                    "default": 4,
                    "min": 1,
                    "max": 32,
                    "step": 1
                }),
                "fixed_width_steps": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 32,
                    "step": 1
                }),
                "fixed_height_steps": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 32,
                    "step": 1
                }),
                "orientation": (["any", "square", "portrait", "landscape"], {
                    "default": "any"
                }),
                "ratio_mode": (["free", "common", "vertical_common", "horizontal_common"], {
                    "default": "common"
                }),
                "randomise_orientation": ("BOOLEAN", {
                    "default": True
                }),
                "seed": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 0xffffffffffffffff
                }),
            }
        }

    RETURN_TYPES = ("INT", "INT", "STRING")
    RETURN_NAMES = ("width", "height", "resolution")
    FUNCTION = "generate"
    CATEGORY = "ComfyNodesPlus/Utils"

    def generate(
        self,
        mode,
        base_unit,
        min_steps,
        max_steps,
        fixed_width_steps,
        fixed_height_steps,
        orientation,
        ratio_mode,
        randomise_orientation,
        seed
    ):
        if min_steps > max_steps:
            min_steps, max_steps = max_steps, min_steps

        rng = random.Random(seed)

        if mode == "fixed":
            width = fixed_width_steps * base_unit
            height = fixed_height_steps * base_unit
            width, height = self._apply_orientation_hint(
                width,
                height,
                orientation,
                rng,
                randomise_orientation
            )
            return (width, height, f"{width} x {height}")

        if mode == "random":
            width_steps = rng.randint(min_steps, max_steps)
            height_steps = rng.randint(min_steps, max_steps)

            width = width_steps * base_unit
            height = height_steps * base_unit

            width, height = self._apply_orientation_hint(
                width,
                height,
                orientation,
                rng,
                randomise_orientation
            )
            return (width, height, f"{width} x {height}")

        # mode == "random_ratio"
        width, height = self._generate_ratio_size(
            rng=rng,
            base_unit=base_unit,
            min_steps=min_steps,
            max_steps=max_steps,
            orientation=orientation,
            ratio_mode=ratio_mode,
            randomise_orientation=randomise_orientation
        )

        return (width, height, f"{width} x {height}")

    def _apply_orientation_hint(self, width, height, orientation, rng, randomise_orientation):
        if orientation == "square":
            side = min(width, height)
            return side, side

        if orientation == "portrait":
            if width > height:
                return height, width
            return width, height

        if orientation == "landscape":
            if height > width:
                return height, width
            return width, height

        # any
        if randomise_orientation and width != height and rng.choice([True, False]):
            return height, width

        return width, height

    def _generate_ratio_size(
        self,
        rng,
        base_unit,
        min_steps,
        max_steps,
        orientation,
        ratio_mode,
        randomise_orientation
    ):
        presets = self._build_presets(
            base_unit=base_unit,
            min_steps=min_steps,
            max_steps=max_steps,
            orientation=orientation,
            ratio_mode=ratio_mode,
            randomise_orientation=randomise_orientation
        )

        if not presets:
            # Safe fallback
            steps = min_steps
            width = steps * base_unit
            height = steps * base_unit
            return width, height

        width, height = rng.choice(presets)
        return width, height

    def _build_presets(
        self,
        base_unit,
        min_steps,
        max_steps,
        orientation,
        ratio_mode,
        randomise_orientation
    ):
        presets = set()

        # Always allow square unless specifically filtered to only vertical/horizontal common
        if ratio_mode in ("free", "common"):
            for steps in range(min_steps, max_steps + 1):
                side = steps * base_unit
                presets.add((side, side))

        if ratio_mode == "free":
            # Any combination within range
            for w_steps in range(min_steps, max_steps + 1):
                for h_steps in range(min_steps, max_steps + 1):
                    presets.add((w_steps * base_unit, h_steps * base_unit))

        else:
            # Common sensible shapes using only integer step pairs
            common_pairs = [
                (1, 2),  # 512 x 1024
                (2, 3),  # 1024 x 1536
                (2, 4),  # 1024 x 2048
                (3, 4),  # 1536 x 2048
                (2, 1),  # 1024 x 512
                (3, 2),  # 1536 x 1024
                (4, 2),  # 2048 x 1024
                (4, 3),  # 2048 x 1536
            ]

            filtered_pairs = []
            for w_steps, h_steps in common_pairs:
                if w_steps < min_steps or w_steps > max_steps:
                    continue
                if h_steps < min_steps or h_steps > max_steps:
                    continue

                if ratio_mode == "vertical_common" and not (h_steps > w_steps):
                    continue
                if ratio_mode == "horizontal_common" and not (w_steps > h_steps):
                    continue

                filtered_pairs.append((w_steps, h_steps))

            for w_steps, h_steps in filtered_pairs:
                presets.add((w_steps * base_unit, h_steps * base_unit))

        # Orientation filter
        filtered = set()
        for width, height in presets:
            if orientation == "square" and width != height:
                continue
            if orientation == "portrait" and not (height > width):
                continue
            if orientation == "landscape" and not (width > height):
                continue
            filtered.add((width, height))

        presets = filtered

        # If orientation is any and randomise_orientation is enabled,
        # mirror non-square sizes so both vertical/horizontal can appear.
        if orientation == "any" and randomise_orientation:
            mirrored = set(presets)
            for width, height in presets:
                mirrored.add((height, width))
            presets = mirrored

        return sorted(presets)