from math import gcd


ASPECT_RATIO_PRESETS = {
    "1:1": (1, 1),
    "16:9": (16, 9),
    "9:16": (9, 16),
    "4:3": (4, 3),
    "3:4": (3, 4),
    "3:2": (3, 2),
    "2:3": (2, 3),
    "21:9": (21, 9),
    "9:21": (9, 21),
    "5:4": (5, 4),
    "4:5": (4, 5),
    "Custom": None,
}


def normalize_ratio(width, height):
    width = max(1, int(width))
    height = max(1, int(height))
    divisor = gcd(width, height)
    return width // divisor, height // divisor


def round_to_multiple(value, multiple):
    return max(multiple, int(round(value / multiple)) * multiple)


def calculate_exact_resolution(ratio_width, ratio_height, multiple, size_step):
    ratio_width, ratio_height = normalize_ratio(ratio_width, ratio_height)
    multiple = max(1, int(multiple))
    size_step = max(1, int(size_step))
    base_width = multiple * ratio_width * size_step
    base_height = multiple * ratio_height * size_step

    return (
        base_width,
        base_height,
        base_width * 2,
        base_height * 2,
        base_width * 4,
        base_height * 4,
    )


def calculate_approx_resolution(ratio_width, ratio_height, multiple, size_step):
    ratio_width, ratio_height = normalize_ratio(ratio_width, ratio_height)
    multiple = max(1, int(multiple))
    size_step = max(1, int(size_step))
    base_side = multiple * size_step

    if ratio_width == ratio_height:
        base_width = base_side
        base_height = base_side
    elif ratio_width > ratio_height:
        base_width = base_side
        base_height = round_to_multiple(base_width * ratio_height / ratio_width, multiple)
    else:
        base_height = base_side
        base_width = round_to_multiple(base_height * ratio_width / ratio_height, multiple)

    upscale_2x_width = base_width * 2
    upscale_2x_height = base_height * 2
    upscale_4x_width = base_width * 4
    upscale_4x_height = base_height * 4

    return (
        base_width,
        base_height,
        upscale_2x_width,
        upscale_2x_height,
        upscale_4x_width,
        upscale_4x_height,
    )


class LTXResolutionHelperNode:
    CATEGORY = "LTX/Helpers"
    FUNCTION = "get_resolutions"
    RETURN_TYPES = ("INT", "INT", "STRING")
    RETURN_NAMES = ("base_width", "base_height", "summary")
    DESCRIPTION = (
        "Calculates a base resolution aligned to a chosen multiple "
        "and shows the resulting sizes after two sequential 2x upscales."
    )

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "aspect_ratio": (list(ASPECT_RATIO_PRESETS.keys()), {"default": "1:1"}),
                "ratio_mode": (["Exact Ratio", "Approximate Max Side"], {"default": "Exact Ratio"}),
                "custom_ratio_width": ("INT", {"default": 16, "min": 1, "max": 999, "step": 1}),
                "custom_ratio_height": ("INT", {"default": 9, "min": 1, "max": 999, "step": 1}),
                "multiple": ("INT", {"default": 32, "min": 1, "max": 512, "step": 1}),
                "size_step": ("INT", {"default": 10, "min": 1, "max": 50, "step": 1}),
            }
        }

    def get_resolutions(
        self,
        aspect_ratio,
        ratio_mode,
        custom_ratio_width,
        custom_ratio_height,
        multiple,
        size_step,
    ):
        if aspect_ratio == "Custom":
            ratio_width, ratio_height = normalize_ratio(custom_ratio_width, custom_ratio_height)
        else:
            ratio_width, ratio_height = ASPECT_RATIO_PRESETS[aspect_ratio]

        calculator = calculate_exact_resolution if ratio_mode == "Exact Ratio" else calculate_approx_resolution
        (
            base_width,
            base_height,
            upscale_2x_width,
            upscale_2x_height,
            upscale_4x_width,
            upscale_4x_height,
        ) = calculator(ratio_width, ratio_height, multiple, size_step)

        summary = (
            f"mode {ratio_mode} | "
            f"ratio {ratio_width}:{ratio_height} | "
            f"base {base_width}x{base_height} | "
            f"2x {upscale_2x_width}x{upscale_2x_height} | "
            f"4x {upscale_4x_width}x{upscale_4x_height}"
        )

        return (base_width, base_height, summary)


NODE_CLASS_MAPPINGS = {
    "LTXResolutionHelperNode": LTXResolutionHelperNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LTXResolutionHelperNode": "LTX Resolution Helper",
}

WEB_DIRECTORY = "./js"