import streamlit as st
import requests
import zipfile
import io
import os
import sys
import torch
import numpy as np
import matplotlib.pyplot as plt
import cv2
from PIL import Image
from diffusers import StableDiffusionInpaintPipeline, EulerDiscreteScheduler
import copy

# ========== Download SAM Repo ==========
def download_sam_repo():
    repo_url = "https://github.com/facebookresearch/segment-anything/archive/refs/heads/main.zip"
    repo_dir = "segment-anything-main"
    if not os.path.exists(repo_dir):
        st.info("🔽 Downloading Segment Anything repo...")
        response = requests.get(repo_url)
        if response.status_code == 200:
            with zipfile.ZipFile(io.BytesIO(response.content)) as zip_ref:
                zip_ref.extractall(".")
            st.success("✅ Segment Anything repo downloaded!")
        else:
            st.error(f"❌ Failed to download repo: {response.status_code}")

download_sam_repo()
sys.path.append(os.path.abspath("segment-anything-main"))

from segment_anything import sam_model_registry, SamAutomaticMaskGenerator

# ========== Download SAM Model Checkpoint ==========
def download_file(url, output_path):
    if not os.path.exists(output_path):
        st.info(f"🔽 Downloading {os.path.basename(output_path)}...")
        response = requests.get(url, stream=True)
        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        st.success(f"✅ Downloaded {os.path.basename(output_path)}")

sam_url = "https://huggingface.co/camenduru/segment_anything/resolve/main/sam_vit_h_4b8939.pth"
download_file(sam_url, "sam_vit_h_4b8939.pth")

# ========== Load SAM Model ==========
@st.cache_resource
def load_sam_model():
    sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
    sam.to("cuda")
    mask_generator = SamAutomaticMaskGenerator(
        model=sam,
        points_per_side=32,
        pred_iou_thresh=0.99,
        stability_score_thresh=0.92,
        crop_n_layers=1,
        crop_n_points_downscale_factor=2,
        min_mask_region_area=100,
    )
    return mask_generator

# ========== Load SD Model ==========
@st.cache_resource
def load_sd_pipeline():
    model_dir = 'stabilityai/stable-diffusion-2-inpainting'
    scheduler = EulerDiscreteScheduler.from_pretrained(model_dir, subfolder='scheduler')
    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        model_dir,
        scheduler=scheduler,
        torch_dtype=torch.float16,
        revision="fp16"
    ).to("cuda")
    pipe.enable_xformers_memory_efficient_attention()
    return pipe

# ========== Display masks ==========
def show_masks(image, masks):
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(image)
    for i, mask in enumerate(masks):
        m = mask['segmentation']
        color = np.random.random((1, 3)).tolist()[0]
        img = np.ones((m.shape[0], m.shape[1], 3))
        for j in range(3):
            img[:, :, j] = color[j]
        ax.imshow(np.dstack((img, m * 0.35)))
        contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            cnt = contours[0]
            M = cv2.moments(cnt)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
                ax.text(cx, cy, str(i), color='white', fontsize=16, ha='center', va='center', fontweight='bold')
    ax.axis('off')
    st.pyplot(fig)

# ========== Image Grid ==========
def create_image_grid(original_image, images, names, rows, columns):
    images = copy.copy(images)
    names = copy.copy(names)

    images.insert(0, original_image)
    names.insert(0, "Original")

    fig, axes = plt.subplots(rows, columns, figsize=(15, 15))
    for idx, (img, name) in enumerate(zip(images, names)):
        row, col = divmod(idx, columns)
        axes[row, col].imshow(img)
        axes[row, col].set_title(name)
        axes[row, col].axis('off')
    for idx in range(len(images), rows * columns):
        row, col = divmod(idx, columns)
        axes[row, col].axis('off')
    plt.tight_layout()
    st.pyplot(fig)

# ========== Streamlit UI ==========
st.title("🎨 Segment & Inpaint with Streamlit")

uploaded_file = st.file_uploader("Upload an Image", type=['png', 'jpg', 'jpeg'])

if uploaded_file:
    source_image = Image.open(uploaded_file).convert("RGB").resize((512, 512))
    st.image(source_image, caption="Uploaded Image", use_column_width=True)

    mask_gen = load_sam_model()
    masks = mask_gen.generate(np.asarray(source_image))
    st.write(f"Number of Segments Found: {len(masks)}")

    show_masks(source_image, masks)

    mask_idx = st.number_input(f"Select Mask Index (0 to {len(masks)-1})", min_value=0, max_value=len(masks)-1, value=0)
    prompt = st.text_input("Enter Inpainting Prompt", "a skirt full of text")
    generate = st.button("Generate Inpainting")

    if generate:
        segmentation_mask = masks[mask_idx]['segmentation']
        stable_mask = Image.fromarray(segmentation_mask * 255).convert("RGB")

        pipe = load_sd_pipeline()
        generator = torch.Generator(device="cuda").manual_seed(77)

        images = []
        for i in range(4):
            result = pipe(
                prompt=prompt,
                guidance_scale=7.5,
                num_inference_steps=50,
                generator=generator,
                image=source_image,
                mask_image=stable_mask
            ).images[0]
            images.append(result)

        create_image_grid(source_image, images, [prompt]*4, 2, 3)