import numpy as np

try:
    from PIL import Image, ImageOps
except ImportError:  # pragma: no cover - PIL ships with ComfyUI
    Image = None
    ImageOps = None

try:
    from ._label_ops import compose_labeled_image
except ImportError:  # importable for tests without package context
    from _label_ops import compose_labeled_image

CATEGORY = "MoonPack/image"


class LabeledLoadImage:
    DESCRIPTION = (
        "Loads an image and bakes an editable caption bar into it (e.g. 'Front View', "
        "'Character A'), so a reference stays identifiable once it's combined into a "
        "bigger Reference Concat/Grid layout. The bar is added as extra space at the "
        "top or bottom - the photo itself is never cropped or resized. A bold arrow in "
        "the bar points back at the photo; text and arrow can each be turned off "
        "independently. An optional border (same color as the bar) frames the whole "
        "result. Use the node's Preview button for an instant, exact look at the result "
        "before queuing."
    )
    SEARCH_ALIASES = ["load image", "label", "caption", "annotate", "reference", "text overlay"]

    @classmethod
    def INPUT_TYPES(cls):
        import os

        import folder_paths

        input_dir = folder_paths.get_input_directory()
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
        files = folder_paths.filter_files_content_types(files, ["image"])
        return {
            "required": {
                "image": (sorted(files), {"image_upload": True}),
                "text": ("STRING", {
                    "default": "", "multiline": False,
                    "tooltip": "Caption drawn in the label bar. Auto-shrinks to fit, then truncates with an ellipsis.",
                }),
                "text_position": (["top", "bottom"], {
                    "default": "top",
                    "tooltip": "Which edge of the image the label bar is added to.",
                }),
                "show_label": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Draw the caption text in the bar.",
                }),
                "show_arrow": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Draw a bold arrow in the bar, pointing back at the photo.",
                }),
                "background_color": ("STRING", {
                    "default": "#000000",
                    "tooltip": "Label bar fill color and border color, as a hex code.",
                }),
                "text_color": ("STRING", {
                    "default": "#ffffff",
                    "tooltip": "Caption text and arrow color, as a hex code.",
                }),
                "border_width": ("INT", {
                    "default": 0, "min": 0, "max": 512, "step": 1,
                    "tooltip": "Inset border thickness in pixels, drawn in background_color around the whole result. 0 = no border.",
                }),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "load_image"
    CATEGORY = CATEGORY

    def load_image(self, image, text, text_position, show_label, show_arrow,
                    background_color, text_color, border_width):
        import folder_paths

        image_path = folder_paths.get_annotated_filepath(image)
        img = Image.open(image_path)
        img = ImageOps.exif_transpose(img).convert("RGB")

        composed = compose_labeled_image(
            img, text, text_position, show_label, show_arrow,
            background_color, text_color, border_width,
        )

        arr = np.array(composed).astype(np.float32) / 255.0
        tensor = arr[None, ...]

        import torch
        return (torch.from_numpy(tensor),)


NODE_CLASS_MAPPINGS = {
    "MoonPack_LabeledLoadImage": LabeledLoadImage,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MoonPack_LabeledLoadImage": "Load Image (Labeled)",
}