import os
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont

class TextOverlay:
    """
    Text overlay node with:
      - Fill/stroke alpha
      - Shadow and background box
      - Pixel-perfect stroke alignment for MULTILINE text
      - Even stroke width for crisper edges
      - Default vertical_alignment = 'middle'
    """

    _horizontal_alignments = ["left", "center", "right"]
    _vertical_alignments = ["top", "middle", "bottom"]

    def __init__(self, device="cpu"):
        self.device = device
        self._loaded_font = None
        # cache (lines, widths, heights, tops, min_top, min_left, block_w, block_h)
        self._cached = None
        self._x = None
        self._y = None

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # keep image first so the node wiring stays obvious
                "image": ("IMAGE",),

                # UI order (top → bottom)
                "text": ("STRING", {"multiline": True, "default": "the quick brown fox\njumps over the lazy dog"}),
                "all_caps": ("BOOLEAN", {"default": False}),

                # font, font-size, font color, font alpha
                "font": ("STRING", {"default": "ariblk.ttf"}),
                "font_size": ("INT", {"default": 32, "min": 1, "max": 9999, "step": 1}),
                "fill_color_hex": ("STRING", {"default": "#FFFFFF"}),
                "fill_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "display": "slider"}),

                # padding, hor_align, x_shift, vert align, y_shift, line_spacing
                "padding": ("INT", {"default": 16, "min": 0, "max": 256, "step": 1}),
                "vertical_alignment": (cls._vertical_alignments, {"default": "middle"}),
                "y_shift": ("INT", {"default": 0, "min": -256, "max": 256, "step": 1}),
                "horizontal_alignment": (cls._horizontal_alignments, {"default": "center"}),
                "x_shift": ("INT", {"default": 0, "min": -256, "max": 256, "step": 1}),
                "line_spacing": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 50.0, "step": 0.5}),

                # all strokes
                "stroke_enable": ("BOOLEAN", {"default": True}),
                "stroke_color_hex": ("STRING", {"default": "#000000"}),
                "stroke_thickness": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01}),
                "stroke_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "display": "slider"}),

                # all boxes
                "bg_enable": ("BOOLEAN", {"default": False}),
                "bg_padding": ("INT", {"default": 8, "min": 0, "max": 256, "step": 1}),
                "bg_radius": ("INT", {"default": 8, "min": 0, "max": 512, "step": 1}),
                "bg_color_hex": ("STRING", {"default": "#000000"}),
                "bg_alpha": ("FLOAT", {"default": 0.50, "min": 0.0, "max": 1.0, "step": 0.01, "display": "slider"}),

                # all shadows
                "shadow_enable": ("BOOLEAN", {"default": False}),
                "shadow_distance": ("INT", {"default": 4, "min": -50, "max": 50, "step": 1}),
                "shadow_color_hex": ("STRING", {"default": "#000000"}),
                "shadow_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "display": "slider"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "batch_process"
    CATEGORY = "image/text"

    # ---------------- helpers ----------------

    def hex_to_rgb(self, hex_color: str):
        hex_color = hex_color.strip().lstrip("#")
        if len(hex_color) == 3:
            hex_color = "".join(ch * 2 for ch in hex_color)
        return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))

    def _normalize_text(self, text: str) -> str:
        # Convert literal "\n" to actual newlines
        return text.replace("\\n", "\n").replace("\\N", "\n")

    def _load_font(self, font, font_size):
        fonts_dir = os.path.join(os.path.dirname(__file__), "fonts")
        font_path = os.path.join(fonts_dir, font)
        if not os.path.exists(font_path):
            font_path = font  # fallback to system path or provided absolute path
        try:
            return ImageFont.truetype(font_path, font_size)
        except Exception as e:
            print(f"Error loading font: {e} — using default font")
            return ImageFont.load_default()

    def _wrap_lines(self, draw, text, font, max_width, padding):
        """
        Greedy wrapping that respects explicit newlines and preserves spaces
        (including multiple consecutive spaces).
        """
        import re

        paragraphs = self._normalize_text(text).split("\n")
        out_lines = []

        for para in paragraphs:
            # Keep empty lines exactly
            if para == "":
                out_lines.append("")
                continue

            # Tokens are either non-space runs or whitespace runs
            tokens = re.findall(r"\S+|\s+", para)

            line = ""
            for tok in tokens:
                candidate = line + tok
                if draw.textlength(candidate, font=font) <= (max_width - 2 * padding):
                    line = candidate
                else:
                    if line == "":
                        # Single token longer than the max width: force-break
                        out_lines.append(candidate)
                        line = ""
                    else:
                        out_lines.append(line)
                        # Start new line with current token (trim leading spaces only)
                        line = tok.lstrip()

            # Push whatever remains (don’t strip—preserve trailing spaces)
            out_lines.append(line)

        return out_lines

    # ---------------- core drawing ----------------

    def _compute_layout(self, img_w, img_h, draw, text, font, stroke_width, padding,
                        h_align, v_align, x_shift, y_shift, line_spacing, use_cache):
        """
        Manual multiline layout so stroke/fill line positions are identical.

        IMPORTANT: We compute the visual extents (min_left / max_right) from each line's
        textbbox (which already includes stroke when stroke_width > 0). This gives a
        correct block width for the background box and eliminates the right-side gap
        when stroke is enabled.

        Returns:
          (lines, widths, heights, tops, min_top, min_left, block_w, block_h, x0, visual_top_y)
        """
        if self._cached is None or not use_cache:
            lines = self._wrap_lines(draw, text, font, img_w, padding)

            widths, heights, tops = [], [], []
            lefts, rights = [], []
            for ln in lines:
                # bbox includes stroke when stroke_width > 0
                l, t, r, b = draw.textbbox((0, 0), ln, font=font, stroke_width=stroke_width)
                widths.append(r - l)
                heights.append(b - t)
                tops.append(t)
                lefts.append(l)
                rights.append(r)

            min_left = min(lefts) if lefts else 0
            max_right = max(rights) if rights else 0
            block_w = max_right - min_left if rights else 0
            block_h = sum(heights) + (len(heights) - 1) * line_spacing if heights else 0
            min_top = min(tops) if tops else 0  # may be negative for ascenders

            # cache results to reuse on subsequent frames
            self._cached = (lines, widths, heights, tops, min_top, min_left, block_w, block_h)

        lines, widths, heights, tops, min_top, min_left, block_w, block_h = self._cached

        # Horizontal anchor for the *visual* block (uses min_left/max_right)
        if h_align == "left":
            x0 = padding
        elif h_align == "center":
            x0 = (img_w - block_w) / 2
        else:  # right
            x0 = img_w - block_w - padding

        # Vertical anchor for the *visual* block
        if v_align == "top":
            visual_top_y = padding
        elif v_align == "middle":
            visual_top_y = (img_h - block_h) / 2
        else:  # bottom
            visual_top_y = img_h - block_h - padding

        # Snap + shifts
        x0 = int(round(x0 + x_shift))
        visual_top_y = int(round(visual_top_y + y_shift))

        return lines, widths, heights, tops, min_top, min_left, block_w, block_h, x0, visual_top_y

    def draw_text(
        self,
        image,
        text,
        all_caps,
        font_size,
        font,
        fill_color_hex,
        fill_alpha,
        stroke_enable,
        stroke_color_hex,
        stroke_alpha,
        stroke_thickness,
        padding,
        horizontal_alignment,
        vertical_alignment,
        x_shift,
        y_shift,
        line_spacing,
        bg_enable,
        bg_color_hex,
        bg_alpha,
        bg_padding,
        bg_radius,
        shadow_enable,
        shadow_color_hex,
        shadow_alpha,
        shadow_distance,
        use_cache=False,
    ):
        # Prepare
        if image.mode != "RGBA":
            image = image.convert("RGBA")
        self._loaded_font = self._load_font(font, font_size)
        draw = ImageDraw.Draw(image, "RGBA")

        # Apply ALL CAPS if requested (after literal "\n" handling via _wrap_lines)
        if all_caps:
            # Uppercase before wrapping so measurement/layout reflect final glyphs
            text = text.upper()

        # Compute even stroke width (recommended for crisp outlines)
        sw = int(round(font_size * stroke_thickness * 0.5)) if stroke_enable else 0
        if sw % 2 == 1 and sw > 0:
            sw += 1

        # Layout (manual lines)
        (lines, widths, heights, tops, min_top, min_left, block_w, block_h,
         x0, visual_top_y) = self._compute_layout(
            image.width, image.height, draw, text, self._loaded_font, sw,
            padding, horizontal_alignment, vertical_alignment, x_shift, y_shift,
            line_spacing, use_cache
        )

        # Convert the visual top into the baseline y for the first line
        first_line_baseline_y = visual_top_y - min_top

        # This is the x where we actually draw glyphs so that the visual left aligns to x0
        x_draw = x0 - min_left

        # Background box (perfectly aligned with visual extents)
        if bg_enable and block_w > 0 and block_h > 0:
            br, bgc, bb = self.hex_to_rgb(bg_color_hex)
            ba = int(max(0.0, min(1.0, bg_alpha)) * 255)
            rect = [x0 - bg_padding, visual_top_y - bg_padding,
                    x0 + block_w + bg_padding, visual_top_y + block_h + bg_padding]
            overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
            od = ImageDraw.Draw(overlay, "RGBA")
            try:
                od.rounded_rectangle(rect, radius=max(0, int(bg_radius)), fill=(br, bgc, bb, ba))
            except Exception:
                od.rectangle(rect, fill=(br, bgc, bb, ba))
            image = Image.alpha_composite(image, overlay)

        # Shadow (per-line, before stroke/fill)
        if shadow_enable and block_w > 0 and block_h > 0:
            sh_r, sh_g, sh_b = self.hex_to_rgb(shadow_color_hex)
            sh_a = int(max(0.0, min(1.0, shadow_alpha)) * 255)
            sdx = sdy = int(shadow_distance)
            overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
            od = ImageDraw.Draw(overlay, "RGBA")

            yy = first_line_baseline_y
            for ln, h in zip(lines, heights):
                od.text((x_draw + sdx, yy + sdy), ln, font=self._loaded_font, fill=(sh_r, sh_g, sh_b, sh_a))
                yy += int(round(h + line_spacing))

            image = Image.alpha_composite(image, overlay)

        # Stroke + Fill (per-line; identical geometry/positions)
        fr, fg, fb = self.hex_to_rgb(fill_color_hex)
        fa = int(max(0.0, min(1.0, fill_alpha)) * 255)
        sr, sg, sb = self.hex_to_rgb(stroke_color_hex)
        sa = int(max(0.0, min(1.0, stroke_alpha)) * 255)

        overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
        od = ImageDraw.Draw(overlay, "RGBA")

        yy = first_line_baseline_y
        for ln, h in zip(lines, heights):
            if sw > 0 and sa > 0:
                od.text(
                    (x_draw, yy),
                    ln,
                    font=self._loaded_font,
                    fill=(sr, sg, sb, sa),
                    stroke_width=sw,
                    stroke_fill=(sr, sg, sb, sa),
                )
            if fa > 0:
                od.text(
                    (x_draw, yy),
                    ln,
                    font=self._loaded_font,
                    fill=(fr, fg, fb, fa),
                )
            yy += int(round(h + line_spacing))

        image = Image.alpha_composite(image, overlay)
        return image.convert("RGB")

    # ---------------- Comfy entrypoint ----------------

    def batch_process(
        self,
        image,
        text,
        all_caps,  # NEW
        font_size,
        font,
        fill_color_hex,
        fill_alpha,
        stroke_enable,
        stroke_color_hex,
        stroke_alpha,
        stroke_thickness,
        padding,
        horizontal_alignment,
        vertical_alignment,
        x_shift,
        y_shift,
        line_spacing,
        bg_enable,
        bg_color_hex,
        bg_alpha,
        bg_padding,
        bg_radius,
        shadow_enable,
        shadow_color_hex,
        shadow_alpha,
        shadow_distance,
    ):
        """Handles both single and batch image processing (with layout caching across frames)."""

        # Single image
        if len(image.shape) == 3:
            np_img = image.cpu().numpy()
            pil_img = Image.fromarray((np_img * 255).astype(np.uint8))
            out_img = self.draw_text(
                pil_img, text, all_caps,
                font_size, font,
                fill_color_hex, fill_alpha,
                stroke_enable,
                stroke_color_hex, stroke_alpha, stroke_thickness,
                padding, horizontal_alignment, vertical_alignment,
                x_shift, y_shift, line_spacing,
                bg_enable, bg_color_hex, bg_alpha, bg_padding, bg_radius,
                shadow_enable, shadow_color_hex, shadow_alpha, shadow_distance,
                use_cache=False,
            )
            out = np.array(out_img).astype(np.float32) / 255.0
            return (torch.tensor(out),)

        # Batch of images
        else:
            np_imgs = image.cpu().numpy()
            outs = []
            use_cache = False
            for arr in np_imgs:
                pil_img = Image.fromarray((arr * 255).astype(np.uint8))
                out_img = self.draw_text(
                    pil_img, text, all_caps,
                    font_size, font,
                    fill_color_hex, fill_alpha,
                    stroke_enable,
                    stroke_color_hex, stroke_alpha, stroke_thickness,
                    padding, horizontal_alignment, vertical_alignment,
                    x_shift, y_shift, line_spacing,
                    bg_enable, bg_color_hex, bg_alpha, bg_padding, bg_radius,
                    shadow_enable, shadow_color_hex, shadow_alpha, shadow_distance,
                    use_cache=use_cache,  # reuse wrapped lines/metrics for next frames
                )
                outs.append(np.array(out_img).astype(np.float32) / 255.0)
                use_cache = True
            return (torch.tensor(np.stack(outs)),)

NODE_CLASS_MAPPINGS = {"Advanced Text Overlay": TextOverlay}