# streamlit_app.py — Fit Studio AI (Fashion & Apparel) v0.6.5 UI
# Fix: "Video (loop)" duplicate — single placeholder (st.empty) updated in-place.

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

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

API_BASE = (
    _env("FITSTUDIO_API")
    or _env("AI_LIGHTBOX_API")
    or _env("LUXFIT_API")
).rstrip("/") if (
    _env("FITSTUDIO_API") or _env("AI_LIGHTBOX_API") or _env("LUXFIT_API")
) else ""

HF_TOKEN = _env("FITSTUDIO_TOKEN") or _env("AI_LIGHTBOX_TOKEN") or _env("LUXFIT_TOKEN")
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}

# === Tile geometry (strict 9:16) ===
TILE_H = int(os.getenv("TILE_H", "640"))
TILE_W = max(200, int(round(TILE_H * 9 / 16)))
PREVIEW_BG = (255, 255, 255)

GIF_MAX_W = int(os.getenv("GIF_MAX_W", "512"))
GIF_FPS   = int(os.getenv("GIF_FPS", "12"))
GIF_MAX_FRAMES = GIF_FPS * 12

# ================= Session State =================
defaults = {"results": None, "job_counter": 0, "duration": "6", "rand": uuid.uuid4().hex[:8]}
for k, v in defaults.items():
    if k not in st.session_state:
        st.session_state[k] = v

# ================= CSS =================
st.markdown(f"""
<style>
h4.tile-title {{ margin: 0 0 8px 0; font-weight: 600; }}
div.tile-box {{
  width: {TILE_W}px; height: {TILE_H}px;
  display: flex; align-items: center; justify-content: center;
  background: #ffffff; border-radius: 12px; overflow: hidden;
  box-shadow: 0 1px 6px rgba(0,0,0,.08);
  margin: 0 0 6px 0;
}}
img.tile-media, video.tile-media {{
  width: 100%; height: 100%; object-fit: contain; background: #ffffff;
  display: block;
}}
.small-cap {{ font-size: 12px; color: #6b7280; margin: 4px 0 8px 2px; }}
</style>
""", unsafe_allow_html=True)

# ================= Helpers =================
def _needs_auth(url: str) -> bool:
    return (API_BASE and (url.startswith(API_BASE) or url.startswith("outputs/")))

@st.cache_data(ttl=300, show_spinner=False)
def fetch_bytes(url_or_path: str | None):
    if not url_or_path:
        return None
    try:
        p = Path(url_or_path)
        if p.is_file():
            return p.read_bytes()
        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

