import gradio as gr
from PIL import Image, ImageOps
from collections import Counter
import time
from functools import wraps

MAX_SIZE_PREVIEW = 320


def is_not_dark(color, threshold=30):
    # color is a tuple like (R, G, B)
    # Avoid dark pixels where all channels are <= threshold
    return not all(c <= threshold for c in color)


def get_dominant_color_exclude_dark(image):
    img_small = image.resize((50, 50))  # smaller for speed
    pixels = list(img_small.getdata())

    # Filter out dark or near-black pixels
    filtered_pixels = [p for p in pixels if is_not_dark(p)]

    if not filtered_pixels:
        # If all pixels are dark, fallback to original list
        filtered_pixels = pixels

    most_common = Counter(filtered_pixels).most_common(1)[0][0]
    return most_common


def parse_color(color):
    """
    Convert color to Pillow-friendly (R, G, B, A) tuple in 0–255 range.
    Supports:
      - tuple/list of floats or ints
      - 'rgba(r, g, b, a)' string
      - 'rgb(r, g, b)' string
      - hex colors: '#RRGGBB' or '#RRGGBBAA'
    """

    if isinstance(color, (tuple, list)):
        parts = [float(c) for c in color]

    elif isinstance(color, str):
        c = color.strip().lower()

        # Hex color
        if c.startswith("#"):
            c = c.lstrip("#")
            if len(c) == 6:  # RRGGBB
                r, g, b = int(c[0:2], 16), int(c[2:4], 16), int(c[4:6], 16)
                return (r, g, b, 255)
            elif len(c) == 8:  # RRGGBBAA
                r, g, b, a = int(c[0:2], 16), int(c[2:4], 16), int(c[4:6], 16), int(c[6:8], 16)
                return (r, g, b, a)
            else:
                raise ValueError(f"Invalid hex color: {color}")

        # RGBA / RGB string
        c = c.replace("rgba", "").replace("rgb", "").replace("(", "").replace(")", "")
        parts = [float(x.strip()) for x in c.split(",")]

    else:
        raise ValueError(f"Unsupported color format: {color}")

    # Handle missing alpha in tuple/list or rgb() string
    if len(parts) == 3:
        parts.append(1.0)  # default alpha=1

    # Convert to int (0–255)
    return (
        int(round(parts[0])),
        int(round(parts[1])),
        int(round(parts[2])),
        int(round(parts[3] * 255 if parts[3] <= 1 else parts[3]))
    )


def to_rgb_with_bg(img, bg_color=(255, 255, 255)):
    """Convert image to RGB, replacing transparent areas with a given background color."""
    if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info):
        background = Image.new("RGB", img.size, bg_color)
        background.paste(img, mask=img.split()[-1])  # use alpha channel as mask
        return background
    else:
        return img.convert("RGB")


def proportional_reviewer(img, max_size=MAX_SIZE_PREVIEW, resampling_filter=Image.NEAREST):
    # First scale
    new_width = int(img.width)
    new_height = int(img.height)

    # Enforce max size
    if new_width > max_size or new_height > max_size:
        scale_factor = max_size / max(new_width, new_height)
        new_width = int(new_width * scale_factor)
        new_height = int(new_height * scale_factor)

    return img.resize((new_width, new_height), resampling_filter)


def get_dominant_color(image):
    img_small = image.resize((50, 50))  # smaller for speed
    pixels = list(img_small.getdata())
    most_common = Counter(pixels).most_common(1)[0][0]
    return most_common


def upload_image(img):
    # Store image and reset padding
    state = {
        "image": img,
        "pad_top": 0,
        "pad_bottom": 0,
        "pad_left": 0,
        "pad_right": 0,
    }
    if img is None:
        return None, state

    dominant_color = get_dominant_color_exclude_dark(img.convert("RGB"))
    state["dominant_color"] = dominant_color
    state["image_rgb"] = to_rgb_with_bg(img, bg_color=dominant_color)

    return proportional_reviewer(state["image_rgb"]), state


