from .nodes import (
    EraserNode,
    GenFillNode,
    ImageExpansionNode,
    ReplaceBgNode,
    RmbgNode,
    RemoveForegroundNode,
    ShotByTextOriginalNode,
    ShotByImageOriginalNode,
    TailoredGenNode,
    TailoredModelInfoNode,
    Text2ImageBaseNode,
    Text2ImageFastNode,
    Text2ImageHDNode,
    TailoredPortraitNode,
    ReimagineNode,
    GenerateImageNodeV2,
    GenerateImageLiteNodeV2,
    RefineImageNodeV2,
    RefineImageLiteNodeV2,
    GenerateStructuredPromptNodeV2,
    GenerateStructuredPromptLiteNodeV2,
    ShotByTextAutomaticNode,
    ShotByImageManualPaddingNode,
    ShotByImageAutomaticAspectRatioNode,
    ShotByImageCustomCoordinatesNode,
    ShotByImageManualPlacementNode,
    ShotByImageAutomaticNode,
    ShotByTextAutomaticAspectRatioNode,
    ShotByTextManualPlacementNode,
    ShotByTextManualPaddingNode,
    ShotByTextCustomCoordinatesNode,
    AttributionByImageNode,
    RemoveVideoBackgroundNode,
    VideoSolidColorBackgroundNode,
    VideoMaskByPromptNode,
    VideoMaskByKeyPointsNode,
    VideoIncreaseResolutionNode,
    VideoEraseElementsNode,
    LoadVideoFramesNode,
    PreviewVideoURLNode
)

# 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,
    "GenerateImageLiteNodeV2": GenerateImageLiteNodeV2,
    "RefineImageNodeV2": RefineImageNodeV2,
    "RefineImageLiteNodeV2": RefineImageLiteNodeV2,
    "GenerateStructuredPromptNodeV2": GenerateStructuredPromptNodeV2,
    "GenerateStructuredPromptLiteNodeV2": GenerateStructuredPromptLiteNodeV2,
    "RemoveVideoBackgroundNode":RemoveVideoBackgroundNode,
    "VideoSolidColorBackgroundNode":VideoSolidColorBackgroundNode,
    "VideoMaskByPromptNode":VideoMaskByPromptNode,
    "VideoMaskByKeyPointsNode":VideoMaskByKeyPointsNode,
    "VideoIncreaseResolutionNode":VideoIncreaseResolutionNode,
    "VideoEraseElementsNode":VideoEraseElementsNode,
    "LoadVideoFramesNode":LoadVideoFramesNode,
    "PreviewVideoURLNode":PreviewVideoURLNode
}
# 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",
    "GenerateImageLiteNodeV2": "Generate Image - Lite",
    "RefineImageNodeV2": "Refine and Regenerate Image",
    "RefineImageLiteNodeV2": "Refine Image - Lite",
    "GenerateStructuredPromptNodeV2": "Generate Structured Prompt",
    "GenerateStructuredPromptLiteNodeV2": "Generate Structured Prompt - Lite",
    "RemoveVideoBackgroundNode": "Bria Remove Video Background",
    "VideoSolidColorBackgroundNode":"Bria SolidColor Background Video",
    "VideoMaskByPromptNode":"Bria Video Mask By Prompt",
    "VideoMaskByKeyPointsNode":"Bria Video Mask By Key Points",
    "VideoIncreaseResolutionNode":"Bria Video Increase Resolution",
    "VideoEraseElementsNode":"Bria Video Erase Elements",
    "LoadVideoFramesNode":"Bria Load Video",
    "PreviewVideoURLNode":"Bria Preview Video"
}