app.py

import os
import math
import random
import tempfile
import gradio as gr
from huggingface_hub import HfApi, hf_hub_download

# ======================================================
# CONFIG
# ======================================================
DEFAULT_REPO = "rahul7star/WAN22-Mar-2026"
FREE_REPO = "rahul7star/ohamlab"
FREE_IMAGE_DIR = "showcase/image"

HF_TOKEN = os.getenv("HF_TOKEN")
APP_PASSWORD = os.getenv("APP_PASSWORD")

if not HF_TOKEN:
    raise RuntimeError("HF_TOKEN not set")
if not APP_PASSWORD:
    raise RuntimeError("APP_PASSWORD not set")

api = HfApi()
PAGE_SIZE = 6

tmp_dir = tempfile.mkdtemp()
images_cache = {}   # (repo, date) -> images

# ======================================================
# FREE PREVIEW
# ======================================================
def load_free_images():
    tmp = tempfile.mkdtemp()
    files = api.list_repo_files(FREE_REPO, token=HF_TOKEN)

    imgs = [
        f for f in files
        if f.startswith(FREE_IMAGE_DIR)
        and f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
    ]

    random.shuffle(imgs)
    imgs = imgs[:5]

    out = []
    for f in imgs:
        try:
            out.append(hf_hub_download(FREE_REPO, f, token=HF_TOKEN, local_dir=tmp))
        except Exception as e:
            print("[FREE IMG ERROR]", e)

    return out


FREE_IMAGES = load_free_images()

# ======================================================
# HELPERS
# ======================================================
def clear_cache():
    images_cache.clear()
    print("[CACHE] cleared all")


def list_date_folders(repo, last_n=20):
    files = api.list_repo_files(repo, token=HF_TOKEN)
    dates = sorted(
        {f.split("/")[0] for f in files if "/" in f},
        reverse=True
    )
    return dates[:last_n]


def download_images(repo, date):
    key = (repo, date)
    if key in images_cache:
        return images_cache[key]

    files = api.list_repo_files(repo, token=HF_TOKEN)
    local = []

    for f in files:
        if not f.startswith(date + "/"):
            continue
        if not f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
            continue

        dst = os.path.join(tmp_dir, repo.replace("/", "_"), f)
        os.makedirs(os.path.dirname(dst), exist_ok=True)

        try:
            img_path = hf_hub_download(
                repo,
                f,
                token=HF_TOKEN,
                local_dir=os.path.dirname(dst),
            )

            # ------------------------------------------
            # NEW: load summary.txt from same folder
            # ------------------------------------------
            summary_text = ""
            summary_file = os.path.join(os.path.dirname(f), "summary.txt")

            if summary_file in files:
                try:
                    summary_path = hf_hub_download(
                        repo,
                        summary_file,
                        token=HF_TOKEN,
                        local_dir=os.path.dirname(dst),
                    )
                    with open(summary_path, "r", encoding="utf-8") as sf:
                        summary_text = sf.read().strip()
                except Exception as e:
                    print("[SUMMARY ERROR]", summary_file, e)

            # Gallery item = (image, caption)
            local.append((img_path, summary_text))

            print("[DOWNLOAD]", repo, f)

        except Exception as e:
            print("[DOWNLOAD ERROR]", f, e)

    images_cache[key] = local
    print(f"[INFO] {repo} {date}: {len(local)} images")
    return local


# ======================================================
# PAGINATION
# ======================================================
def render_page(repo, date, page):
    images = download_images(repo, date)
    total_pages = max(1, math.ceil(len(images) / PAGE_SIZE))

    page = max(0, min(int(page), total_pages - 1))
    start = page * PAGE_SIZE
    end = start + PAGE_SIZE

    return (
        images[start:end],
        page,
        f"Page {page+1} / {total_pages}",
    )


def nav(delta, repo, date, page):
    return render_page(repo, date, page + delta)


def change_date(repo, date):
    clear_cache()
    return render_page(repo, date, 0)


def change_repo(repo):
    clear_cache()
    dates = list_date_folders(repo)
    first = dates[0] if dates else ""
    imgs, page, info = render_page(repo, first, 0)
    return dates, imgs, page, first, info


# ======================================================
# AUTH
# ======================================================
def check_password(pwd):
    if pwd != APP_PASSWORD:
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            [],
            0,
            "",
            "❌ Wrong password",
        )

    dates = list_date_folders(DEFAULT_REPO)
    latest = dates[0]

    imgs, page, info = render_page(DEFAULT_REPO, latest, 0)

    return (
        gr.update(visible=False),
        gr.update(visible=True),
        imgs,
        page,
        latest,
        info,
    )


# ======================================================
# CSS
# ======================================================
css = """
footer, .gradio-footer { display:none!important; }
#ohamlab-footer {
  position:fixed;
  bottom:0;
  width:100%;
  background:#f8f9fb;
  font-size:12px;
  padding:8px;
  border-top:1px solid #e5e7eb;
  text-align:center;
}
"""

# ======================================================
# UI
# ======================================================
with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.Markdown("# 🔐 OhamLab Image Showcase (PRO)")

    # ---------- LOGIN ----------
    with gr.Column(visible=True) as login_box:
        pwd = gr.Textbox(type="password", label="Password")
        login_btn = gr.Button("Unlock")

        gr.Markdown("### 🎁 Free Preview")
        with gr.Row():
            for img in FREE_IMAGES:
                gr.Image(img, show_label=False)

    # ---------- APP ----------
    with gr.Column(visible=False) as app:
        repo_state = gr.State(DEFAULT_REPO)
        date_state = gr.State("")
        page_state = gr.State(0)

        repo_input = gr.Textbox(
            value=DEFAULT_REPO,
            label="Base HuggingFace Repo",
        )

        date_dropdown = gr.Dropdown(label="Select Date")

        manual_date = gr.Textbox(
            label="Or enter date manually (YYYY-MM-DD)"
        )

        gallery = gr.Gallery(columns=3, height="auto")
        page_info = gr.Markdown()

        with gr.Row():
            prev_btn = gr.Button("⬅ Previous")
            next_btn = gr.Button("Next ➡")

    # ---------- EVENTS ----------
    login_btn.click(
        check_password,
        inputs=pwd,
        outputs=[
            login_box,
            app,
            gallery,
            page_state,
            date_state,
            page_info,
        ],
    )

    repo_input.change(
        change_repo,
        inputs=repo_input,
        outputs=[
            date_dropdown,
            gallery,
            page_state,
            date_state,
            page_info,
        ],
    ).then(
        lambda r: r,
        inputs=repo_input,
        outputs=repo_state,
    )

    date_dropdown.change(
        lambda r, d: change_date(r, d),
        inputs=[repo_state, date_dropdown],
        outputs=[gallery, page_state, page_info],
    ).then(
        lambda d: d,
        inputs=date_dropdown,
        outputs=date_state,
    )

    manual_date.submit(
        lambda r, d: change_date(r, d),
        inputs=[repo_state, manual_date],
        outputs=[gallery, page_state, page_info],
    ).then(
        lambda d: d,
        inputs=manual_date,
        outputs=date_state,
    )

    prev_btn.click(
        lambda r, d, p: nav(-1, r, d, p),
        inputs=[repo_state, date_state, page_state],
        outputs=[gallery, page_state, page_info],
    )

    next_btn.click(
        lambda r, d, p: nav(1, r, d, p),
        inputs=[repo_state, date_state, page_state],
        outputs=[gallery, page_state, page_info],
    )

    gr.HTML("<div id='ohamlab-footer'>© OhamLab Copyright</div>")

demo.launch()