# ComfyUI-QuickRatio — aspect-ratio calculator node for ComfyUI.
# Copyright (C) 2025 Dragon7108
#
# This program is free software: you can redistribute it and/or modify it under the terms
# of the GNU General Public License as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later version. It is distributed WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
# PARTICULAR PURPOSE. See the GNU General Public License for more details, a copy of which
# ships in LICENSE and is available at <https://www.gnu.org/licenses/>.

class QuickRatioNode:
    CATEGORY = "utilities"
    DESCRIPTION = "Quick aspect ratio calculator with live dimension preview"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "preset": (["1:1", "11:17", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21", "3:2", "2:3", "5:4", "4:5", "custom"], {"default": "1:1"}),
                "base_resolution": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}),
                "divisible_by": ("INT", {"default": 8, "min": 1, "max": 256, "step": 8}),
            },
            "optional": {
                "custom_width_ratio": ("INT", {"default": 1, "min": 1, "max": 100}),
                "custom_height_ratio": ("INT", {"default": 1, "min": 1, "max": 100}),
            },
        }
    
    RETURN_TYPES = ("INT", "INT")
    RETURN_NAMES = ("width", "height")
    FUNCTION = "get_dimensions"
    OUTPUT_NODE = False
    
    def get_dimensions(self, preset, base_resolution, divisible_by=8, custom_width_ratio=1, custom_height_ratio=1):
        print(f"QuickRatio DEBUG: preset={preset}, base_resolution={base_resolution}, divisible_by={divisible_by}, custom_width_ratio={custom_width_ratio}, custom_height_ratio={custom_height_ratio}")
        
        # Define aspect ratio presets.
        # KEEP IN SYNC with: the `preset` dropdown above, and the presets map in
        # web/js/quickratio_live.js (the live-badge preview duplicates this).
        presets = {
            "1:1": (1, 1),
            "11:17": (11, 17),
            "16:9": (16, 9),
            "9:16": (9, 16),
            "4:3": (4, 3),
            "3:4": (3, 4),
            "21:9": (21, 9),
            "9:21": (9, 21),
            "3:2": (3, 2),
            "2:3": (2, 3),
            "5:4": (5, 4),
            "4:5": (4, 5)
        }
        
        # Get ratio values
        if preset == "custom":
            ratio_w, ratio_h = custom_width_ratio, custom_height_ratio
        else:
            ratio_w, ratio_h = presets.get(preset, (1, 1))
        
        # Calculate dimensions, snapping the computed side to `divisible_by`
        ratio = ratio_w / ratio_h
        d = max(1, divisible_by)
        
        if ratio >= 1:
            # Landscape or square: base resolution is width
            width = base_resolution
            height = round(base_resolution / ratio / d) * d
        else:
            # Portrait: base resolution is height
            height = base_resolution
            width = round(base_resolution * ratio / d) * d
        
        print(f"QuickRatio DEBUG: calculated width={width}, height={height}")
        return (width, height)

NODE_CLASS_MAPPINGS = {
    "QuickRatio": QuickRatioNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "QuickRatio": "Quick Ratio"
}