import io
import json
import os
import tempfile
import zipfile
from dataclasses import dataclass
from pathlib import Path

import cairosvg
import cv2
import gradio as gr
import numpy as np
from PIL import Image

MODEL_AVAILABLE = False
MODEL_ERROR = None
processor = None
seg_pipeline = None

try:
    from transformers import pipeline
    MODEL_AVAILABLE = True
except Exception as e:
    MODEL_ERROR = str(e)


@dataclass
class PlanConfig:
    min_room_area_ratio: float = 0.02
    wall_thickness: int = 6
    variation: str = "balanced"
    max_rooms: int = 6
    use_vision_model: bool = False


def pil_to_bgr(img: Image.Image) -> np.ndarray:
    arr = np.array(img.convert("RGB"))
    return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)


def bgr_to_pil(arr: np.ndarray) -> Image.Image:
    rgb = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(rgb)


def load_any_image(image_file):
    if image_file is None:
        raise gr.Error("Upload a PNG, JPG, or SVG file first.")
    path = image_file if isinstance(image_file, str) else getattr(image_file, "name", None)
    if not path:
        raise gr.Error("Could not read uploaded file.")
    ext = os.path.splitext(path)[1].lower()
    if ext == ".svg":
        png_bytes = cairosvg.svg2png(url=path)
        return Image.open(io.BytesIO(png_bytes)).convert("RGB"), ext
    return Image.open(path).convert("RGB"), ext


def init_segmentation_pipeline():
    global seg_pipeline, MODEL_AVAILABLE, MODEL_ERROR
    if seg_pipeline is not None:
        return seg_pipeline
    if not MODEL_AVAILABLE:
        return None
    try:
        seg_pipeline = pipeline(
            task="image-segmentation",
            model="nvidia/segformer-b0-finetuned-ade-512-512",
            device=-1,
        )
        return seg_pipeline
    except Exception as e:
        MODEL_ERROR = str(e)
        return None


def preprocess_and_extract_boundary(img_bgr: np.ndarray):
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    _, th = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    kernel = np.ones((5, 5), np.uint8)
    th = cv2.morphologyEx(th, cv2.MORPH_CLOSE, kernel, iterations=2)
    th = cv2.morphologyEx(th, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=1)
    contours, _ = cv2.findContours(th, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        raise ValueError("No closed shape detected. Use a darker outline on light background.")
    contour = max(contours, key=cv2.contourArea)
    peri = cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, 0.01 * peri, True)
    mask = np.zeros_like(gray)
    cv2.drawContours(mask, [contour], -1, 255, thickness=cv2.FILLED)
    return th, contour[:, 0, :], approx[:, 0, :], mask


def mask_from_vision_model(image: Image.Image, boundary_mask: np.ndarray):
    pipe = init_segmentation_pipeline()
    if pipe is None:
        return None, f"Vision model unavailable: {MODEL_ERROR}" if MODEL_ERROR else "Vision model unavailable"
    try:
        preds = pipe(image)
        if not preds:
            return None, "Vision model returned no masks"
        room_mask = np.zeros(boundary_mask.shape, dtype=np.uint8)
        for pred in preds:
            label = str(pred.get("label", "")).lower()
            if any(k in label for k in ["floor", "rug", "carpet", "room", "interior"]):
                seg = pred.get("mask")
                if isinstance(seg, Image.Image):
                    seg = np.array(seg.convert("L"))
                else:
                    seg = np.array(seg)
                if seg.ndim == 3:
                    seg = seg[..., 0]
                seg = cv2.resize(seg.astype(np.uint8), (boundary_mask.shape[1], boundary_mask.shape[0]))
                room_mask = np.maximum(room_mask, (seg > 0).astype(np.uint8) * 255)
        room_mask = cv2.bitwise_and(room_mask, boundary_mask)
        if np.count_nonzero(room_mask) == 0:
            return None, "Vision model found no interior-like regions"
        return room_mask, "Vision model used for interior proposal"
    except Exception as e:
        return None, f"Vision inference failed: {e}"


