import gradio as gr
import numpy as np
from PIL import Image, ImageFilter, ImageOps
from transformers import pipeline
from pathlib import Path
import zipfile
import shutil
import os
import cv2

print("DEPTH AAA ENGINE (CLEAN REBUILD)")

# -------------------------
# CONFIG
# -------------------------
MODEL_ID = "depth-anything/Depth-Anything-V2-Base-hf"

# Use cuda if available, otherwise cpu
device = 0 if os.environ.get("CUDA_VISIBLE_DEVICES", "") != "" else -1

pipe = pipeline(
    task="depth-estimation",
    model=MODEL_ID,
    device=device
)

# -------------------------
# HELPERS
# -------------------------
def normalize_depth(depth_img: Image.Image) -> np.ndarray:
    d = np.array(depth_img).astype(np.float32)
    d_min = d.min()
    d_max = d.max()
    d = (d - d_min) / (d_max - d_min + 1e-6)
    return d

def refine_depth(depth_np: np.ndarray) -> Image.Image:
    img = Image.fromarray((depth_np * 255).astype(np.uint8))

    # leve blur para suavizar banding
    img = img.filter(ImageFilter.GaussianBlur(1.0))

    # melhora contraste sem destruir detalhe
    img = ImageOps.autocontrast(img, cutoff=0.3)

    return img

def colorize_depth(depth_gray: Image.Image) -> Image.Image:
    arr = np.array(depth_gray)
    color = cv2.applyColorMap(arr, cv2.COLORMAP_INFERNO)
    color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)
    return Image.fromarray(color)

def process(files):
    if not files:
        return None, None

    out_dir = Path("depth_output")
    if out_dir.exists():
        shutil.rmtree(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    zip_path = "DEPTH_RESULT.zip"

    preview_items = []

    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for item in files:
            path = item if isinstance(item, str) else getattr(item, "name", item)
            stem = Path(path).stem

            img = Image.open(path).convert("RGB")

            # depth-estimation pipeline returns a dict with depth image
            result = pipe(img)
            depth_raw = result["depth"]

            depth_np = normalize_depth(depth_raw)
            depth_final = refine_depth(depth_np)

            out_path = out_dir / f"{stem}_depth.png"
            depth_final.save(out_path)
            zipf.write(out_path, out_path.name)

            preview_items.append((str(out_path), f"{stem}_depth"))

    # show first result as preview if available
    preview = preview_items[0][0] if preview_items else None
    return zip_path, preview

# -------------------------
# UI
# -------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 DEPTH AAA (CLEAN REBUILD)")

    inp = gr.File(file_count="multiple", type="filepath")
    out_zip = gr.File(label="ZIP")
    preview = gr.Image(label="Preview", type="filepath")

    btn = gr.Button("GERAR DEPTH")
    btn.click(fn=process, inputs=inp, outputs=[out_zip, preview])

demo.launch()