def add_padding(direction, state, pad_value, mode="Add padding"):
    if state is None or state["image"] is None:
        return None, state

    if mode == "Remove padding":
        pad_value = -pad_value

    state[f"pad_{direction}"] += pad_value
    state[f"pad_{direction}"] = max(state[f"pad_{direction}"], 0)

    return preview_image(state), state


def timing(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f"Elapsed time: {end - start:.4f} seconds")
        return result
    return wrapper


# @timing
def preview_image(state):
    img = state["image"]
    if img is None:
        return None

    padded = ImageOps.expand(
        state["image_rgb"],
        border=(state["pad_left"], state["pad_top"], state["pad_right"], state["pad_bottom"]),
        fill=state["dominant_color"]
    )

    return proportional_reviewer(padded)


@timing
def generate_result(state, img_source, overlap=.1, color=None, max_resolution=0, expand_top=False, expand_left=False, expand_right=False, expand_bottom=False):
    img = state["image"]
    if img is None:
        return None, None

    # Handle transparency replacement inside original image
    has_alpha = img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info)
    width, height = img.size

    # Add padding with dominant color
    if color in [None, ""]:
        img_rgb = state["image_rgb"]
    else:
        img_rgb = to_rgb_with_bg(img, bg_color=parse_color(color))

    padded = ImageOps.expand(
        img_rgb,
        border=(state["pad_left"], state["pad_top"], state["pad_right"], state["pad_bottom"]),
        fill=(state["dominant_color"] if color in [None, ""] else parse_color(color))
    )

    # Calculate original black box area
    orig_box = (
        state["pad_left"],
        state["pad_top"],
        state["pad_left"] + width,
        state["pad_top"] + height
    )

    # Apply overlap: shrink black box by overlap percentage
    if overlap > 0:
        overlap_x = int(width * overlap)
        overlap_y = int(height * overlap)

        left_shift = overlap_x if expand_left else 0
        top_shift = overlap_y if expand_top else 0
        right_shift = overlap_x if expand_right else 0
        bottom_shift = overlap_y if expand_bottom else 0

        orig_box = (
            orig_box[0] + left_shift,
            orig_box[1] + top_shift,
            orig_box[2] - right_shift,
            orig_box[3] - bottom_shift
        )

    # Start mask all white, then paste black for (shrunken) image area
    mask = Image.new("L", padded.size, 255)
    mask.paste(0, orig_box)

    # If original had transparency, fix transparent pixels to white
    has_alpha = state["image"].mode in ("RGBA", "LA") or (
        state["image"].mode == "P" and "transparency" in state["image"].info
    )
    if has_alpha:
        alpha = state["image"].convert("RGBA").split()[-1]
        trans_mask = Image.new("L", alpha.size, 0)
        trans_mask.paste(255, mask=alpha.point(lambda a: 0 if a > 0 else 255))
        mask.paste(255, (state["pad_left"], state["pad_top"]), trans_mask)

    if max_resolution:
        padded = proportional_reviewer(padded, max_resolution, Image.LANCZOS)
        mask = proportional_reviewer(mask, max_resolution, Image.LANCZOS)

    return padded, mask


CSS = """
#img_input {
  height: 200px;
  width: 200px;
  max-width: 200px;
  max-height: 200px;
  margin-left: 0;      /* no left margin */
  margin-right: auto;  /* push any extra space to the right */
}

#img_output img {
    object-fit: contain !important; /* keep whole image visible */
    width: 100% !important;
    height: 100% !important;
}

#container {
    display: flex;
    align-items: center;
}

#btn_left, #btn_right {
    height: 400px; /* match img_output height */
    transform-origin: center;
    display: inline-flex;
    align-items: center;
    justify-content: center;
}

#btn_top, #btn_botton {
    width: 600px;
    max-width: 600px;
    margin: 0 auto;  /* center block horizontally */
}
"""

