import gradio as gr
import cv2
import numpy as np
import os
import uuid
import json
from PIL import Image
from pathlib import Path
import imagehash
import requests
from io import BytesIO
from rembg import remove, new_session

# Lazy global for background-removal session (small model).
_bg_session = None
import mediapipe as mp  # still used for future extensions; not required for detection now

# Paths
ROOT = Path(__file__).resolve().parent
OUTPUT_ROOT = ROOT.parent / "outputs"
OUTPUT_ROOT.mkdir(exist_ok=True)

# Quality thresholds
VAR_LAP_MIN = 60  # focus
BRIGHT_MIN, BRIGHT_MAX = 30, 230
FACE_CONF_MIN = 0.85
CLEAN_BG_COLOR = (128, 128, 128, 255)

# OpenCV DNN face detector
MODEL_DIR = ROOT / "models"
PROTOTXT = MODEL_DIR / "deploy.prototxt"
CAFFEMODEL = MODEL_DIR / "res10_300x300_ssd_iter_140000.caffemodel"
if not PROTOTXT.exists() or not CAFFEMODEL.exists():
    raise FileNotFoundError("Face detector weights not found. Expected deploy.prototxt and res10_300x300_ssd_iter_140000.caffemodel in preprocess/models/")
FACE_NET = cv2.dnn.readNetFromCaffe(str(PROTOTXT), str(CAFFEMODEL))


def pil_to_cv(img: Image.Image):
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)


def cv_to_pil(img: np.ndarray):
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))


def variance_of_laplacian(image):
    return cv2.Laplacian(image, cv2.CV_64F).var()


def basic_qc(pil_img):
    cv_img = pil_to_cv(pil_img)
    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
    focus = variance_of_laplacian(gray)
    mean = gray.mean()
    ok_focus = focus >= VAR_LAP_MIN
    ok_bright = BRIGHT_MIN <= mean <= BRIGHT_MAX
    return {
        "focus": round(float(focus), 2),
        "brightness": round(float(mean), 2),
        "ok_focus": ok_focus,
        "ok_brightness": ok_bright,
    }


def face_detect(pil_img):
    cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    h, w, _ = cv_img.shape
    blob = cv2.dnn.blobFromImage(cv_img, 1.0, (300, 300), (104.0, 177.0, 123.0))
    FACE_NET.setInput(blob)
    detections = FACE_NET.forward()
    boxes = []
    for i in range(detections.shape[2]):
        score = detections[0, 0, i, 2]
        if score < FACE_CONF_MIN:
            continue
        box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
        x1, y1, x2, y2 = box.astype(int)
        boxes.append({
            "x1": int(max(0, x1)),
            "y1": int(max(0, y1)),
            "x2": int(min(w, x2)),
            "y2": int(min(h, y2)),
            "score": float(score),
        })
    return boxes


