import os

# ── Set cache/config dirs ──
hf_home = "/data/.cache/huggingface"
yolo_cfg = "/data/ultralytics"
os.makedirs(hf_home, exist_ok=True)
os.makedirs(yolo_cfg, exist_ok=True)
os.environ["HF_HOME"] = hf_home
os.environ["YOLO_CONFIG_DIR"] = yolo_cfg

from ultralytics import YOLO
import numpy as np
import torch
from PIL import Image
import cv2
from diffusers import StableDiffusionXLInpaintPipeline
import gradio as gr

# ---- utils ----
def pil_to_cv2(pil_img):
    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

def cv2_to_pil(cv_img):
    return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))

# ---- load models ----
yolo = YOLO("yolov8x-seg.pt")
inpaint_pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
    torch_dtype=torch.float16,
    use_safetensors=True,
    use_auth_token=os.getenv("HF_TOKEN")
).to("cuda")

# ---- processing logic ----
def run_background_removal_and_inpaint(image, prompt, negative_prompt):
    if image is None:
        raise gr.Error("Please upload an image.")

    img_cv = pil_to_cv2(image)
    results = yolo(img_cv)

    if not results or not results[0].masks or len(results[0].masks.data) == 0:
        raise gr.Error("No subject detected in the image. Please upload a clearer photo.")

    mask = results[0].masks.data[0].cpu().numpy()
    binary = (mask > 0.5).astype(np.uint8)
    background_mask = 1 - binary
    kernel = np.ones((15, 15), np.uint8)
    dilated = cv2.dilate(background_mask, kernel, iterations=1)
    inpaint_mask = (dilated * 255).astype(np.uint8)

    mask_pil = cv2_to_pil(inpaint_mask).resize((1024, 1024)).convert("L")
    img_pil = image.resize((1024, 1024)).convert("RGB")

    result = inpaint_pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or "",
        image=img_pil,
        mask_image=mask_pil,
        guidance_scale=10,
        num_inference_steps=40
    ).images[0]

    return result

# ---- Gradio interface ----
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Remove & Replace Background")
    gr.Markdown("Upload a headshot, and describe the desired new background.")

    with gr.Row():
        input_img = gr.Image(type="pil", label="Upload Image")
        output_img = gr.Image(type="pil", label="Result")

    with gr.Row():
        prompt = gr.Textbox(
            label="New Background Prompt",
            value="modern open-plan office, soft natural light, minimalistic decor"
        )
        neg_prompt = gr.Textbox(
            label="Negative Prompt",
            value="cartoon, fantasy, dark lighting, painting, anime"
        )

    error_box = gr.Markdown()

    def safe_run(img, prompt, neg_prompt):
        try:
            result = run_background_removal_and_inpaint(img, prompt, neg_prompt)
            return result, ""
        except Exception as e:
            print(f"[ERROR] {type(e).__name__}: {e}")
            return None, f"**❌ Error:** {type(e).__name__}: {e}"

    run_btn = gr.Button("Run Background Inpaint")
    run_btn.click(
        fn=safe_run,
        inputs=[input_img, prompt, neg_prompt],
        outputs=[output_img, error_box]
    )

demo.launch(debug=True)