def axis_aligned_split(bounds, orientation, pos, gap=8):
    x1, y1, x2, y2 = bounds
    if orientation == "v":
        return [(x1, y1, pos - gap, y2), (pos + gap, y1, x2, y2)]
    return [(x1, y1, x2, pos - gap), (x1, pos + gap, x2, y2)]


def rect_inside_mask(rect, mask, fill_ratio=0.72):
    x1, y1, x2, y2 = [int(v) for v in rect]
    if x2 <= x1 or y2 <= y1:
        return False
    roi = mask[max(0, y1):max(0, y2), max(0, x1):max(0, x2)]
    if roi.size == 0:
        return False
    return (roi > 0).mean() >= fill_ratio


def generate_rooms(mask: np.ndarray, cfg: PlanConfig):
    ys, xs = np.where(mask > 0)
    x1, y1, x2, y2 = xs.min(), ys.min(), xs.max(), ys.max()
    rooms = [(x1, y1, x2, y2)]
    target = {"compact": 6, "balanced": 5, "open": 3}.get(cfg.variation, 5)
    target = min(target, cfg.max_rooms)
    for _ in range(target * 4):
        if len(rooms) >= target:
            break
        areas = [max(0, (r[2] - r[0]) * (r[3] - r[1])) for r in rooms]
        idx = int(np.argmax(areas))
        r = rooms.pop(idx)
        rx1, ry1, rx2, ry2 = r
        w, h = rx2 - rx1, ry2 - ry1
        if w * h < mask.shape[0] * mask.shape[1] * cfg.min_room_area_ratio:
            rooms.append(r)
            continue
        if cfg.variation == "open":
            orientation = "v" if w > h * 1.25 else "h"
            cuts = [0.5]
        elif cfg.variation == "compact":
            orientation = "h" if h > w else "v"
            cuts = [0.45, 0.55]
        else:
            orientation = "v" if w >= h else "h"
            cuts = [0.5, 0.4, 0.6]
        split_done = False
        for c in cuts:
            pos = int(rx1 + w * c) if orientation == "v" else int(ry1 + h * c)
            kids = axis_aligned_split(r, orientation, pos, gap=cfg.wall_thickness)
            valid = [k for k in kids if rect_inside_mask(k, mask)]
            if len(valid) == 2:
                rooms.extend(valid)
                split_done = True
                break
        if not split_done:
            rooms.append(r)
    return rooms