@st.cache_data(ttl=300, show_spinner=False)
def image_to_9_16_canvas(img_bytes: bytes, w: int = TILE_W, h: int = TILE_H, bg_rgb=PREVIEW_BG) -> bytes:
    im = Image.open(io.BytesIO(img_bytes)).convert("RGBA")
    iw, ih = im.size
    scale = min(w / iw, h / ih)
    nw, nh = max(1, int(iw*scale)), max(1, int(ih*scale))
    im_resized = im.resize((nw, nh), Image.LANCZOS)
    canvas = Image.new("RGBA", (w, h), (*bg_rgb, 255))
    off = ((w - nw)//2, (h - nh)//2)
    canvas.paste(im_resized, off, im_resized)
    out = io.BytesIO(); canvas.save(out, format="PNG"); return out.getvalue()

@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"

def unique_key(prefix: str) -> str:
    return f"{prefix}_{st.session_state['job_counter']}_{st.session_state['rand']}_{uuid.uuid4().hex[:6]}"

def _b64(data: bytes, mime: str) -> str:
    return f"data:{mime};base64," + base64.b64encode(data).decode("ascii")

def show_image_tile(col, title: str, src: str|bytes|None, filename_stub="image", key_prefix=""):
    col.markdown(f"<h4 class='tile-title'>{title}</h4>", unsafe_allow_html=True)
    if not src:
        placeholder = Image.new("RGBA", (TILE_W, TILE_H), (240,240,240,255))
        buf = io.BytesIO(); placeholder.save(buf, "PNG")
        html = f"<div class='tile-box'><img class='tile-media' src='{_b64(buf.getvalue(),'image/png')}'/></div>"
        col.markdown(html, unsafe_allow_html=True)
        col.markdown("<div class='small-cap'>—</div>", unsafe_allow_html=True)
        return None, None

    b = src if isinstance(src, (bytes, bytearray)) else fetch_bytes(src)
    if not b:
        col.warning("Image could not be loaded"); return None, None

    tile_png = image_to_9_16_canvas(b, w=TILE_W, h=TILE_H, bg_rgb=PREVIEW_BG)
    html = f"<div class='tile-box'><img class='tile-media' src='{_b64(tile_png,'image/png')}'/></div>"
    col.markdown(html, unsafe_allow_html=True)

    sz = image_size_from_bytes(b)
    if sz:
        col.markdown(f"<div class='small-cap'>Source: {sz[0]}×{sz[1]} px</div>", unsafe_allow_html=True)

    ext = infer_ext(b)
    mime = "image/jpeg" if ext in [".jpg",".jpeg"] else ("image/png" if ext==".png" else "image/webp")
    col.download_button(
        "Download (full quality)",
        data=b,
        file_name=f"{filename_stub}{ext}",
        mime=mime,
        key=unique_key(f"{key_prefix}_{filename_stub}_dl")
    )
    return src, b

def _video_html_src(src: str) -> str:
    return f"<div class='tile-box'><video class='tile-media' src='{src}' autoplay muted loop playsinline controls></video></div>"

# --- Video rendering in a single placeholder (prevents duplicate heading) ---
def render_video_placeholder(slot, video_url: str|None, key_prefix=""):
    with slot.container():
        st.markdown("<h4 class='tile-title'>Video (loop)</h4>", unsafe_allow_html=True)
        if not video_url:
            st.markdown("<div class='small-cap'>Not generated yet.</div>", unsafe_allow_html=True)
            return None
        vb = fetch_bytes(video_url)
        vsrc = _b64(vb, "video/mp4") if vb else video_url
        st.markdown(_video_html_src(vsrc), unsafe_allow_html=True)
        if vb:
            st.download_button(
                "Download video (MP4)",
                data=vb,
                file_name="fitstudio_video.mp4",
                mime="video/mp4",
                key=unique_key(f"{key_prefix}_video_dl")
            )
        return vb

# ---- MP4 → GIF ----
@st.cache_data(show_spinner=False)
def mp4_to_gif_bytes(mp4_bytes: bytes, target_w: int = GIF_MAX_W, fps: int = GIF_FPS, max_frames: int = GIF_MAX_FRAMES) -> bytes:
    try:
        import imageio.v3 as iio
        from PIL import Image
    except Exception as e:
        raise RuntimeError("imageio.v3 + PIL required for GIF conversion") from e

    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as f:
        f.write(mp4_bytes); mp4_path = f.name

    try:
        import imageio.v3 as iio2  # same pkg; just to be explicit
        meta = {}
        try: meta = iio2.immeta(mp4_path)
        except Exception: meta = {}
        src_fps = meta.get("fps", fps)
        step = max(1, int(round(src_fps / fps))) if isinstance(src_fps, (int, float)) and src_fps > 0 else 1
    except Exception:
        step = 1

    frames = []
    try:
        for idx, frame in enumerate(iio.imiter(mp4_path)):
            if idx % step != 0: continue
            im = Image.fromarray(frame)
            if im.width > target_w:
                new_h = max(1, int(im.height * target_w / im.width))
                im = im.resize((target_w, new_h), Image.LANCZOS)
            frames.append(im.convert("P", palette=Image.ADAPTIVE))
            if len(frames) >= max_frames: break
    except Exception as e:
        raise RuntimeError("Failed to decode frames for GIF. ffmpeg/pyav may be missing.") from e
    finally:
        try: Path(mp4_path).unlink(missing_ok=True)
        except Exception: pass

    if not frames: raise RuntimeError("No frames decoded for GIF.")
    out = io.BytesIO()
    frames[0].save(out, format="GIF", save_all=True, append_images=frames[1:],
                   loop=0, duration=max(10, int(1000 / fps)), disposal=2)
    return out.getvalue()

def make_zip(named_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()

def backend_ok() -> bool:
    if not API_BASE: return False
    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_image_edit(data: dict, files_payload):
    r = requests.post(f"{API_BASE}/v1/image/edit", 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_chain(data: dict, files_payload):
    payload = {**data, "to_video": "false"}
    r = requests.post(f"{API_BASE}/v1/tryon/chain", data=payload, 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, gender="auto", age_group="auto"):
    payload = {"image_url": image_url, "gender": gender, "age_group": age_group, "frame_aspect": "9:16"}
    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()

# ================= Page =================
st.set_page_config(page_title="Fit Studio AI — by Vahit Feryad", layout="wide", page_icon="🧵")

left, right = st.columns([0.7, 0.3], vertical_alignment="center")
with left:
    st.title("🧵 Fit Studio AI — Fashion & Apparel")
    st.caption("Built by **Vahit Feryad** — Demo for Groove Jones")
with right:
    badge = "🟢 Online" if backend_ok() else "🔴 Offline"
    st.metric("Backend", badge)

st.markdown("**9:16** commercial stills (VTO) → optional runway-style video (looping).  \nNo persistence: outputs are remote URLs only. Gender & Age controls steer prompts server-side.")

if not API_BASE:
    st.error("API_BASE not set. Define env var FITSTUDIO_API (or AI_LIGHTBOX_API / LUXFIT_API).")
    st.stop()

# ================= Controls =================
c0, c_gender, c_age, c2, c3, c4 = st.columns([1.0, 1.0, 1.2, 1.8, 1.4, 1.6])
with c0:       category = st.selectbox("Category", ["general","underwear"], index=0, key="category")
with c_gender: gender   = st.selectbox("Gender", ["auto","female","male","unisex"], index=0, key="gender")
with c_age:    age_group= st.selectbox("Age group", ["auto","teen","young_adult","adult","mature"], index=0, key="age_group")
with c2:       file_list= st.file_uploader("Garment image (upload)", type=["jpg","jpeg","png","webp"], accept_multiple_files=True, key="file_upl")
with c3:       image_url= st.text_input("or Garment URL", key="image_url")
with c4:
    num_images    = st.selectbox("Number of looks", [1,2,3,4], index=0, key="num_images")
    custom_prompt = st.text_area("Custom prompt (blank → default)", value="", height=90, key="custom_prompt")

bA, bB = st.columns([1.0, 1.0])
with bA: st.info(f"API: {API_BASE}", icon="🔌")
with bB:
    run_image = st.button("Run (Image only)", type="primary", use_container_width=True, key="run_img_btn")
    run_chain = st.button("Chain (Image → prep Video)", use_container_width=True, key="run_chain_btn")

# ================= Demo preview =================
st.markdown("---"); st.subheader("Demo preview")
dp_in, dp_sp, dp_v = st.columns([1,1,1])
show_image_tile(dp_in, "Input", fetch_bytes("input.jpg"), filename_stub="input", key_prefix="demo_in")
show_image_tile(dp_sp, "Sample Output", fetch_bytes("tryon.jpg"), filename_stub="tryon", key_prefix="demo_out")
# keep separate label to avoid confusion with "Video (loop)"
demo_vbytes = fetch_bytes("fitstudio_video.mp4")
if demo_vbytes:
    dp_v.markdown("<h4 class='tile-title'>Sample Video (loop)</h4>", unsafe_allow_html=True)
    dp_v.markdown(f"<div class='tile-box'><video class='tile-media' src='{_b64(demo_vbytes,'video/mp4')}' autoplay muted loop playsinline controls></video></div>", unsafe_allow_html=True)

# ================= Input preview =================
st.markdown("---")
g1, _, _, _ = st.columns([1,1,1,1])
input_preview = None
if file_list:
    try: input_preview = file_list[0].getvalue()
    except Exception: input_preview = None
elif image_url:
    input_preview = fetch_bytes(image_url)
show_image_tile(g1, "Input (preview)", input_preview, filename_stub="preview", key_prefix="preview")

# ================ Run helpers ================
def _build_files_payload():
    files_payload = []
    if file_list:
        for f in file_list:
            files_payload.append(("files", (f.name, f.getvalue(), f.type or "image/jpeg")))
    return files_payload

def _common_form_data():
    data = {
        "category": st.session_state["category"],
        "gender": st.session_state["gender"],
        "age_group": st.session_state["age_group"],
        "frame_aspect": "9:16",
        "num_images": str(st.session_state["num_images"]),
    }
    if st.session_state.get("custom_prompt","").strip():
        data["prompt"] = st.session_state["custom_prompt"].strip()
    if image_url:
        data["image_urls"] = image_url.strip()
    return data

if run_image or run_chain:
    if not backend_ok():
        st.error("Backend not reachable. Check API_BASE / token.")
    elif not (file_list or image_url):
        st.error("Provide at least one garment image (upload or URL).")
    else:
        try:
            files_payload = _build_files_payload()
            form = _common_form_data()
            with st.spinner("Running…"):
                if run_image:
                    out = post_image_edit(form, files_payload)
                    image_json = out
                else:
                    out = post_chain(form, files_payload)   # to_video=false
                    image_json = out.get("image_step") or out

            res_block = (image_json.get("result") or {})
            imgs_norm = res_block.get("images_9_16") or []
            imgs_raw  = res_block.get("images") or []

            urls_norm, urls_raw = [], []
            for it in imgs_norm:
                u = it.get("url"); 
                if u: urls_norm.append(u)
            for it in imgs_raw:
                u = it.get("url"); 
                if u: urls_raw.append(u)

            st.session_state["job_counter"] += 1
            st.session_state["results"] = {
                "job_id": (image_json.get("job_id") or out.get("job_id")),
                "schema_version": (image_json.get("schema_version") or out.get("schema_version")),
                "image_urls_norm": urls_norm,
                "image_urls_raw": urls_raw,
                "video_url": None,
            }

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

# ===== Helper: dual download (normalized + native) =====
def show_dual_image_tile(col, title: str, norm_src: str|bytes|None, native_src: str|bytes|None,
                         filename_stub="image", key_prefix=""):
    preview_src = norm_src or native_src
    _, preview_bytes = show_image_tile(col, title, preview_src, filename_stub=filename_stub, key_prefix=key_prefix)

    norm_bytes = None; native_bytes = None
    if isinstance(norm_src, (bytes, bytearray)): norm_bytes = bytes(norm_src)
    elif isinstance(norm_src, str): norm_bytes = fetch_bytes(norm_src)
    if isinstance(native_src, (bytes, bytearray)): native_bytes = bytes(native_src)
    elif isinstance(native_src, str): native_bytes = fetch_bytes(native_src)

    if norm_bytes:
        extn = infer_ext(norm_bytes)
        mime = "image/jpeg" if extn in [".jpg",".jpeg"] else ("image/png" if extn==".png" else "image/webp")
        col.download_button(
            "Download 9:16 (normalized)",
            data=norm_bytes,
            file_name=f"{filename_stub}_9x16{extn}",
            mime=mime,
            key=unique_key(f"{key_prefix}_{filename_stub}_norm_dl")
        )
    if native_bytes:
        ext = infer_ext(native_bytes)
        mime = "image/jpeg" if ext in [".jpg",".jpeg"] else ("image/png" if ext==".png" else "image/webp")
        col.download_button(
            "Download native",
            data=native_bytes,
            file_name=f"{filename_stub}_native{ext}",
            mime=mime,
            key=unique_key(f"{key_prefix}_{filename_stub}_raw_dl")
        )
    return preview_bytes, norm_bytes, native_bytes

# ================ Results ================
res = st.session_state.get("results") or {}
if res:
    key_prefix = f"job{st.session_state['job_counter']}"
    img_urls_norm = res.get("image_urls_norm") or []
    img_urls_raw  = res.get("image_urls_raw") or []
    video_url = res.get("video_url")

    st.markdown("### Results")
    c1, c2, c3, c4 = st.columns([1,1,1,1])
    named = []

    def pair_at(i):
        norm = img_urls_norm[i] if i < len(img_urls_norm) else None
        raw  = img_urls_raw[i]  if i < len(img_urls_raw)  else None
        return norm, raw

    for idx, col in enumerate([c1, c2, c3]):
        norm_u, raw_u = pair_at(idx)
        if norm_u or raw_u:
            title = "Look" if idx == 0 else f"Look #{idx+1}"
            _, b_norm, b_raw = show_dual_image_tile(
                col, title, norm_u, raw_u, filename_stub=f"look_{idx+1}", key_prefix=f"{key_prefix}_{idx}"
            )
            if b_norm: named.append((f"look_{idx+1}_9x16{infer_ext(b_norm)}", b_norm))
            if b_raw:  named.append((f"look_{idx+1}_native{infer_ext(b_raw)}", b_raw))
        else:
            show_image_tile(col, f"Look #{idx+1}", None, key_prefix=f"{key_prefix}_{idx}")

    # --- Single, persistent video slot ---
    video_slot = c4.empty()
    vbytes = render_video_placeholder(video_slot, video_url, key_prefix=key_prefix)

    # Controls
    st.markdown("---")
    colv1, colv2 = st.columns([1.2, 1.2])
    with colv1:
        gen_video = st.checkbox("Generate video from first look", value=False, key="video_toggle")
    with colv2:
        go_video = st.button("Create Video", use_container_width=True, key="video_btn")

    if gen_video and go_video:
        src_for_video = (img_urls_norm[0] if img_urls_norm else (img_urls_raw[0] if img_urls_raw else None))
        if not src_for_video:
            st.warning("No source image to create video.")
        else:
            with st.spinner("Generating video…"):
                v = post_video(src_for_video, gender=st.session_state.get("gender","auto"),
                               age_group=st.session_state.get("age_group","auto"))
            vurl = (v.get("result") or {}).get("video", {}).get("url")
            if vurl:
                st.session_state["results"]["video_url"] = vurl
                # Update same placeholder in-place → no duplicate heading
                vbytes = render_video_placeholder(video_slot, vurl, key_prefix=key_prefix)
            else:
                st.info("Video URL not returned.")

    # ZIP + GIF
    if named or vbytes:
        if vbytes:
            named.append(("video.mp4", vbytes))
        zip_bytes = make_zip(named)
        st.download_button(
            "Download All (ZIP)",
            data=zip_bytes,
            file_name="fitstudio_outputs.zip",
            mime="application/zip",
            key=unique_key(f"{key_prefix}_zip_dl")
        )
        if vbytes:
            try:
                with st.spinner("Preparing GIF…"):
                    gif_bytes = mp4_to_gif_bytes(vbytes, target_w=GIF_MAX_W, fps=GIF_FPS, max_frames=GIF_MAX_FRAMES)
                st.download_button(
                    f"Download GIF ({GIF_MAX_W}px, {GIF_FPS}fps)",
                    data=gif_bytes,
                    file_name="fitstudio_video.gif",
                    mime="image/gif",
                    key=unique_key(f"{key_prefix}_gif_dl")
                )
            except Exception as e:
                st.info(f"GIF conversion not available: {e}. Try installing ffmpeg / imageio-ffmpeg.")

# ================= Sidebar =================
with st.sidebar:
    st.caption(f"Resolved API_BASE: {API_BASE}")
    st.caption("Auth header: " + ("ON" if HF_TOKEN else "OFF"))
    st.markdown("---")
    st.markdown("**About this demo**")
    st.caption("Production-style, no local persistence. Built to show scalable image→video chaining with guardrails for consistent 9:16 outputs.")
    if st.button("Ping /health", key=unique_key("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", key=unique_key("clear_cache")):
        st.cache_data.clear(); st.success("Cache cleared.")
    if st.button("Clear results", key=unique_key("clear_results")):
        st.session_state["results"] = None; st.success("Results cleared.")