from __future__ import annotations

import html
import os
from dataclasses import dataclass
from functools import lru_cache
from typing import Iterable

import cv2
import gradio as gr
import numpy as np
from PIL import Image, ImageFilter

try:
    import spaces
except ImportError:
    spaces = None


APP_TITLE = "PlanPalette"
APP_SUBTITLE = "Generate a furnished top-down architectural floor plan render from a reference palette and CAD plan."
BASE_MODEL_ID = os.getenv("PLANPALETTE_BASE_MODEL", "Lykon/dreamshaper-xl-lightning")
IS_HF_SPACE = bool(os.getenv("SPACE_ID"))


@dataclass
class PaletteColor:
    rgb: tuple[int, int, int]
    percent: float
    material: str


def pil_to_rgb_array(image: Image.Image) -> np.ndarray:
    return np.asarray(image.convert("RGB"), dtype=np.uint8)


def rgb_to_hex(rgb: Iterable[int]) -> str:
    r, g, b = [int(v) for v in rgb]
    return f"#{r:02X}{g:02X}{b:02X}"


def infer_material_name(rgb: tuple[int, int, int]) -> str:
    color = np.uint8([[list(rgb)]])
    hsv = cv2.cvtColor(color, cv2.COLOR_RGB2HSV)[0, 0]
    hue, sat, val = int(hsv[0]), int(hsv[1]), int(hsv[2])

    if val < 70:
        return "charcoal line / deep accent"
    if sat < 35 and val > 205:
        return "plaster / light stone"
    if sat < 45:
        return "concrete / neutral finish"
    if 18 <= hue <= 38:
        return "wood / warm flooring"
    if 39 <= hue <= 82:
        return "planting / landscape"
    if 83 <= hue <= 104:
        return "mint glass / cool surface"
    if 105 <= hue <= 135:
        return "water / blue finish"
    if 136 <= hue <= 165:
        return "soft fabric / feature zone"
    return "accent material"


def sample_reference_pixels(image: np.ndarray, max_samples: int = 26000) -> np.ndarray:
    pixels = image.reshape(-1, 3)
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY).reshape(-1)
    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV).reshape(-1, 3)

    # Keep meaningful color and neutral finish pixels, but avoid pure paper and
    # black linework so the palette reflects the reference styling.
    not_white = gray < 244
    not_black_line = gray > 35
    has_visual_weight = (hsv[:, 1] > 18) | (gray < 225)
    candidates = pixels[not_white & not_black_line & has_visual_weight]
    if len(candidates) < 64:
        candidates = pixels[(gray > 25) & (gray < 248)]
    if len(candidates) == 0:
        candidates = pixels

    if len(candidates) > max_samples:
        rng = np.random.default_rng(42)
        candidates = candidates[rng.choice(len(candidates), max_samples, replace=False)]
    return candidates.astype(np.float32)


def extract_palette(image: np.ndarray, k: int = 6) -> list[PaletteColor]:
    samples = sample_reference_pixels(image)
    k = int(max(2, min(k, len(samples), 8)))

    criteria = (
        cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
        35,
        0.8,
    )
    compactness, labels, centers = cv2.kmeans(
        samples,
        k,
        None,
        criteria,
        3,
        cv2.KMEANS_PP_CENTERS,
    )
    del compactness

    counts = np.bincount(labels.flatten(), minlength=k).astype(np.float32)
    order = np.argsort(counts)[::-1]
    palette: list[PaletteColor] = []
    total = float(counts.sum()) or 1.0

    for idx in order:
        rgb = tuple(np.clip(np.rint(centers[idx]), 0, 255).astype(int).tolist())
        palette.append(
            PaletteColor(
                rgb=rgb,
                percent=float(counts[idx] / total),
                material=infer_material_name(rgb),
            )
        )
    return palette


