# streamlit_app.py — AI LightBox · Beauty (v0.3 UI)
# Packshot → Padding (square/canvas) → Try-on (lips/eyes/face/nails/hair/brow/body/lifestyle) · SR · (optional) Video

import os, io, zipfile, requests, streamlit as st
from PIL import Image
from pathlib import Path

# ================= Env & API (robust) =================
def _env(k):
    return (os.getenv(k) or "").strip().strip("'\"")

# Desteklenen env anahtarları (öncelik sırası)
API_BASE = (
    _env("AI_BEAUTYBOX_API") or
    _env("AI_LIGHTBOX_API")  or
    _env("BeautyBoxAI_API")  or
    _env("AI_BEAUTY_API")
).rstrip("/")

HF_TOKEN = (
    _env("AI_BEAUTYBOX_TOKEN") or
    _env("AI_LIGHTBOX_TOKEN")  or
    _env("BeautyBoxAI_TOKEN")  or
    _env("AI_BEAUTY_TOKEN")
)

if not API_BASE:
    st.error("API_BASE bulunamadı. Settings > Variables: AI_BEAUTYBOX_API veya AI_LIGHTBOX_API tanımlayın.")
    st.stop()

HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
PREVIEW_SIZE = int(os.getenv("PREVIEW_SIZE", "512"))

# ================= Session State (ilk yükte) =================
defaults = {
    "results": None,
    "job_counter": 0,
    "duration": "6",
}
for k, v in defaults.items():
    if k not in st.session_state:
        st.session_state[k] = v

# ================= Kategori/Region eşleme =================
LIPS = {"lipstick","lipgloss","lipliner"}
EYES = {"eyeshadow","eyeliner","mascara","brow","false_lashes"}
FACE = {"foundation","concealer","blush","bronzer","contour","highlighter","body_makeup"}
HANDS = {"nail_polish"}
LIFESTYLE = {"skincare","perfume","tools","set","haircare","shampoo","conditioner","body_lotion","body_care"}
HAIR = {"hair_color"}

REGIONS = ["auto","lips","eyes","face","nails","hand","hair","brow","body","lifestyle"]

def tryon_label(category: str, region_override: str, sr: bool) -> str:
    # Region override önce gelir
    r = (region_override or "auto").lower()
    if r != "auto":
        base = {
            "lips":"On-Lips","eyes":"On-Eyes","face":"On-Face","nails":"On-Nails",
            "hand":"On-Hand","hair":"On-Hair","brow":"On-Brows","body":"On-Body",
            "lifestyle":"Lifestyle"
        }.get(r, "Try-on")
        return f"{base} (SR)" if sr else base

    c = (category or "auto").lower()
    if c in LIPS: base = "On-Lips"
    elif c in EYES: base = "On-Eyes"
    elif c in HANDS: base = "On-Nails"
    elif c in HAIR: base = "On-Hair"
    elif c in FACE: base = "On-Face"
    elif c in LIFESTYLE: base = "Lifestyle"
    else: base = "Try-on"
    return f"{base} (SR)" if sr else base

# ================= Yardımcılar =================
def _needs_auth(url: str) -> bool:
    return url.startswith(API_BASE) or url.startswith("outputs/")

@st.cache_data(ttl=300, show_spinner=False)
def fetch_bytes(url_or_path: str):
    """URL, outputs/ yolu ya da mevcut klasördeki yerel dosyayı okuyabilir."""
    if not url_or_path:
        return None
    try:
        # 1) Mevcut klasörde yerel dosya?
        p = Path(url_or_path)
        if p.is_file():
            return p.read_bytes()
        # 2) outputs/ için API_BASE ön ekini uygula
        url = f"{API_BASE}/{url_or_path}" if url_or_path.startswith("outputs/") else url_or_path
        headers = (HEADERS if (_needs_auth(url) and HF_TOKEN) else None)
        r = requests.get(url, headers=headers, timeout=180)
        r.raise_for_status()
        return r.content
    except Exception:
        return None

def backend_ok() -> bool:
    try:
        r = requests.get(f"{API_BASE}/health", headers=HEADERS if HF_TOKEN else None, timeout=10)
        r.raise_for_status()
        return True
    except Exception:
        return False

def post_chain(data: dict, files_payload):
    data = {**data, "to_video": "false"}
    r = requests.post(f"{API_BASE}/v1/tryon/chain", data=data, files=files_payload or None,
                      headers=HEADERS if HF_TOKEN else None, timeout=600)
    r.raise_for_status()
    return r.json()