def polygon_from_mask_component(component_mask):
    contours, _ = cv2.findContours(component_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None
    c = max(contours, key=cv2.contourArea)
    return c


def draw_floorplan(boundary_mask, room_seed_mask, rooms, contour_points, approx_points, cfg: PlanConfig, vision_note: str):
    h, w = boundary_mask.shape
    canvas = np.full((h, w, 3), 255, np.uint8)
    palette = [(239, 246, 255), (236, 253, 245), (255, 247, 237), (250, 245, 255), (254, 242, 242), (240, 253, 250)]
    room_json = []

    for i, r in enumerate(rooms):
        x1, y1, x2, y2 = [int(v) for v in r]
        room_mask = np.zeros_like(boundary_mask)
        cv2.rectangle(room_mask, (x1, y1), (x2, y2), 255, thickness=-1)
        room_mask = cv2.bitwise_and(room_mask, room_seed_mask)
        c = polygon_from_mask_component(room_mask)
        if c is None:
            continue
        cv2.drawContours(canvas, [c], -1, palette[i % len(palette)], thickness=cv2.FILLED)
        cv2.drawContours(canvas, [c], -1, (0, 0, 0), thickness=2)
        M = cv2.moments(c)
        if M["m00"]:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            cv2.putText(canvas, f"Room {i+1}", (cx - 28, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (20, 20, 20), 1, cv2.LINE_AA)
        room_json.append({"id": f"room_{i+1}", "bbox": [x1, y1, x2, y2], "area_px": int(cv2.contourArea(c)), "polygon": c[:, 0, :].tolist()})

    cv2.drawContours(canvas, [contour_points.reshape(-1, 1, 2)], -1, (0, 0, 0), thickness=max(3, cfg.wall_thickness + 2))
    cv2.putText(canvas, vision_note, (16, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (60, 60, 60), 1, cv2.LINE_AA)

    result = {
        "boundary": contour_points.tolist(),
        "boundary_simplified": approx_points.tolist(),
        "rooms": room_json,
        "config": cfg.__dict__,
        "meta": {"image_size": [int(w), int(h)], "vision_note": vision_note},
    }
    return canvas, result


def export_svg(result, width, height):
    def pts_to_path(pts):
        if not pts:
            return ""
        start = f"M {pts[0][0]} {pts[0][1]}"
        rest = " ".join([f"L {p[0]} {p[1]}" for p in pts[1:]])
        return f"{start} {rest} Z"
    colors = ["#dbeafe", "#dcfce7", "#ffedd5", "#f3e8ff", "#fee2e2", "#ccfbf1"]
    parts = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">', '<rect width="100%" height="100%" fill="white"/>']
    for i, room in enumerate(result["rooms"]):
        parts.append(f'<path d="{pts_to_path(room["polygon"])}" fill="{colors[i % len(colors)]}" stroke="#111" stroke-width="2" />')
    parts.append(f'<path d="{pts_to_path(result["boundary"])}" fill="none" stroke="#000" stroke-width="6" />')
    parts.append('</svg>')
    return "\n".join(parts)


def process(image_file, variation, wall_thickness, max_rooms, use_vision_model):
    image, source_ext = load_any_image(image_file)
    cfg = PlanConfig(variation=variation, wall_thickness=wall_thickness, max_rooms=max_rooms, use_vision_model=use_vision_model)
    img_bgr = pil_to_bgr(image)
    binary, contour, approx, boundary_mask = preprocess_and_extract_boundary(img_bgr)

    room_seed_mask = boundary_mask.copy()
    vision_note = "Heuristic partition mode"
    if use_vision_model:
        proposed_mask, note = mask_from_vision_model(image, boundary_mask)
        if proposed_mask is not None:
            room_seed_mask = proposed_mask
        vision_note = note

    rooms = generate_rooms(room_seed_mask, cfg)
    rendered, result = draw_floorplan(boundary_mask, room_seed_mask, rooms, contour, approx, cfg, vision_note)
    result["meta"]["source_ext"] = source_ext

    svg = export_svg(result, rendered.shape[1], rendered.shape[0])
    json_str = json.dumps(result, indent=2)
    preview = bgr_to_pil(rendered)
    binary_preview = Image.fromarray(room_seed_mask)

    tmpdir = Path(tempfile.mkdtemp(prefix="floorplan_space_"))
    zip_path = tmpdir / "floorplan_bundle.zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("floorplan.json", json_str)
        zf.writestr("floorplan.svg", svg)

    return preview, binary_preview, json_str, str(zip_path)


with gr.Blocks(title="Closed Shape to Floor Plan") as demo:
    gr.Markdown(
        "# Closed Shape to Floor Plan\n"
        "Upload a closed-outline PNG, JPG, or SVG. Enable the vision-model option to let a Hugging Face segmentation model propose interior regions before room partitioning."
    )
    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.File(file_types=[".png", ".jpg", ".jpeg", ".svg"], label="Closed shape file (PNG/JPG/SVG)")
            variation = gr.Radio(["balanced", "compact", "open"], value="balanced", label="Variation")
            wall_thickness = gr.Slider(2, 16, value=6, step=1, label="Wall thickness")
            max_rooms = gr.Slider(2, 8, value=5, step=1, label="Max rooms")
            use_vision_model = gr.Checkbox(value=True, label="Use vision model for interior proposal")
            btn = gr.Button("Generate floor plan", variant="primary")
        with gr.Column(scale=1):
            out = gr.Image(label="Generated floor plan")
            binary = gr.Image(label="Interior / room seed mask")
    structured = gr.Code(label="Structured JSON", language="json")
    archive = gr.File(label="Download JSON + SVG")

    btn.click(process, inputs=[inp, variation, wall_thickness, max_rooms, use_vision_model], outputs=[out, binary, structured, archive])

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), ssr_mode=False)