import io
from pathlib import Path

import cv2
import numpy as np
import pandas as pd
import streamlit as st
from PIL import Image

st.set_page_config(page_title="Saree Pixel Fill Assistant", layout="wide")

SAMPLE_DIR = Path(__file__).parent / "sample_images"  # images directly inside src/


def pil_to_bgr(pil_img):
    arr = np.array(pil_img.convert("RGB"))
    return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)


def bgr_to_pil(bgr):
    return Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))


def make_line_mask(bgr, threshold=215, close_px=1):
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    line_mask = (gray < threshold).astype(np.uint8) * 255

    if close_px > 0:
        k = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE,
            (2 * close_px + 1, 2 * close_px + 1)
        )
        line_mask = cv2.morphologyEx(line_mask, cv2.MORPH_CLOSE, k, iterations=1)

    return line_mask


def find_enclosed_regions(line_mask, min_area=20, max_area=5000):
    h, w = line_mask.shape

    white = (line_mask == 0).astype(np.uint8) * 255
    flood = white.copy()
    mask = np.zeros((h + 2, w + 2), np.uint8)

    cv2.floodFill(flood, mask, (0, 0), 128)

    enclosed = (flood == 255).astype(np.uint8) * 255

    n, labels, stats, _ = cv2.connectedComponentsWithStats(enclosed, 8)

    region_mask = np.zeros_like(enclosed)
    regions = []

    for i in range(1, n):
        area = int(stats[i, cv2.CC_STAT_AREA])
        x = int(stats[i, cv2.CC_STAT_LEFT])
        y = int(stats[i, cv2.CC_STAT_TOP])
        ww = int(stats[i, cv2.CC_STAT_WIDTH])
        hh = int(stats[i, cv2.CC_STAT_HEIGHT])

        touches_border = x == 0 or y == 0 or (x + ww) >= w or (y + hh) >= h

        if not touches_border and min_area <= area <= max_area:
            region_mask[labels == i] = 255
            regions.append({
                "id": i,
                "area_px": area,
                "x": x,
                "y": y,
                "w": ww,
                "h": hh
            })

    return region_mask, regions


def generate_step3(bgr, threshold, close_px, min_area, max_area, fill_regions, fill_color):
    line_mask = make_line_mask(bgr, threshold, close_px)
    region_mask, regions = find_enclosed_regions(line_mask, min_area, max_area)

    out = np.ones_like(bgr) * 255

    out[line_mask > 0] = fill_color

    if fill_regions:
        out[region_mask > 0] = fill_color

    return out, line_mask, region_mask, regions


def overlay_regions(bgr, region_mask, line_mask):
    out = bgr.copy()
    overlay = out.copy()

    overlay[region_mask > 0] = (0, 255, 255)
    overlay[line_mask > 0] = (0, 0, 255)

    return cv2.addWeighted(overlay, 0.45, out, 0.55, 0)


def resize_pixel_art(bgr, scale):
    h, w = bgr.shape[:2]
    return cv2.resize(
        bgr,
        (w * scale, h * scale),
        interpolation=cv2.INTER_NEAREST
    )


def to_bytes_bmp(bgr):
    bio = io.BytesIO()
    bgr_to_pil(bgr).save(bio, format="BMP")
    return bio.getvalue()


st.title("Procelevate Saree Pixel Fill Assistant")
st.caption("Prototype: Step 2 line design → Step 3 pixel separation → Step 4 resize preview")

with st.sidebar:
    st.header("Controls")

    threshold = st.slider("Line detection threshold", 80, 255, 215)
    close_px = st.slider("Close tiny line gaps", 0, 5, 1)
    min_area = st.slider("Minimum fill region area", 1, 1000, 20)
    max_area = st.number_input("Maximum fill region area", min_value=100, value=5000)

    fill_mode = st.selectbox(
        "Fill mode",
        [
            "Line only",
            "Fill enclosed regions",
            "Fill motif body"
        ]
    )

    colour = st.color_picker("Step 3 fill/line colour", "#FF0000")

    scale = st.slider("Step 4 resize scale", 2, 12, 4)

fill_rgb = tuple(int(colour.lstrip("#")[i:i + 2], 16) for i in (0, 2, 4))
fill_bgr = (fill_rgb[2], fill_rgb[1], fill_rgb[0])

sample_files = sorted([
    p for p in SAMPLE_DIR.iterdir()
    if p.suffix.lower() in [".bmp", ".png", ".jpg", ".jpeg"]
])

st.subheader("Choose Step 2 image")

source = st.radio(
    "Image source",
    ["Use sample image from Space", "Upload image"],
    horizontal=True
)

pil = None

if source == "Use sample image from Space":
    if sample_files:
        selected = st.selectbox(
            "Select sample Step 2 image",
            sample_files,
            format_func=lambda p: p.name
        )
        pil = Image.open(selected).convert("RGB")
    else:
        st.warning("No sample images found inside src folder.")
else:
    uploaded = st.file_uploader(
        "Upload Step 2 image",
        type=["bmp", "png", "jpg", "jpeg"]
    )
    if uploaded:
        pil = Image.open(uploaded).convert("RGB")

if pil:
    bgr = pil_to_bgr(pil)

    fill_regions = fill_mode != "Line only"

    step3, line_mask, region_mask, regions = generate_step3(
        bgr,
        threshold,
        close_px,
        min_area,
        int(max_area),
        fill_regions,
        fill_bgr
    )

    overlay = overlay_regions(bgr, region_mask, line_mask)
    resized = resize_pixel_art(step3, scale)

    c1, c2, c3 = st.columns(3)

    with c1:
        st.subheader("Input Step 2")
        st.image(pil, use_column_width=True)

    with c2:
        st.subheader("Detected Regions")
        st.image(bgr_to_pil(overlay), use_column_width=True)

    with c3:
        st.subheader("Generated Step 3")
        st.image(bgr_to_pil(step3), use_column_width=True)

    st.subheader("Step 4 Resize Preview")
    st.image(bgr_to_pil(resized), use_column_width=True)

    m1, m2, m3 = st.columns(3)
    m1.metric("Detected regions", len(regions))
    m2.metric("Line pixels", int((line_mask > 0).sum()))
    m3.metric("Filled pixels", int((region_mask > 0).sum()))

    st.download_button(
        "Download Step 3 BMP",
        to_bytes_bmp(step3),
        file_name="generated_step3.bmp"
    )

    st.download_button(
        "Download Step 4 Preview BMP",
        to_bytes_bmp(resized),
        file_name="generated_step4_preview.bmp"
    )

    with st.expander("Region details"):
        if regions:
            st.dataframe(pd.DataFrame(regions).head(500), use_container_width=True)
        else:
            st.info("No enclosed regions detected with current settings.")