def crop_with_margin(pil_img, box, margin=0.15, square=True, target=1024):
    w, h = pil_img.size
    x1, y1, x2, y2 = box
    bw, bh = x2 - x1, y2 - y1
    mx, my = int(bw * margin), int(bh * margin)
    x1 = max(0, x1 - mx)
    y1 = max(0, y1 - my)
    x2 = min(w, x2 + mx)
    y2 = min(h, y2 + my)
    if square:
        # make square by padding
        bw, bh = x2 - x1, y2 - y1
        if bw > bh:
            diff = bw - bh
            y1 = max(0, y1 - diff // 2)
            y2 = min(h, y2 + diff - diff // 2)
        else:
            diff = bh - bw
            x1 = max(0, x1 - diff // 2)
            x2 = min(w, x2 + diff - diff // 2)
    cropped = pil_img.crop((x1, y1, x2, y2)).resize((target, target), Image.Resampling.LANCZOS)
    return cropped


def get_bg_session():
    global _bg_session
    if _bg_session is None:
        # u2netp is ~4 MB; default u2net is ~176 MB.
        _bg_session = new_session("u2netp")
    return _bg_session


def segment_person(pil_img):
    # Use lightweight model; pass session explicitly to avoid arg conflicts.
    rgba = remove(pil_img.convert("RGB"), session=get_bg_session())
    return rgba.convert("RGBA")


def face_identity_crop(img: Image.Image, face_box):
    """
    Mask everything outside the face box to neutral gray to reduce LoRA memorizing clothes/jewelry.
    """
    x1, y1, x2, y2 = face_box
    w, h = img.size
    pad_x = int((x2 - x1) * 0.15)
    pad_y = int((y2 - y1) * 0.20)
    fx1 = max(0, x1 - pad_x)
    fy1 = max(0, y1 - pad_y)
    fx2 = min(w, x2 + pad_x)
    fy2 = min(h, y2 + pad_y)
    mask_np = np.zeros((h, w), dtype=np.uint8)
    mask_np[fy1:fy2, fx1:fx2] = 255
    mask = Image.fromarray(mask_np)
    base = Image.new("RGBA", img.size, CLEAN_BG_COLOR)
    identity = Image.composite(img.convert("RGBA"), base, mask)
    return identity.convert("RGB")


def florence_caption(pil_img: Image.Image):
    """
    Optional Florence-2-Large captioning. Set FLORENCE_URL and FLORENCE_KEY env vars.
    """
    url = os.environ.get("FLORENCE_URL")
    key = os.environ.get("FLORENCE_KEY")
    if not url or not key:
        return None
    try:
        buf = BytesIO()
        pil_img.save(buf, format="PNG")
        resp = requests.post(
            url,
            headers={"Authorization": f"Bearer {key}"},
            files={"file": ("image.png", buf.getvalue(), "image/png")},
            timeout=15,
        )
        if resp.ok:
            data = resp.json()
            return data.get("caption") or data.get("text")
    except Exception:
        return None
    return None


def dedupe_hashes(images):
    hashes = []
    unique = []
    for img in images:
        h = imagehash.phash(img)
        dup = False
        for seen in hashes:
            if h - seen <= 4:  # Hamming distance threshold
                dup = True
                break
        if not dup:
            hashes.append(h)
            unique.append(img)
    return unique


def process_images(files):
    if not files:
        return "No images uploaded", None, None

    imgs = [Image.open(f.name if hasattr(f, 'name') else f).convert("RGB") for f in files]

    # Deduplicate
    imgs = dedupe_hashes(imgs)

    processed = []
    manifest = []

    for idx, img in enumerate(imgs):
        qc = basic_qc(img)
        faces = face_detect(img)
        # prefer largest face if multiple
        if faces:
            face = sorted(faces, key=lambda b: b["score"] * ((b["x2"]-b["x1"])*(b["y2"]-b["y1"])), reverse=True)[0]
            crop = crop_with_margin(img, (face["x1"], face["y1"], face["x2"], face["y2"]), margin=0.15, square=True, target=1024)
            caption = florence_caption(crop) or "photo of sks_person, close-up, neutral lighting, straight face"
        else:
            # fallback: center crop
            w, h = img.size
            side = min(w, h)
            x1 = (w - side)//2
            y1 = (h - side)//2
            crop = img.crop((x1, y1, x1+side, y1+side)).resize((1024, 1024), Image.Resampling.LANCZOS)
            caption = florence_caption(crop) or "photo of sks_person"

        # segmentation variant
        cutout = segment_person(crop)
        # flat background
        bg = Image.new("RGBA", cutout.size, CLEAN_BG_COLOR)
        clean = Image.alpha_composite(bg, cutout).convert("RGB")

        # face-neutralized variant to suppress clothes/jewelry
        if faces:
            face_box = (face["x1"], face["y1"], face["x2"], face["y2"])
            identity_only = face_identity_crop(clean, face_box)
        else:
            identity_only = clean

        uid = uuid.uuid4().hex[:8]
        base_dir = OUTPUT_ROOT / "preprocessed"
        base_dir.mkdir(parents=True, exist_ok=True)
        out_clean = base_dir / f"img_{idx}_{uid}_clean.png"
        out_orig = base_dir / f"img_{idx}_{uid}_orig.png"
        out_face = base_dir / f"img_{idx}_{uid}_face.png"

        clean.save(out_clean)
        crop.save(out_orig)
        identity_only.save(out_face)

        manifest.append({
            "image": str(out_clean),
            "identity_masked": str(out_face),
            "caption": caption,
            "focus": qc["focus"],
            "brightness": qc["brightness"],
            "has_face": bool(faces),
        })

        processed.append((crop, clean, identity_only, caption, qc))

    manifest_path = OUTPUT_ROOT / "preprocessed" / "captions.jsonl"
    with open(manifest_path, "w") as f:
        for row in manifest:
            f.write(json.dumps(row) + "\n")

    # Build gallery display
    gallery = []
    for crop, clean, identity_only, cap, qc in processed:
        gallery.append((crop, f"Aligned | {cap}"))
        gallery.append((clean, f"Clean bg | {cap} | focus {qc['focus']} | bright {qc['brightness']}"))
        gallery.append((identity_only, f"Face-only | {cap} | focus {qc['focus']}"))

    summary = f"Saved {len(processed)} images to {manifest_path.parent}. captions.jsonl created."
    return summary, gallery, str(manifest_path)


def build_ui():
    with gr.Blocks(title="LoRA Prep") as demo:
        gr.Markdown("# Image Preprocess for LoRA\nUpload ~8 photos (selfies + body).")
        with gr.Row():
            inp = gr.File(file_count="multiple", file_types=["image"], label="Upload images")
        run_btn = gr.Button("Run Preprocess")
        status = gr.Markdown()
        gallery = gr.Gallery(label="Crops & Cleaned", columns=2, height=600)
        manifest_path = gr.Textbox(label="captions.jsonl path", interactive=False)

        run_btn.click(fn=process_images, inputs=inp, outputs=[status, gallery, manifest_path])
    return demo


if __name__ == "__main__":
    app = build_ui()
    # Expose locally and offer a Gradio share link for easy access if localhost is blocked.
    launch_info = app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True,
        inbrowser=False,
    )
    try:
        share_url = getattr(launch_info, "share_url", None)
        if share_url:
            print(f"GRADIO_SHARE_URL: {share_url}", flush=True)
    except Exception:
        pass