with gr.Blocks(css=CSS) as demo:
    gr.Markdown("## Outpaint Mask Maker")
    gr.Markdown("More spaces: [inpaint-mask-maker](https://huggingface.co/spaces/r3gm/inpaint-mask-maker)")

    state = gr.State({"image": None, "pad_top": 0, "pad_bottom": 0, "pad_left": 0, "pad_right": 0})

    with gr.Row():
        img_input = gr.Image(type="pil", label="Upload Image", elem_id="img_input", sources=['upload', 'clipboard'], image_mode="RGBA")

    with gr.Column():
        with gr.Row():
            btn_top = gr.Button("⬆️ Top Padding", elem_id="btn_top", variant="huggingface")
        with gr.Row():
            btn_left = gr.Button("⬅️ Left Padding", elem_id="btn_left", scale=1, variant="huggingface")
            img_output = gr.Image(label="Preview", height=400, width=400, scale=10, elem_id="img_output", show_download_button=False)
            btn_right = gr.Button("➡️ Right Padding", elem_id="btn_right", scale=1, variant="huggingface")
        with gr.Row():
            btn_bottom = gr.Button("⬇️ Bottom Padding", elem_id="btn_botton", variant="huggingface")

    padding_pixels = gr.Slider(1, 1000, value=100, step=1, label="Padding Amount (pixels)", info="How many pixels of padding to add each time you tap the `Padding Buttons`")
    color_ = gr.ColorPicker(label="Color Used to Fill Padding", info="Choose the color used to fill the extended padding area. If not defined, the most common color in the image will be used.")

    with gr.Accordion("Settings", open=False):
        padding_mode = gr.Dropdown(["Add padding", "Remove padding"], label="Padding Mode", info="Add or remove space around the image.")
        max_resolution_ = gr.Slider(
            0,
            8192,
            value=0,
            step=64,
            label="Max Resolution (px)",
            info="Set the exact maximum size in pixels while keeping the aspect ratio. Set to 0 to disable.",
        )
        with gr.Row():
            mask_expansion = gr.Slider(.0, .5, value=0.1, step=0.01, label="Mask Expansion (%)", info="Change how far the mask reaches from the edges of the image.")
        with gr.Row():
            expand_top = gr.Checkbox(label="Expand Top", value=True)
            expand_left = gr.Checkbox(label="Expand Left", value=True)
            expand_right = gr.Checkbox(label="Expand Right", value=True)
            expand_bottom = gr.Checkbox(label="Expand Bottom", value=True)

    btn_result = gr.Button("Generate Result", variant="primary")

    with gr.Row():
        result_img = gr.Image(label="Final Padded Image", format="png", image_mode="RGB")
        result_mask = gr.Image(label="Mask", format="png", image_mode="RGB")

    img_input.change(upload_image, inputs=img_input, outputs=[img_output, state])
    btn_top.click(lambda s, p, m: add_padding("top", s, p, m), inputs=[state, padding_pixels, padding_mode], outputs=[img_output, state], show_progress="minimal")
    btn_bottom.click(lambda s, p, m: add_padding("bottom", s, p, m), inputs=[state, padding_pixels, padding_mode], outputs=[img_output, state], show_progress="minimal")
    btn_left.click(lambda s, p, m: add_padding("left", s, p, m), inputs=[state, padding_pixels, padding_mode], outputs=[img_output, state], show_progress="minimal")
    btn_right.click(lambda s, p, m: add_padding("right", s, p, m), inputs=[state, padding_pixels, padding_mode], outputs=[img_output, state], show_progress="minimal")

    btn_result.click(generate_result, inputs=[state, img_input, mask_expansion, color_, max_resolution_, expand_top, expand_left, expand_right, expand_bottom], outputs=[result_img, result_mask])

demo.launch(share=False, debug=True, show_error=True)