from .nodes import (
    EraserNode,
    GenFillNode,
    ImageExpansionNode,
    ReplaceBgNode,
    RmbgNode,
    RemoveForegroundNode,
    ShotByTextOriginalNode,
    ShotByImageOriginalNode,
    TailoredGenNode,
    TailoredModelInfoNode,
    Text2ImageBaseNode,
    Text2ImageFastNode,
    Text2ImageHDNode,
    TailoredPortraitNode,
    ReimagineNode,
    GenerateImageNodeV2,
    RefineImageNodeV2,
    ShotByTextAutomaticNode,
    ShotByImageManualPaddingNode,
    ShotByImageAutomaticAspectRatioNode,
    ShotByImageCustomCoordinatesNode,
    ShotByImageManualPlacementNode,
    ShotByImageAutomaticNode,
    ShotByTextAutomaticAspectRatioNode,
    ShotByTextManualPlacementNode,
    ShotByTextManualPaddingNode,
    ShotByTextCustomCoordinatesNode,
    AttributionByImageNode
)

# Map the node class to a name used internally by ComfyUI
NODE_CLASS_MAPPINGS = {
    "BriaEraser": EraserNode,  # Return the class, not an instance
    "BriaGenFill": GenFillNode,
    "ImageExpansionNode": ImageExpansionNode,
    "ReplaceBgNode": ReplaceBgNode,
    "RmbgNode": RmbgNode,
    "RemoveForegroundNode": RemoveForegroundNode,
    "ShotByTextOriginal": ShotByTextOriginalNode,
    "ShotByImageOriginal": ShotByImageOriginalNode,
    "ShotByTextAutomatic": ShotByTextAutomaticNode,
    "ShotByTextManualPlacement": ShotByTextManualPlacementNode,
    "ShotByTextCustomCoordinates": ShotByTextCustomCoordinatesNode,
    "ShotByTextManualPadding": ShotByTextManualPaddingNode,
    "ShotByTextAutomaticAspectRatio": ShotByTextAutomaticAspectRatioNode,
    "ShotByImageAutomatic": ShotByImageAutomaticNode,
    "ShotByImageManualPlacement": ShotByImageManualPlacementNode,
    "ShotByImageCustomCoordinates": ShotByImageCustomCoordinatesNode,
    "ShotByImageManualPadding": ShotByImageManualPaddingNode,
    "ShotByImageAutomaticAspectRatio": ShotByImageAutomaticAspectRatioNode,
    "BriaTailoredGen": TailoredGenNode,
    "TailoredModelInfoNode": TailoredModelInfoNode,
    "TailoredPortraitNode": TailoredPortraitNode,
    "Text2ImageBaseNode": Text2ImageBaseNode,
    "Text2ImageFastNode": Text2ImageFastNode,
    "Text2ImageHDNode": Text2ImageHDNode,
    "ReimagineNode": ReimagineNode,
    "AttributionByImageNode": AttributionByImageNode,
    "GenerateImageNodeV2": GenerateImageNodeV2,
    "RefineImageNodeV2": RefineImageNodeV2,
}
# Map the node display name to the one shown in the ComfyUI node interface
NODE_DISPLAY_NAME_MAPPINGS = {
    "BriaEraser": "Bria Eraser",
    "BriaGenFill": "Bria GenFill",
    "ImageExpansionNode": "Bria Image Expansion",
    "ReplaceBgNode": "Bria Replace Background",
    "RmbgNode": "Bria RMBG",
    "RemoveForegroundNode": "Bria Remove Foreground",
    "ShotByTextOriginal": "Shot by Text - Original",
    "ShotByImageOriginal": "Shot by Image - Original",
    "ShotByTextAutomatic": "Shot by Text - Automatic",
    "ShotByTextManualPlacement": "Shot by Text - Manual Placement",
    "ShotByTextCustomCoordinates": "Shot by Text - Custom Coordinates",
    "ShotByTextManualPadding": "Shot by Text - Manual Padding",
    "ShotByTextAutomaticAspectRatio": "Shot by Text - Automatic Aspect Ratio",
    "ShotByImageAutomatic": "Shot by Image - Automatic",
    "ShotByImageManualPlacement": "Shot by Image - Manual Placement",
    "ShotByImageCustomCoordinates": "Shot by Image - Custom Coordinates",
    "ShotByImageManualPadding": "Shot by Image - Manual Padding",
    "ShotByImageAutomaticAspectRatio": "Shot by Image - Automatic Aspect Ratio",
    "BriaTailoredGen": "Bria Tailored Gen",
    "TailoredModelInfoNode": "Bria Tailored Model Info",
    "TailoredPortraitNode": "Bria Restyle Portrait",
    "Text2ImageBaseNode": "Bria Text2Image Base",
    "Text2ImageFastNode": "Bria Text2Image Fast",
    "Text2ImageHDNode": "Bria Text2Image HD",
    "ReimagineNode": "Bria Reimagine",
    "AttributionByImageNode": "Attribution By Image Node",
    "GenerateImageNodeV2": "Generate Image",
    "RefineImageNodeV2": "Refine and Regenerate Image",
}