def post_video(image_url: str, duration="6", resolution="768P", prompt_optimizer=False):
    payload = {
        "image_url": image_url,
        "duration": duration,
        "resolution": resolution,
        "prompt_optimizer": "true" if prompt_optimizer else "false",
    }
    r = requests.post(f"{API_BASE}/v1/video/from-image", data=payload,
                      headers=HEADERS if HF_TOKEN else None, timeout=600)
    r.raise_for_status()
    return r.json()

@st.cache_data(ttl=300, show_spinner=False)
def image_size_from_bytes(b: bytes):
    try:
        im = Image.open(io.BytesIO(b))
        return im.size
    except Exception:
        return None

def infer_ext(data: bytes) -> str:
    try:
        im = Image.open(io.BytesIO(data))
        fmt = (im.format or "JPEG").lower()
        return "." + {"jpeg": "jpg", "jpg": "jpg", "png": "png", "webp": "webp"}.get(fmt, "jpg")
    except Exception:
        return ".jpg"

@st.cache_data(ttl=300, show_spinner=False)
def square_preview_bytes(img_bytes: bytes, size: int = PREVIEW_SIZE, bg_rgb=(255, 255, 255)) -> bytes:
    im = Image.open(io.BytesIO(img_bytes)).convert("RGBA")
    w, h = im.size
    scale = min(size / w, size / h)
    new_w, new_h = max(1, int(w * scale)), max(1, int(h * scale))
    im_resized = im.resize((new_w, new_h), Image.LANCZOS)
    canvas = Image.new("RGBA", (size, size), (*bg_rgb, 255))
    off = ((size - new_w) // 2, (size - new_h) // 2)
    canvas.paste(im_resized, off, im_resized)
    buf = io.BytesIO(); canvas.save(buf, format="PNG"); return buf.getvalue()

def clamp(v: float, lo: float, hi: float) -> float:
    try:
        v = float(v)
    except Exception:
        v = (lo + hi) / 2.0
    return max(lo, min(hi, v))

def first_present(d: dict, keys: list, default=None):
    for k in keys:
        v = d.get(k)
        if v:
            return v
    return default

def show_tile(col, title, url: str | None, bg_rgb=(255,255,255), filename_stub="image", key_prefix=""):
    col.subheader(title)
    if not url:
        placeholder = Image.new("RGBA", (PREVIEW_SIZE, PREVIEW_SIZE), (0, 0, 0, 0))
        col.image(placeholder, use_container_width=True)
        col.caption("—")
        return None, None

    b = fetch_bytes(url)
    if not b:
        col.warning("Görsel yüklenemedi")
        return None, None

    pv = square_preview_bytes(b, size=PREVIEW_SIZE, bg_rgb=bg_rgb)
    col.image(pv, use_container_width=True)
    sz = image_size_from_bytes(b)
    if sz:
        col.caption(f"Gerçek çözünürlük: {sz[0]}×{sz[1]} px")

    ext = infer_ext(b)
    mime = "image/jpeg" if ext.lower() in [".jpg", ".jpeg"] else ("image/png" if ext.lower()==".png" else "image/webp")
    col.download_button(
        "İndir (tam kalite)",
        data=b,
        file_name=f"{filename_stub}{ext}",
        mime=mime,
        key=f"{key_prefix}_{filename_stub}_dl"
    )
    return url, b

def make_zip(named_bytes: list[tuple[str, bytes]]) -> bytes:
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as z:
        for fname, b in named_bytes:
            if b:
                z.writestr(fname, b)
    buf.seek(0)
    return buf.read()

# ================= Sayfa =================
st.set_page_config(page_title="AI LightBox · Beauty", layout="wide", page_icon="💄")
st.title("💄 AI LightBox · Beauty")
st.caption("Packshot → Padding (canvas) → Try-on (lips/eyes/face/nails/hair/brow/body/lifestyle) · Super Resolution · (optional) Video")

# ================= Üst Kontrol Şeridi =================
ok = backend_ok()
col1, col2, col3, col4, col5, col6, col7 = st.columns([1.8, 1.6, 1.4, 1.3, 0.9, 1.1, 0.9])

with col1:
    file_list = st.file_uploader("Product image (upload)", type=["jpg","jpeg","png","webp"], accept_multiple_files=True, key="file_upl")
with col2:
    image_url = st.text_input("or Product URL", key="image_url")
with col3:
    category = st.selectbox(
        "Category",
        [
            # Lips
            "auto","lipstick","lipgloss","lipliner",
            # Eyes
            "eyeshadow","eyeliner","mascara","brow","false_lashes",
            # Face/complexion
            "foundation","concealer","blush","bronzer","contour","highlighter","body_makeup",
            # Nails
            "nail_polish",
            # Lifestyle
            "skincare","perfume","tools","set",
            # Hair/Body care
            "hair_color","haircare","shampoo","conditioner","body_lotion","body_care",
        ],
        index=0, key="category"
    )
with col4:
    padding_ratio_val = st.number_input(
        "Content scale (0.30–0.95)",
        min_value=0.30, max_value=0.95, value=0.50, step=0.01,
        help="Kare tuvalde ürünün uzun kenarı = tuval × bu oran.",
        key="padding_ratio"
    )
with col5:
    upscale = st.checkbox("Super Resolution", False, key="upscale")
with col6:
    upscale_stage = st.selectbox("SR stage", ["both","final","packshot"], index=0, help="Default: both", key="upscale_stage")
with col7:
    upscale_factor = st.selectbox("SR factor", ["2","4"], index=0, key="upscale_factor")

# == Tuval & Try-on kontrolleri ==
t1, t2, t3, t4 = st.columns([1.15, 1.15, 1.2, 1.5])
with t1:
    canvas_policy = st.selectbox("Canvas", ["fixed_1200","match_long_edge","keep_input"], index=0, key="canvas_policy")
with t2:
    canvas_size = st.number_input("Canvas size (px)", min_value=256, max_value=8192, value=1200, step=64, key="canvas_size")
with t3:
    apply_region = st.selectbox("Apply region", REGIONS, index=0, key="apply_region")
with t4:
    shade_hex = st.text_input("Shade HEX (optional, e.g. #C6A27A)", value="", key="shade_hex")

b1, b2, b3 = st.columns([1.6, 1.6, 1.0])
with b1:
    shade_swatch_url = st.text_input("Shade swatch URL (optional)", value="", key="shade_swatch_url")
with b2:
    model_ref_url = st.text_input("Model image URL (face/hand/eyes/hair/body — optional)", key="face_url")
with b3:
    identity_lock = st.checkbox("Identity lock", True, key="identity_lock")

# ============ DEMO PREVIEW (mevcut klasörden input.jpg & tryon.jpg) ============
st.markdown("---")
st.subheader("Demo preview (çalıştırmadan önce örnek)")

d_in, d_pack, d_pad, d_try = st.columns(4)

# Input preview
in_bytes = fetch_bytes("input.jpg")
if in_bytes:
    d_in.subheader("Input")
    d_in.image(square_preview_bytes(in_bytes, size=PREVIEW_SIZE, bg_rgb=(255,255,255)), use_container_width=True)
    sz = image_size_from_bytes(in_bytes)
    if sz:
        d_in.caption(f"Input resolution: {sz[0]}×{sz[1]} px")
    d_in.download_button(
        "İndir (tam kalite)", data=in_bytes,
        file_name=f"input{infer_ext(in_bytes)}",
        mime="image/jpeg", key="demo_input_dl"
    )
else:
    show_tile(d_in, "Input", None)

# Packshot & Padded columnları bu demoda boş
show_tile(d_pack, "Packshot", None)
show_tile(d_pad, "Padded", None)

# On-Hair / final try-on preview
try_bytes = fetch_bytes("tryon.jpg")
if try_bytes:
    d_try.subheader("On-Hair")
    d_try.image(square_preview_bytes(try_bytes, size=PREVIEW_SIZE, bg_rgb=(255,255,255)), use_container_width=True)
    sz2 = image_size_from_bytes(try_bytes)
    if sz2:
        d_try.caption(f"Gerçek çözünürlük: {sz2[0]}×{sz2[1]} px")
    d_try.download_button(
        "İndir (tam kalite)", data=try_bytes,
        file_name=f"tryon{infer_ext(try_bytes)}",
        mime="image/jpeg", key="demo_tryon_dl"
    )
else:
    show_tile(d_try, "On-Hair", None)

# ================= Alt kontrol şeridi =================
cA, cB = st.columns([1.0, 1.0])
with cA:
    st.info(f"Backend: {'Online' if ok else 'Offline'}", icon="🔌")
with cB:
    run = st.button("Run", type="primary", use_container_width=True, key="run_btn")

# ================= Girdi Önizleme =================
c_in, c_pack, c_pad, c_try = st.columns(4)
white_bg = (255, 255, 255)

input_preview_bytes = None
if file_list:
    try:
        input_preview_bytes = file_list[0].getvalue()
    except Exception:
        input_preview_bytes = None
elif image_url:
    input_preview_bytes = fetch_bytes(image_url)

if input_preview_bytes:
    c_in.subheader("Input")
    c_in.image(square_preview_bytes(input_preview_bytes, size=PREVIEW_SIZE, bg_rgb=white_bg), use_container_width=True)
    try:
        sz = Image.open(io.BytesIO(input_preview_bytes)).size
        c_in.caption(f"Input resolution: {sz[0]}×{sz[1]} px")
    except Exception:
        c_in.caption("—")
else:
    show_tile(c_in, "Input", None)

# ================ Koş & Sonuçları Kaydet ================
if run:
    if not ok:
        st.error("Backend erişilemiyor. API_BASE / TOKEN kontrol edin.")
    elif not (file_list or image_url):
        st.error("En az bir product görseli girin (upload veya URL).")
    else:
        try:
            padding_ratio = clamp(padding_ratio_val, 0.30, 0.95)
            data = {
                "category": category,
                "edit_prompt": "",
                "num_images": "1",
                "image_urls": (image_url or ""),
                "mannequin_image_url": (model_ref_url or ""),
                "padding_ratio": str(padding_ratio),
                "identity_lock": "true" if identity_lock else "false",
                # Canvas
                "canvas_policy": canvas_policy,
                "canvas_size": str(int(canvas_size or 1200)),
                # Shade / Region
                "shade_hex": (shade_hex or "").strip() or None,
                "shade_swatch_url": (shade_swatch_url or "").strip() or None,
                "apply_region": apply_region,
                # SR
                "upscale": "true" if upscale else "false",
                "upscale_factor": st.session_state["upscale_factor"],
                "upscale_stage": st.session_state["upscale_stage"],
            }
            files_payload = []
            if file_list:
                for f in file_list:
                    files_payload.append(("files", (f.name, f.getvalue(), f.type or "image/jpeg")))

            with st.spinner("Running chain…"):
                out = post_chain(data, files_payload)

            # ---- Çıktıları topla ----
            packshot_url = None
            try:
                imgs = out["packshot_step"]["result"]["images"]
                if imgs:
                    packshot_url = imgs[0].get("url")
            except Exception:
                pass

            padded_url = None
            try:
                padded_url = out.get("padding_step", {}).get("saved_file")
                if not padded_url:
                    pads = out.get("packshot_step", {}).get("padded_urls") or out.get("padded_urls") or []
                    if pads:
                        padded_url = pads[0]
            except Exception:
                pass

            packshot_sr_list = first_present(out.get("packshot_step", {}), [
                "upscaled_packshot_urls", "upscaled_packshot_files"
            ], default=[]) or []
            packshot_sr_url = packshot_sr_list[0] if packshot_sr_list else None

            placed_url = None
            try:
                imgs = out["image_step"]["result"]["images"]
                if imgs:
                    placed_url = imgs[0].get("url")
            except Exception:
                pass

            final_sr_list = first_present(out.get("image_step", {}), [
                "upscaled_final_urls", "upscaled_final_files"
            ], default=[]) or []
            placed_sr_url = final_sr_list[0] if final_sr_list else None

            st.session_state["job_counter"] += 1
            st.session_state["results"] = {
                "job_id": out.get("job_id"),
                "schema_version": out.get("schema_version"),
                "final_canvas": (
                    out.get("packshot_step", {}).get("final_canvas")
                    or out.get("padding_step", {}).get("final_canvas")
                ),
                "packshot_url": packshot_url,
                "packshot_sr_url": packshot_sr_url,
                "padded_url": padded_url,
                "placed_url": placed_url,
                "placed_sr_url": placed_sr_url,
                "video_url": None,
                "api_base": API_BASE,
                "has_auth_header": bool(HEADERS),
                # UI state snapshot (etiket/isimlendirme için)
                "category": category,
                "apply_region": apply_region,
            }

        except requests.HTTPError as e:
            st.error(f"HTTP {e.response.status_code}")
        except Exception as e:
            st.error(f"Hata: {e}")

# ================ Sonuçları Göster (kalıcı) ================
res = st.session_state.get("results")
vbytes = None
if res:
    key_prefix = f"job{st.session_state['job_counter']}"

    p_title = "Packshot (SR)" if res.get("packshot_sr_url") else "Packshot"
    p_url, p_bytes = show_tile(
        c_pack, p_title, res.get("packshot_sr_url") or res.get("packshot_url"),
        bg_rgb=white_bg, filename_stub="packshot", key_prefix=key_prefix
    )

    pad_url, pad_bytes = show_tile(
        c_pad, "Padded", res.get("padded_url"),
        bg_rgb=white_bg, filename_stub="padded", key_prefix=key_prefix
    )

    shown_final_url = res.get("placed_sr_url") or res.get("placed_url")
    m_title = tryon_label(res.get("category","auto"), res.get("apply_region","auto"), bool(res.get("placed_sr_url")))
    m_url, m_bytes = show_tile(
        c_try, m_title, shown_final_url,
        bg_rgb=white_bg, filename_stub="tryon", key_prefix=key_prefix
    )

    # ---- Video (final görselden) ----
    st.markdown("---")
    colv1, colv2, colv3 = st.columns([1.1, 1.0, 1.2])
    with colv1:
        gen_video = st.checkbox("Generate video from this result", value=False, key="video_toggle")
    with colv2:
        st.selectbox("Video sec", ["6","10"], index=(0 if st.session_state["duration"]=="6" else 1), key="duration_select")
    with colv3:
        go_video = st.button("Create video", use_container_width=True, key="video_btn")

    if gen_video and go_video:
        src_for_video = res.get("placed_sr_url") or res.get("placed_url")
        if not src_for_video:
            st.warning("Video için final görsel yok.")
        else:
            with st.spinner("Generating video…"):
                v = post_video(
                    src_for_video,
                    duration=st.session_state.get("duration_select","6"),
                    resolution="768P",
                    prompt_optimizer=False
                )
            vurl = (v.get("result") or {}).get("video", {}).get("url")
            if vurl:
                st.session_state["results"]["video_url"] = vurl
            else:
                st.info("Video URL gelmedi.")

    vurl = st.session_state["results"].get("video_url")
    if vurl:
        st.subheader("Video")
        st.video(vurl, format="video/mp4")
        vb = fetch_bytes(vurl)
        if vb:
            vbytes = vb
            st.download_button(
                "Download video (MP4)",
                data=vb,
                file_name="ai_lightbox_video.mp4",
                mime="video/mp4",
                key=f"{key_prefix}_video_dl"
            )
        else:
            st.info("Video URL var ama içerik indirilemedi.")

    # ---- ZIP indir ----
    named = []
    if p_bytes: named.append((f"packshot{infer_ext(p_bytes)}", p_bytes))
    if pad_bytes: named.append((f"padded{infer_ext(pad_bytes)}", pad_bytes))
    if m_bytes: named.append((f"tryon{infer_ext(m_bytes)}", m_bytes))
    if vbytes: named.append(("video.mp4", vbytes))
    if named:
        zip_bytes = make_zip(named)
        st.download_button(
            "Download all (ZIP)",
            data=zip_bytes,
            file_name="ai_lightbox_beauty_outputs.zip",
            mime="application/zip",
            key=f"{key_prefix}_zip_dl"
        )

    # Debug
    with st.expander("Debug"):
        st.json({
            "job_id": res.get("job_id"),
            "schema_version": res.get("schema_version"),
            "final_canvas": res.get("final_canvas"),
            "api_base": res.get("api_base"),
            "has_auth_header": res.get("has_auth_header"),
            "category": res.get("category"),
            "apply_region": res.get("apply_region"),
            "packshot_url": res.get("packshot_url"),
            "packshot_sr_url": res.get("packshot_sr_url"),
            "padded_url": res.get("padded_url"),
            "placed_url": res.get("placed_url"),
            "placed_sr_url": res.get("placed_sr_url"),
            "video_url": res.get("video_url"),
        })

# ================= Sidebar (teşhis) =================
with st.sidebar:
    st.caption(f"Resolved API_BASE: {API_BASE}")
    st.caption("Auth header: " + ("ON" if HF_TOKEN else "OFF"))
    if st.button("Ping /health"):
        try:
            r = requests.get(f"{API_BASE}/health",
                             headers=HEADERS if HF_TOKEN else None,
                             timeout=10)
            st.write(r.status_code)
            try:
                st.json(r.json())
            except Exception:
                st.code((r.text or "")[:1000])
        except Exception as e:
            st.error(f"Health error: {e}")
    if st.button("Clear cache"):
        st.cache_data.clear(); st.success("Cache cleared.")
    if st.button("Clear results"):
        st.session_state["results"] = None; st.success("Results cleared.")