def make_line_mask(cad_rgb: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(cad_rgb, cv2.COLOR_RGB2GRAY)
    gray = cv2.GaussianBlur(gray, (3, 3), 0)

    _, otsu = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    adaptive = cv2.adaptiveThreshold(
        gray,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        31,
        9,
    )
    canny = cv2.Canny(gray, 60, 170)

    line_mask = cv2.bitwise_or(otsu, adaptive)
    line_mask = cv2.bitwise_or(line_mask, canny)
    line_mask = cv2.morphologyEx(line_mask, cv2.MORPH_CLOSE, np.ones((2, 2), np.uint8), iterations=1)

    # Keep text, thin walls, and hatch marks visible while preventing tiny specks
    # from driving segmentation.
    line_mask = cv2.dilate(line_mask, np.ones((2, 2), np.uint8), iterations=1)
    return line_mask > 0


def resize_for_sdxl(image: Image.Image, max_side: int = 1024, min_side: int = 512) -> Image.Image:
    width, height = image.size
    scale = max(min_side / min(width, height), 1.0)
    if max(width, height) * scale > max_side:
        scale = max_side / max(width, height)

    new_width = max(8, int(width * scale))
    new_height = max(8, int(height * scale))
    new_width = int(round(new_width / 8) * 8)
    new_height = int(round(new_height / 8) * 8)
    return image.convert("RGB").resize((new_width, new_height), Image.Resampling.LANCZOS)


def make_canny_control_image(cad_image: Image.Image) -> Image.Image:
    cad_rgb = pil_to_rgb_array(cad_image)
    gray = cv2.cvtColor(cad_rgb, cv2.COLOR_RGB2GRAY)
    gray = cv2.GaussianBlur(gray, (3, 3), 0)
    edges = cv2.Canny(gray, 70, 180)
    edges = cv2.dilate(edges, np.ones((2, 2), np.uint8), iterations=1)
    control = np.stack([edges, edges, edges], axis=-1)
    return Image.fromarray(control, mode="RGB")


def make_palette_style_canvas(size: tuple[int, int], palette: list[PaletteColor]) -> Image.Image:
    width, height = size
    palette_rgbs = [item.rgb for item in palette[:6]] or [
        (232, 221, 199),
        (204, 222, 214),
        (215, 224, 235),
        (224, 208, 212),
    ]

    rng = np.random.default_rng(42)
    low_w = max(16, width // 48)
    low_h = max(16, height // 48)
    palette_array = np.array(palette_rgbs, dtype=np.float32)
    weights = np.array([max(item.percent, 0.04) for item in palette[: len(palette_rgbs)]], dtype=np.float32)
    if len(weights) != len(palette_rgbs):
        weights = np.ones(len(palette_rgbs), dtype=np.float32)
    weights = weights / weights.sum()

    color_indices = rng.choice(len(palette_rgbs), size=(low_h, low_w), p=weights)
    canvas = palette_array[color_indices]
    canvas = cv2.resize(canvas, (width, height), interpolation=cv2.INTER_CUBIC)
    canvas = cv2.GaussianBlur(canvas, (0, 0), 18)

    white = np.full_like(canvas, 255)
    canvas = canvas * 0.68 + white * 0.32

    paper_noise = rng.normal(0, 3.5, size=canvas.shape).astype(np.float32)
    canvas = np.clip(canvas + paper_noise, 0, 255).astype(np.uint8)
    return Image.fromarray(canvas, mode="RGB").filter(ImageFilter.GaussianBlur(radius=0.6))


def overlay_original_linework(base_image: Image.Image, cad_image: Image.Image, strength: float) -> Image.Image:
    if strength <= 0:
        return base_image.convert("RGB")

    cad_resized = cad_image.convert("RGB").resize(base_image.size, Image.Resampling.LANCZOS)
    cad_rgb = pil_to_rgb_array(cad_resized)
    base_rgb = pil_to_rgb_array(base_image).astype(np.float32)
    line_mask = make_line_mask(cad_rgb)

    line_alpha = cv2.GaussianBlur(line_mask.astype(np.float32), (0, 0), 0.55)[..., None] * float(strength)
    line_tone = np.minimum(cad_rgb.astype(np.float32), 55)
    composited = base_rgb * (1 - line_alpha) + line_tone * line_alpha
    return Image.fromarray(np.clip(composited, 0, 255).astype(np.uint8), mode="RGB")


def palette_prompt_fragment(palette: list[PaletteColor]) -> str:
    colors = ", ".join(rgb_to_hex(item.rgb) for item in palette[:6])
    materials = ", ".join(item.material for item in palette[:4])
    return f"reference palette colors {colors}; material mood: {materials}"


def describe_plan_canvas(cad_image: Image.Image) -> str:
    width, height = cad_image.size
    aspect = width / max(height, 1)
    if aspect > 1.55:
        return "wide horizontal multi-unit floor plan composition"
    if aspect < 0.8:
        return "tall vertical architectural floor plan composition"
    return "balanced architectural floor plan composition"


def build_ai_prompt(palette: list[PaletteColor], prompt_hint: str, cad_image: Image.Image) -> str:
    user_hint = prompt_hint.strip() if prompt_hint else "top-down furnished real estate floor plan render"
    return (
        f"{user_hint}, high quality top-down architectural visualization, furnished apartment plan, "
        "white walls, wood flooring, marble and tile floors, beds, sofas, dining tables, kitchen counters, "
        "bathroom fixtures, plants, balconies, realistic material textures, clean real estate marketing plan, "
        "orthographic top view, crisp room boundaries, bright professional render, "
        f"{describe_plan_canvas(cad_image)}, "
        "render the floor plan as a finished colored marketing image, not as a CAD drawing, "
        "avoid black blueprint linework, avoid engineering symbols, avoid title blocks, avoid logos, "
        f"{palette_prompt_fragment(palette)}"
    )


@lru_cache(maxsize=1)
def load_text_to_image_pipeline():
    import torch
    from diffusers import AutoPipelineForText2Image

    use_cuda = torch.cuda.is_available()
    if not use_cuda and IS_HF_SPACE:
        raise RuntimeError("AI mode needs GPU or ZeroGPU hardware. Please switch the Hugging Face Space hardware.")
    if not use_cuda and os.getenv("PLANPALETTE_ALLOW_CPU", "1") != "1":
        raise RuntimeError("No CUDA GPU found. Set PLANPALETTE_ALLOW_CPU=1 to try very slow CPU inference.")

    dtype = torch.float16 if use_cuda else torch.float32
    pipe = AutoPipelineForText2Image.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=dtype,
        use_safetensors=True,
    )
    if use_cuda:
        pipe.enable_model_cpu_offload()
    else:
        pipe.to("cpu")
    pipe.enable_attention_slicing()
    return pipe


def _ai_colorize_floor_plan(
    reference_image: Image.Image,
    cad_image: Image.Image,
    palette: list[PaletteColor],
    prompt_hint: str,
    steps: int,
    linework_strength: float,
) -> Image.Image:
    del reference_image

    pipe = load_text_to_image_pipeline()
    default_max_side = "1024" if IS_HF_SPACE else "640"
    model_cad = resize_for_sdxl(cad_image, max_side=int(os.getenv("PLANPALETTE_MAX_SIDE", default_max_side)))
    prompt = build_ai_prompt(palette, prompt_hint, model_cad)

    result = pipe(
        prompt=prompt,
        num_inference_steps=int(steps),
        guidance_scale=1.0,
        width=model_cad.width,
        height=model_cad.height,
    ).images[0]

    return overlay_original_linework(result, model_cad, linework_strength)


if spaces is not None and IS_HF_SPACE:
    ai_colorize_floor_plan = spaces.GPU(duration=60)(_ai_colorize_floor_plan)
else:
    ai_colorize_floor_plan = _ai_colorize_floor_plan


def connected_region_map(line_mask: np.ndarray) -> tuple[np.ndarray, int]:
    height, width = line_mask.shape
    gap_closed_lines = cv2.dilate(line_mask.astype(np.uint8) * 255, np.ones((5, 5), np.uint8), iterations=1)
    fillable = cv2.bitwise_not(gap_closed_lines)

    fillable = cv2.morphologyEx(fillable, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=1)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(fillable, connectivity=8)

    min_area = max(220, int(height * width * 0.004))
    region_map = np.zeros((height, width), dtype=np.int32)
    region_id = 1

    image_area = height * width
    for label in range(1, num_labels):
        area = int(stats[label, cv2.CC_STAT_AREA])
        if area < min_area or area > int(image_area * 0.92):
            continue

        component = labels == label
        component = cv2.morphologyEx(component.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((9, 9), np.uint8), iterations=1)
        component = component.astype(bool) & ~line_mask
        if int(component.sum()) < min_area:
            continue
        region_map[component] = region_id
        region_id += 1

    if region_id <= 2:
        region_map, region_id = fallback_grid_regions(line_mask)

    return region_map, region_id - 1


def fallback_grid_regions(line_mask: np.ndarray) -> tuple[np.ndarray, int]:
    height, width = line_mask.shape
    region_map = np.zeros((height, width), dtype=np.int32)
    fillable = ~line_mask
    region_id = 1
    rows, cols = 4, 4
    min_area = max(120, int(height * width * 0.002))

    for row in range(rows):
        for col in range(cols):
            y0 = int(row * height / rows)
            y1 = int((row + 1) * height / rows)
            x0 = int(col * width / cols)
            x1 = int((col + 1) * width / cols)
            tile = fillable[y0:y1, x0:x1]
            if int(tile.sum()) < min_area:
                continue
            region_map[y0:y1, x0:x1][tile] = region_id
            region_id += 1

    return region_map, region_id


def soften_palette_color(rgb: tuple[int, int, int], index: int) -> np.ndarray:
    color = np.array(rgb, dtype=np.float32)
    white = np.array([255, 255, 255], dtype=np.float32)
    softened = color * 0.54 + white * 0.46

    # Slight alternating warmth/coolness keeps adjacent rooms readable even when
    # the source palette has several near-neutrals.
    offsets = np.array(
        [
            [10, 5, -2],
            [-4, 5, 10],
            [6, -2, 5],
            [-2, 9, -3],
            [8, 2, 8],
            [-5, 4, 4],
        ],
        dtype=np.float32,
    )
    return np.clip(softened + offsets[index % len(offsets)], 0, 255)


def colorize_regions(cad_rgb: np.ndarray, line_mask: np.ndarray, region_map: np.ndarray, palette: list[PaletteColor]) -> np.ndarray:
    height, width = line_mask.shape
    fill_layer = np.full((height, width, 3), 255, dtype=np.float32)
    palette_rgbs = [item.rgb for item in palette] or [(218, 205, 184), (188, 210, 198), (201, 213, 228)]

    region_ids = [idx for idx in np.unique(region_map) if idx > 0]
    for assignment_index, region_id in enumerate(region_ids):
        mask = region_map == region_id
        ys, xs = np.where(mask)
        if len(xs) == 0:
            continue

        centroid_bias = int((xs.mean() / max(width, 1)) * 2 + (ys.mean() / max(height, 1)) * 3)
        palette_index = (assignment_index + centroid_bias) % len(palette_rgbs)
        base_color = soften_palette_color(palette_rgbs[palette_index], assignment_index)
        fill_layer[mask] = base_color

    region_alpha = (region_map > 0).astype(np.float32)
    region_alpha = cv2.GaussianBlur(region_alpha, (0, 0), 1.35)
    region_alpha = np.clip(region_alpha[..., None] * 0.78, 0, 0.78)

    cad_float = cad_rgb.astype(np.float32)
    brightened_cad = cad_float * 0.45 + 255 * 0.55
    colorized = brightened_cad * (1 - region_alpha) + fill_layer * region_alpha

    subtle_shadow = cv2.GaussianBlur(line_mask.astype(np.float32), (0, 0), 2.2)[..., None]
    colorized = colorized * (1 - subtle_shadow * 0.08)

    line_alpha = cv2.GaussianBlur(line_mask.astype(np.float32), (0, 0), 0.45)[..., None]
    original_line_tone = np.minimum(cad_float, 35)
    composited = colorized * (1 - line_alpha) + original_line_tone * line_alpha
    return np.clip(composited, 0, 255).astype(np.uint8)


def build_legend_html(palette: list[PaletteColor], region_count: int | None = None) -> str:
    if not palette:
        return "<div class='legend-empty'>Upload a reference image to extract a palette.</div>"

    swatches = []
    for item in palette:
        hex_color = rgb_to_hex(item.rgb)
        label = html.escape(item.material.title())
        swatches.append(
            f"""
            <div class="swatch-row">
                <span class="swatch" style="background:{hex_color};"></span>
                <div class="swatch-copy">
                    <strong>{hex_color}</strong>
                    <span>{label} - {item.percent * 100:.1f}%</span>
                </div>
            </div>
            """
        )

    return f"""
    <section class="legend-panel">
        <div class="legend-stat">
            <strong>{len(palette)}</strong>
            <span>reference colors guiding the image model</span>
        </div>
        <div class="legend-list">
            {''.join(swatches)}
        </div>
    </section>
    """


def transfer_style(
    reference_image: Image.Image | None,
    cad_image: Image.Image | None,
    palette_size: int,
    prompt_hint: str,
    steps: int,
    linework_strength: float,
) -> tuple[Image.Image | None, str]:
    if reference_image is None or cad_image is None:
        return None, "<div class='legend-empty'>Upload both floor plans, then run PlanPalette.</div>"

    reference_rgb = pil_to_rgb_array(reference_image)
    palette = extract_palette(reference_rgb, k=palette_size)

    try:
        final = ai_colorize_floor_plan(
            reference_image,
            cad_image,
            palette,
            prompt_hint,
            steps,
            linework_strength,
        )
    except Exception as exc:
        escaped = html.escape(str(exc))
        return None, f"<div class='legend-empty'>AI generation failed: {escaped}</div>"

    return final, build_legend_html(palette)


CUSTOM_CSS = """
:root {
    --pp-ink: #171717;
    --pp-muted: #5c646f;
    --pp-line: #d8dde3;
    --pp-surface: #f8f7f4;
    --pp-accent: #1f7a6d;
    --pp-accent-strong: #145a51;
}

.gradio-container {
    max-width: 1180px !important;
    margin: 0 auto;
    color: var(--pp-ink);
    background:
        linear-gradient(180deg, rgba(248, 247, 244, 0.98), rgba(246, 248, 249, 0.98));
}

.pp-header {
    padding: 18px 0 8px;
    border-bottom: 1px solid var(--pp-line);
    margin-bottom: 14px;
}

.pp-title {
    margin: 0;
    font-size: clamp(2rem, 3vw, 3.2rem);
    line-height: 1.02;
    font-weight: 780;
    letter-spacing: 0;
}

.pp-subtitle {
    margin: 8px 0 0;
    max-width: 760px;
    color: var(--pp-muted);
    font-size: 1rem;
    line-height: 1.5;
}

.pp-panel {
    border: 1px solid var(--pp-line) !important;
    border-radius: 8px !important;
    background: rgba(255, 255, 255, 0.82) !important;
}

.pp-run-button {
    min-height: 46px;
    border-radius: 6px !important;
    background: var(--pp-accent) !important;
    border-color: var(--pp-accent) !important;
    color: white !important;
    font-weight: 700 !important;
}

.pp-run-button:hover {
    background: var(--pp-accent-strong) !important;
}

.legend-panel {
    border: 1px solid var(--pp-line);
    border-radius: 8px;
    background: #ffffff;
    padding: 14px;
}

.legend-stat {
    display: flex;
    align-items: baseline;
    gap: 10px;
    padding-bottom: 12px;
    margin-bottom: 12px;
    border-bottom: 1px solid var(--pp-line);
}

.legend-stat strong {
    font-size: 1.75rem;
    line-height: 1;
}

.legend-stat span,
.swatch-copy span,
.legend-empty {
    color: var(--pp-muted);
}

.legend-list {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(190px, 1fr));
    gap: 10px;
}

.swatch-row {
    display: flex;
    gap: 10px;
    align-items: center;
    min-width: 0;
}

.swatch {
    width: 36px;
    height: 36px;
    flex: 0 0 auto;
    border-radius: 6px;
    border: 1px solid rgba(0, 0, 0, 0.12);
    box-shadow: inset 0 0 0 1px rgba(255, 255, 255, 0.32);
}

.swatch-copy {
    min-width: 0;
    display: flex;
    flex-direction: column;
    gap: 2px;
}

.swatch-copy strong {
    font-size: 0.92rem;
}

.swatch-copy span {
    font-size: 0.84rem;
    line-height: 1.25;
}

.legend-empty {
    border: 1px dashed var(--pp-line);
    border-radius: 8px;
    background: #ffffff;
    padding: 16px;
}
"""


with gr.Blocks(title=APP_TITLE, css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="teal", neutral_hue="slate")) as demo:
    gr.HTML(
        f"""
        <header class="pp-header">
            <h1 class="pp-title">{APP_TITLE}</h1>
            <p class="pp-subtitle">{APP_SUBTITLE}</p>
        </header>
        """
    )

    with gr.Row(equal_height=True):
        with gr.Column(scale=1, elem_classes=["pp-panel"]):
            reference_input = gr.Image(
                label="Reference Styled Floor Plan",
                type="pil",
                image_mode="RGB",
                height=360,
            )
        with gr.Column(scale=1, elem_classes=["pp-panel"]):
            cad_input = gr.Image(
                label="Raw CAD Floor Plan",
                type="pil",
                image_mode="RGB",
                height=360,
            )

    with gr.Row():
        palette_size = gr.Slider(
            minimum=3,
            maximum=8,
            value=6,
            step=1,
            label="Palette Size",
            info="Number of dominant reference colors to transfer.",
        )
        steps = gr.Slider(
            minimum=2,
            maximum=8,
            value=4,
            step=1,
            label="AI Steps",
            info="Lightning/turbo models work best at low step counts.",
        )
        linework_strength = gr.Slider(
            minimum=0,
            maximum=0.6,
            value=0,
            step=0.02,
            label="CAD Line Overlay",
            info="Set to 0 for pure AI render.",
        )

    with gr.Row():
        prompt_hint = gr.Textbox(
            label="Style Hint",
            value="top-down furnished real estate floor plan render like an architectural marketing brochure",
            lines=2,
        )
        run_button = gr.Button("Generate Colorized Plan", variant="primary", elem_classes=["pp-run-button"])

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Final PNG",
                type="pil",
                image_mode="RGB",
                format="png",
                height=460,
            )
        with gr.Column(scale=1):
            legend_output = gr.HTML(
                value="<div class='legend-empty'>Upload both floor plans, then run PlanPalette.</div>",
                label="Palette / Material Legend",
            )

    run_button.click(
        fn=transfer_style,
        inputs=[reference_input, cad_input, palette_size, prompt_hint, steps, linework_strength],
        outputs=[output_image, legend_output],
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)