"""ShotCraft — AI Shot Director for e-commerce.
HF x Gradio Build Small Hackathon 2026.
Stage 1: MiniCPM-V-2_6 (8B) | Stage 2: FLUX.1-schnell (12B)
"""
from __future__ import annotations
import io, json, os, tempfile, zipfile
import gradio as gr
from PIL import Image
from schemas import CATEGORIES, STYLE_PRESETS
from director import generate_concepts
from frames import generate_frame, generate_frames, seed_for
from model_runtime import API_URL, BackendError, health

EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")

# Force Gradio's light palette so stock widgets match the light design system.
# Plain <script> in head= runs before the app mounts (launch(js=...) proved unreliable).
FORCE_THEME_HEAD = """
<script>
(function () {
    var url = new URL(window.location.href);
    if (url.searchParams.get('__theme') !== 'light') {
        url.searchParams.set('__theme', 'light');
        window.location.replace(url.href);
    }
})();
</script>
<script>
// Stage rail = primary navigation. The tab pill bar is parked offscreen
// (NOT display:none - Gradio 6 builds tab buttons from measured width, so a
// zero-width tablist renders no buttons at all). Rail clicks forward to the
// real tab buttons; delegated so it survives rail re-renders.
document.addEventListener('click', function (e) {
    var card = e.target.closest('.stage-card');
    if (!card || !card.dataset.target) return;
    var map = { direct: 0, render: 1, export: 1 };
    var idx = map[card.dataset.target];
    if (idx === undefined) return;
    var tabs = document.querySelectorAll('#shotcraft-tabs [role="tab"]');
    if (tabs[idx]) tabs[idx].click();
});
document.addEventListener('keydown', function (e) {
    if (e.key !== 'Enter' && e.key !== ' ') return;
    var card = e.target.closest && e.target.closest('.stage-card');
    if (card) { e.preventDefault(); card.click(); }
});
</script>
"""

# Native theme base: lets Gradio's own variables style internals custom CSS
# can't reach (sliders, checkboxes, focus rings) and loads Inter without a
# CSS @import (no flash of unstyled text). cyan-600 #0891b2 is the single
# accent — closest Tailwind hue to the original brand teal #0fa3b1.
THEME = gr.themes.Base(
    primary_hue="cyan",
    secondary_hue="cyan",
    neutral_hue="slate",
    radius_size="md",
    font=gr.themes.GoogleFont("Inter"),
    font_mono=gr.themes.GoogleFont("JetBrains Mono"),
).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_500",
    button_primary_text_color="white",
    slider_color="*primary_600",
)

# ShotCraft design system — premium light theme, single cyan accent.
# 1px borders for separation, layered soft shadows, Inter with tight heading
# tracking, 200ms ease-out micro-interactions, skeleton shimmer on pending
# outputs, visible dual-ring focus. Lines marked KEEP are load-bearing
# Gradio fixes — do not remove.
SHOTCRAFT_CSS = """
:root {
    --sc-bg: #f8fafc;
    --sc-card: #ffffff;
    --sc-ink: #0f172a;
    --sc-body: #475569;
    --sc-muted: #94a3b8;
    --sc-border: rgba(15, 23, 42, 0.08);
    --sc-border-strong: rgba(15, 23, 42, 0.16);
    --sc-accent: #0891b2;
    --sc-accent-strong: #0e7490;
    --sc-accent-soft: rgba(8, 145, 178, 0.10);
    --sc-accent-ring: rgba(8, 145, 178, 0.22);
    --sc-danger: #dc2626;
    --sc-shadow-sm: 0 1px 2px rgba(2, 8, 23, 0.05);
    --sc-shadow-md: 0 1px 2px rgba(2, 8, 23, 0.05), 0 10px 30px rgba(2, 8, 23, 0.07);
    --sc-shadow-lg: 0 2px 4px rgba(2, 8, 23, 0.06), 0 18px 50px rgba(2, 8, 23, 0.10);
    --sc-ease: cubic-bezier(0.22, 1, 0.36, 1);
}

/* KEEP: prevent horizontal scroll on small screens */
/* clip, not hidden: overflow-x:hidden computes overflow-y to auto, turning
   body into a non-scrolling scroll container that captures position:sticky. */
html, body {
    overflow-x: clip !important;
}

body,
.gradio-container {
    background:
        radial-gradient(1100px 480px at 85% -8%, rgba(8, 145, 178, 0.07), transparent 60%),
        linear-gradient(180deg, #fbfcfe 0%, var(--sc-bg) 100%) !important;
    color: var(--sc-body) !important;
    font-family: 'Inter', ui-sans-serif, system-ui, sans-serif !important;
}

.gradio-container {
    max-width: 1180px !important;
    margin: 0 auto !important;
    padding: 28px 22px 48px !important;
}

/* KEEP: let the flex/grid chain shrink below content min-width on small screens */
.gradio-container, .gradio-container .main, .gradio-container .app,
.gradio-container .contain, .gradio-container .column, .gradio-container .row,
.gradio-container .block, .gradio-container .form,
.gradio-container .tab-container, .gradio-container [role="tablist"] {
    min-width: 0 !important;
    max-width: 100% !important;
}

.prose h1, .prose h2, .prose h3 {
    color: var(--sc-ink) !important;
    letter-spacing: -0.02em;
    font-weight: 700;
}

/* ---------- Interaction primitives ---------- */
button, [role="tab"], .thumbnail-item, .label-wrap, tbody tr {
    cursor: pointer !important;
}

button {
    transition: transform 150ms var(--sc-ease), box-shadow 200ms var(--sc-ease),
                background 200ms var(--sc-ease), border-color 200ms var(--sc-ease),
                color 200ms var(--sc-ease) !important;
}

button:active {
    transform: scale(0.98);
}

button:focus-visible, [role="tab"]:focus-visible {
    outline: none !important;
    box-shadow: 0 0 0 2px #ffffff, 0 0 0 4px var(--sc-accent) !important;
}

input[type="checkbox"], input[type="radio"] {
    accent-color: var(--sc-accent);
}

/* ---------- Hero ---------- */
.shotcraft-hero {
    position: relative;
    overflow: hidden;
    padding: 36px 36px 30px;
    border: 1px solid var(--sc-border);
    border-radius: 18px;
    background: rgba(255, 255, 255, 0.82);
    backdrop-filter: blur(12px); /* safe here: hero contains no dropdowns */
    box-shadow: var(--sc-shadow-md);
}

.shotcraft-hero::before {
    content: "";
    position: absolute;
    top: 0; left: 0; right: 0;
    height: 3px;
    background: linear-gradient(90deg, var(--sc-accent), rgba(8, 145, 178, 0.25));
}

.shotcraft-kicker {
    color: var(--sc-accent);
    font-size: 11px;
    font-weight: 600;
    letter-spacing: 0.14em;
    text-transform: uppercase;
    margin-bottom: 14px;
}

.shotcraft-hero h1 {
    max-width: 720px;
    margin: 0;
    color: var(--sc-ink);
    font-weight: 700;
    font-size: clamp(30px, 4.4vw, 46px);
    line-height: 1.08;
    letter-spacing: -0.03em;
}

.shotcraft-hero p {
    max-width: 640px;
    margin: 14px 0 22px;
    color: var(--sc-body);
    font-size: 15px;
    line-height: 1.65;
}

.hero-chips {
    display: flex;
    flex-wrap: wrap;
    gap: 8px;
}

.hero-chip {
    display: inline-flex;
    align-items: center;
    gap: 6px;
    padding: 6px 13px;
    border: 1px solid var(--sc-border);
    border-radius: 999px;
    background: #ffffff;
    color: var(--sc-body);
    font-size: 12px;
    font-weight: 500;
    letter-spacing: 0.01em;
    transition: border-color 200ms var(--sc-ease), box-shadow 200ms var(--sc-ease);
}

.hero-chip:hover {
    border-color: var(--sc-accent-ring);
    box-shadow: var(--sc-shadow-sm);
}

.hero-chip b {
    color: var(--sc-accent);
    font-weight: 700;
}

/* ---------- Status bar ---------- */
.shotcraft-status {
    display: flex;
    align-items: center;
    gap: 10px;
    margin: 14px 0 16px;
    padding: 10px 14px;
    border-radius: 10px;
    font-size: 12px;
    font-weight: 600;
    letter-spacing: 0.02em;
}

.shotcraft-status.ok {
    border: 1px solid var(--sc-accent-ring);
    background: var(--sc-accent-soft);
    color: var(--sc-accent-strong);
}

.shotcraft-status.fail {
    border: 1px solid rgba(220, 38, 38, 0.30);
    background: rgba(220, 38, 38, 0.07);
    color: var(--sc-danger);
}

.shotcraft-status.checking {
    border: 1px dashed var(--sc-border-strong);
    background: #ffffff;
    color: var(--sc-muted);
}

.status-dot {
    width: 8px;
    height: 8px;
    border-radius: 50%;
    background: currentColor;
    animation: dotPulse 1.8s var(--sc-ease) infinite;
}

@keyframes dotPulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.45; }
}

/* ---------- Stage rail ---------- */
.stage-rail {
    display: grid;
    grid-template-columns: repeat(3, 1fr);
    gap: 12px;
    margin: 16px 0 20px;
}

.stage-card {
    position: relative;
    overflow: hidden;
    padding: 16px 18px 14px;
    border: 1px solid var(--sc-border);
    border-radius: 14px;
    background: var(--sc-card);
    box-shadow: var(--sc-shadow-sm);
    transition: transform 200ms var(--sc-ease), box-shadow 200ms var(--sc-ease),
                border-color 200ms var(--sc-ease), opacity 200ms var(--sc-ease);
}

.stage-card::before {
    content: "";
    position: absolute;
    top: 0; left: 0; right: 0;
    height: 3px;
    background: transparent;
    transition: background 200ms var(--sc-ease);
}

.stage-card.pending { opacity: 0.55; }

.stage-card.done::before { background: var(--sc-accent-ring); }
.stage-card.done .stage-label { color: var(--sc-accent-strong); }

.stage-card.active {
    transform: translateY(-2px);
    border-color: var(--sc-accent-ring);
    box-shadow: var(--sc-shadow-md), 0 8px 24px rgba(8, 145, 178, 0.12);
    opacity: 1;
}

.stage-card.active::before { background: var(--sc-accent); }
.stage-card.active .stage-label { color: var(--sc-accent); }

.stage-label {
    display: block;
    color: var(--sc-muted);
    font-size: 10px;
    font-weight: 600;
    letter-spacing: 0.12em;
    text-transform: uppercase;
}

.stage-card strong {
    display: block;
    margin: 7px 0 3px;
    color: var(--sc-ink);
    font-size: 16px;
    font-weight: 700;
    letter-spacing: -0.01em;
}

.stage-card span:last-child {
    color: var(--sc-muted);
    font-size: 13px;
    line-height: 1.5;
}

/* ---------- Tabs: segmented control ---------- */
#shotcraft-tabs [role="tablist"] {
    gap: 4px !important;
    border: 1px solid var(--sc-border) !important;
    border-radius: 12px !important;
    background: #eef2f6 !important;
    padding: 4px !important;
    box-shadow: none !important;
}

#shotcraft-tabs button[role="tab"] {
    border-radius: 9px !important;
    color: var(--sc-muted) !important;
    font-weight: 600 !important;
    font-size: 13px !important;
    letter-spacing: 0 !important;
    border: 1px solid transparent !important;
    transition: all 200ms var(--sc-ease) !important;
}

#shotcraft-tabs button[role="tab"]:hover {
    color: var(--sc-ink) !important;
    background: rgba(255, 255, 255, 0.6) !important;
}

#shotcraft-tabs button[role="tab"][aria-selected="true"] {
    background: #ffffff !important;
    color: var(--sc-ink) !important;
    border-color: var(--sc-border) !important;
    box-shadow: var(--sc-shadow-sm) !important;
}

.stage-tab {
    /* KEEP opacity-only: transform/filter (even identity, kept by fill-mode)
       make the tab a containing block for fixed elements -> dropdown menus
       misplaced. */
    animation: tabReveal 320ms var(--sc-ease);
}

@keyframes tabReveal {
    from { opacity: 0; }
    to { opacity: 1; }
}

/* ---------- Panels & cards ---------- */
.glass-panel,
.concept-card,
.render-panel {
    border: 1px solid var(--sc-border) !important;
    border-radius: 16px !important;
    background: var(--sc-card) !important;
    box-shadow: var(--sc-shadow-md) !important;
    /* KEEP: NO backdrop-filter here — it turns the panel into a containing
       block for position:fixed children, which teleports Gradio dropdown
       menus. */
}

.glass-panel, .render-panel {
    padding: 20px !important;
}

.concept-card {
    margin-bottom: 10px !important;
    box-shadow: var(--sc-shadow-sm) !important;
    transition: transform 200ms var(--sc-ease), border-color 200ms var(--sc-ease),
                box-shadow 200ms var(--sc-ease);
}

.concept-card:hover {
    transform: translateY(-1px);
    border-color: var(--sc-border-strong) !important;
    box-shadow: var(--sc-shadow-md) !important;
}

.concept-card > div:first-child {
    border-radius: 14px !important;
    color: var(--sc-ink) !important;
    font-weight: 600 !important;
}

.analysis-card {
    border: 1px solid var(--sc-accent-ring) !important;
    border-left: 3px solid var(--sc-accent) !important;
    border-radius: 12px !important;
    background: var(--sc-accent-soft) !important;
    padding: 14px 16px !important;
    color: var(--sc-body) !important;
}

/* KEEP: Gradio group internals — no gray block backgrounds inside cards */
.glass-panel .block, .glass-panel .form, .glass-panel .styler,
.render-panel .block, .render-panel .form, .render-panel .styler,
.concept-card .block, .concept-card .form, .concept-card .styler {
    background: transparent !important;
}

/* ---------- Forms ---------- */
label,
.block label,
.form label {
    color: var(--sc-body) !important;
    font-size: 12px !important;
    font-weight: 600 !important;
    letter-spacing: 0 !important;
    text-transform: none !important;
}

textarea,
input,
select {
    border-color: rgba(15, 23, 42, 0.12) !important;
    border-radius: 8px !important;
    background: #ffffff !important;
    color: var(--sc-ink) !important;
    transition: border-color 150ms var(--sc-ease), box-shadow 150ms var(--sc-ease) !important;
}

textarea:focus,
input:focus {
    border-color: var(--sc-accent) !important;
    box-shadow: 0 0 0 3px var(--sc-accent-ring) !important;
}

/* KEEP scoped to ul.options only — [role="listbox"] also matches the dropdown
   INPUT wrap and painted a stray pill border inside the field. */
ul.options {
    z-index: 9999 !important;
    border: 1px solid var(--sc-border) !important;
    border-radius: 10px !important;
    background: #ffffff !important;
    box-shadow: var(--sc-shadow-lg) !important;
}

ul.options [role="option"], ul.options li {
    background: transparent !important;
    color: var(--sc-ink) !important;
}

ul.options [role="option"]:hover, ul.options li:hover,
ul.options [role="option"][aria-selected="true"], ul.options li.selected {
    background: var(--sc-accent-soft) !important;
    color: var(--sc-accent-strong) !important;
}

/* Examples table */
.table-wrap, table, thead, tbody, tr, th, td {
    background: #ffffff !important;
    color: var(--sc-body) !important;
    border-color: var(--sc-border) !important;
}

tbody tr { transition: background 150ms var(--sc-ease) !important; }
tbody tr:hover, tbody tr:hover td { background: var(--sc-accent-soft) !important; }

/* ---------- Buttons ---------- */
button.primary,
.shotcraft-primary button,
button.shotcraft-primary {
    border: 0 !important;
    border-radius: 10px !important;
    background: linear-gradient(180deg, #0aa3c4, var(--sc-accent)) !important;
    color: #ffffff !important;
    font-weight: 600 !important;
    letter-spacing: 0 !important;
    box-shadow: var(--sc-shadow-sm), inset 0 1px 0 rgba(255, 255, 255, 0.18) !important;
}

button.primary:hover,
.shotcraft-primary button:hover,
button.shotcraft-primary:hover {
    transform: translateY(-1px);
    background: linear-gradient(180deg, #0db0d2, #0a9cc0) !important;
    box-shadow: 0 8px 22px rgba(8, 145, 178, 0.28), inset 0 1px 0 rgba(255, 255, 255, 0.18) !important;
}

button.primary:active,
.shotcraft-primary button:active,
button.shotcraft-primary:active {
    transform: translateY(0) scale(0.98);
}

.shotcraft-secondary button,
button.shotcraft-secondary {
    border: 1px solid rgba(15, 23, 42, 0.14) !important;
    border-radius: 10px !important;
    background: #ffffff !important;
    color: var(--sc-ink) !important;
    font-weight: 600 !important;
    box-shadow: var(--sc-shadow-sm) !important;
}

.shotcraft-secondary button:hover,
button.shotcraft-secondary:hover {
    transform: translateY(-1px);
    border-color: var(--sc-accent-ring) !important;
    background: #fbfdfe !important;
    color: var(--sc-accent-strong) !important;
}

.continue-cta {
    margin-top: 16px !important;
}

.continue-cta button,
button.continue-cta {
    width: 100% !important;
    font-size: 15px !important;
    padding: 13px 22px !important;
    animation: ctaGlow 2.4s var(--sc-ease) infinite;
}

@keyframes ctaGlow {
    0%, 100% { box-shadow: 0 4px 14px rgba(8, 145, 178, 0.18); }
    50% { box-shadow: 0 6px 26px rgba(8, 145, 178, 0.38); }
}

/* ---------- Gallery & media ---------- */
.shotcraft-gallery, .source-preview {
    border-radius: 14px !important;
    background: #ffffff !important;
}

.shotcraft-gallery .thumbnail-item {
    border-radius: 10px !important;
    transition: transform 200ms var(--sc-ease), box-shadow 200ms var(--sc-ease) !important;
}

.shotcraft-gallery .thumbnail-item:hover {
    transform: translateY(-2px);
    box-shadow: var(--sc-shadow-md) !important;
}

.shotcraft-gallery .thumbnail-item.selected {
    outline: 2px solid var(--sc-accent) !important;
    outline-offset: 2px;
}

/* Empty states: dashed, quiet — never a bare gray void */
.gradio-container .empty {
    border: 1.5px dashed rgba(15, 23, 42, 0.14) !important;
    border-radius: 12px !important;
    background: rgba(248, 250, 252, 0.7) !important;
    color: var(--sc-muted) !important;
}

/* ---------- Pending state: skeleton shimmer + status pill ----------
   Hooks Gradio's built-in .generating pending class — pure CSS, no JS.
   position:relative is safe for the fixed-position pill (only transform/
   filter/backdrop-filter create containing blocks for fixed children). */
.generating {
    position: relative !important;
    overflow: hidden !important;
    border-color: var(--sc-accent-ring) !important;
}

.generating::after {
    content: "";
    position: absolute;
    inset: 0;
    z-index: 5;
    pointer-events: none;
    background: linear-gradient(100deg, transparent 30%, rgba(255, 255, 255, 0.55) 50%, transparent 70%);
    background-size: 200% 100%;
    animation: shimmer 1.6s linear infinite;
}

@keyframes shimmer {
    from { background-position: 150% 0; }
    to { background-position: -50% 0; }
}

.generating::before {
    content: "\\25CF  Generating";
    position: fixed;
    top: 16px;
    right: 20px;
    z-index: 1000;
    padding: 6px 14px;
    border: 1px solid var(--sc-accent-ring);
    border-radius: 999px;
    background: rgba(255, 255, 255, 0.92);
    color: var(--sc-accent-strong);
    font-size: 12px;
    font-weight: 600;
    letter-spacing: 0.04em;
    box-shadow: var(--sc-shadow-lg);
    animation: dotPulse 1.4s var(--sc-ease) infinite;
    pointer-events: none;
}

/* ---------- Toasts ---------- */
.toast-body {
    border: 1px solid var(--sc-border) !important;
    border-radius: 12px !important;
    background: rgba(255, 255, 255, 0.94) !important;
    backdrop-filter: blur(10px);
    box-shadow: var(--sc-shadow-lg) !important;
}

.visually-hidden {
    position: absolute !important;
    left: -9999px !important;
    width: 1px !important;
    height: 1px !important;
    overflow: hidden !important;
}

/* ---------- Stage rail = navigation (tab pills hidden) ---------- */
#shotcraft-tabs [role="tablist"] {
    position: absolute !important;
    left: -9999px !important;
    top: 0 !important;
    width: 600px !important;  /* real width: zero-width tablists render no tab buttons */
    pointer-events: none !important;
}

.stage-card {
    cursor: pointer;
    user-select: none;
}

.stage-card:hover {
    transform: translateY(-2px);
    border-color: var(--sc-accent-ring);
    box-shadow: var(--sc-shadow-md);
}

.stage-card:focus-visible {
    outline: 2px solid var(--sc-accent);
    outline-offset: 2px;
}

/* ---------- Sticky continue CTA ---------- */
/* Gradio sets overflow:hidden on the container, which captures position:sticky
   (a non-scrolling scroll container). overflow:clip clips identically but
   does NOT create a scroll container, so the CTA can stick to the real
   scrollport (Gradio's inner overflow-y:auto wrapper). */
.gradio-container {
    overflow: clip !important;
}

.continue-cta,
button.continue-cta {
    position: sticky !important;
    bottom: 14px;
    z-index: 60;
}

/* ---------- Locked product look ---------- */
.locked-box {
    border-left: 4px solid var(--sc-accent) !important;
    border-radius: 12px !important;
    background: rgba(8, 145, 178, 0.05) !important;
    padding: 4px 10px !important;
}

.locked-box textarea {
    font-weight: 600 !important;
}

/* ---------- Palette chips ---------- */
.palette-chips {
    display: flex;
    align-items: center;
    gap: 6px;
    padding: 6px 0 2px;
    min-height: 24px;
}

.palette-chip {
    width: 22px;
    height: 22px;
    border-radius: 7px;
    border: 1px solid rgba(15, 23, 42, 0.18);
    box-shadow: inset 0 0 0 2px rgba(255, 255, 255, 0.35);
    display: inline-block;
}

/* Hide the analysis strip until Stage 1 fills it (it's an empty pill before) */
.analysis-card.hide-container,
.analysis-card:not(:has(.prose)) {
    display: none !important;
}

.boards-toolbar button {
    padding: 6px 14px !important;
    font-size: 12.5px !important;
    white-space: nowrap !important;
}

/* ---------- Mobile ---------- */
@media (max-width: 760px) {
    .gradio-container {
        padding: 14px 10px 30px !important;
    }

    .stage-rail {
        grid-template-columns: 1fr;
    }

    .shotcraft-hero {
        padding: 22px 18px;
    }

    .shotcraft-hero h1 {
        font-size: clamp(24px, 8.5vw, 34px);
        overflow-wrap: break-word;
    }

    .shotcraft-hero p {
        font-size: 14px;
    }

    .shotcraft-status {
        white-space: normal;
        word-break: break-word;
    }

    .shotcraft-gallery .grid-container {
        grid-template-columns: repeat(2, minmax(0, 1fr)) !important;
    }

    .table-wrap {
        overflow-x: auto !important;
        max-width: 100% !important;
    }

    .prose h3, .prose h2 {
        white-space: normal !important;
        overflow-wrap: break-word;
    }
}
"""


def _hero_html() -> str:
    return """
    <section class="shotcraft-hero">
        <div class="shotcraft-kicker">ShotCraft / AI campaign console</div>
        <h1>Direct the shot. Render the campaign.</h1>
        <p>
            Upload one product photo, let MiniCPM draft five grounded art-direction
            boards, then move into FLUX rendering with editable prompts and
            deterministic hero-frame control.
        </p>
        <div class="hero-chips">
            <span class="hero-chip"><b>01</b> Vision director</span>
            <span class="hero-chip"><b>02</b> Five shot concepts</span>
            <span class="hero-chip"><b>03</b> FLUX render reel</span>
            <span class="hero-chip"><b>04</b> Export package</span>
        </div>
    </section>
    """


def _backend_status_html(status: dict) -> str:
    if status.get("status") == "ok":
        return (
            '<div class="shotcraft-status ok"><span class="status-dot"></span>'
            f'BACKEND ONLINE / {status.get("stage1", "?")} / '
            f'{status.get("stage2", "?")} / MODAL</div>'
        )
    return (
        '<div class="shotcraft-status fail"><span class="status-dot"></span>'
        f'BACKEND UNREACHABLE / {API_URL} / deploy Modal before generating</div>'
    )


def _status_checking_html() -> str:
    return (
        '<div class="shotcraft-status checking"><span class="status-dot"></span>'
        'CHECKING BACKEND…</div>'
    )


def _stage_rail(active: str = "direct") -> str:
    steps = [
        ("direct", "01", "Direct", "Analyze product and draft concepts"),
        ("render", "02", "Render", "Generate the five-frame reel"),
        ("export", "03", "Package", "Pick heroes and download assets"),
    ]
    active_index = next((i for i, step in enumerate(steps) if step[0] == active), 0)
    cards = []
    for index, (key, num, title, desc) in enumerate(steps):
        state = "active" if key == active else "done" if index < active_index else "pending"
        cards.append(
            f'<div class="stage-card {state}" data-target="{key}" role="button" '
            f'tabindex="0" title="Go to {title}">'
            f'<span class="stage-label">{num} / {state}</span>'
            f'<strong>{title}</strong>'
            f'<span>{desc}</span>'
            '</div>'
        )
    return '<div class="stage-rail">' + "".join(cards) + "</div>"


def _example_brand() -> str:
    try:
        with open(os.path.join(EXAMPLES_DIR, "brand_brief.txt")) as f:
            return f.read().strip()
    except OSError:
        return ""

def run_stage1(image, brand, category, preset):
    if image is None:
        raise gr.Error("Please upload a product photo first.")
    if not brand or not brand.strip():
        raise gr.Error("Please add a short brand description.")
    if getattr(image, "width", 0) * getattr(image, "height", 0) > 36_000_000:
        raise gr.Error("Image too large — please use a photo under ~36 megapixels.")
    try:
        pkg = generate_concepts(image, brand.strip()[:300], category, preset)
    except (RuntimeError, BackendError) as e:
        raise gr.Error(f"Concept generation failed: {e}")
    pa = pkg.product_analysis
    analysis = (f"**Detected:** {pa.product_type} | **Materials:** {pa.materials} | "
                f"**Colors:** {', '.join(pa.colors)} | {pa.distinguishing_features}")
    fields = []
    for s in pkg.shots:
        fields.extend([s.concept_name, s.scene, s.camera_angle, s.lighting,
                       ", ".join(s.color_palette), s.props, s.marketing_angle,
                       s.image_prompt])
    return (pkg, analysis, *fields, gr.update(interactive=True))

FIELD_KEYS = ["concept_name", "scene", "camera_angle", "lighting",
              "color_palette", "props", "marketing_angle", "image_prompt"]

def apply_edits(pkg, values):
    """FR-1.3 + Q-3: write every edited concept field back into the package."""
    n = len(FIELD_KEYS)
    for i, s in enumerate(pkg.shots):
        chunk = values[i * n:(i + 1) * n]
        if len(chunk) < n:
            break
        for key, val in zip(FIELD_KEYS, chunk):
            if val is None or not str(val).strip():
                continue
            v = str(val).strip()
            if key == "color_palette":
                setattr(s, key, [c.strip() for c in v.split(",") if c.strip()])
            else:
                setattr(s, key, v)

def run_stage2(pkg, preset, aspect, *field_values, progress=gr.Progress()):
    if pkg is None:
        raise gr.Error("Generate shot concepts first (Step 1).")
    apply_edits(pkg, field_values)  # FR-1.3/FR-2.1 + Q-3: all edits persist
    progress((0, 5), desc="Rendering 5 shots in one GPU lease…")
    frames = generate_frames(pkg.shots, preset, aspect,
                             product_desc=getattr(pkg.product_analysis,
                                                  "canonical_description", ""))  # ONE backend call
    images = [(img, f"Shot {s.id}: {s.concept_name}")
              for s, img in zip(pkg.shots, frames)]
    regen_state = {s.id: 0 for s in pkg.shots}
    progress((5, 5), desc="Done")
    return images, pkg, regen_state


def _palette_chips(palette_str) -> str:
    """Small color swatches next to the palette field (hex or CSS color names)."""
    tokens = [t.strip() for t in str(palette_str or "").split(",") if t.strip()][:6]
    if not tokens:
        return ""
    spans = "".join(
        f'<span class="palette-chip" title="{t}" style="background:{t}"></span>'
        for t in tokens)
    return f'<div class="palette-chips">{spans}</div>'


def _shot_summary(s) -> str:
    """One-line accordion header so collapsed boards stay scannable."""
    angle = (s.marketing_angle or s.scene or "").strip()
    if len(angle) > 70:
        angle = angle[:67].rstrip() + "…"
    return f"Shot {s.id} — {s.concept_name}" + (f" · {angle}" if angle else "")


def _apply_locked(pkg, locked):
    """Editable locked-look box is the source of truth for the render prefix."""
    if pkg is not None and locked and str(locked).strip():
        pkg.product_analysis.canonical_description = str(locked).strip()


def run_stage1_ui(image, brand, category, preset, progress=gr.Progress()):
    """UI wrapper: preserve run_stage1() test contract while revealing the next-step CTA.

    Stays on the Direct tab so concepts can be reviewed/edited, then the
    'Continue to Render →' CTA advances the flow (no manual tab hunting).
    Also mirrors the uploaded product photo into the Render tab, fills the
    editable locked-look box, retitles the concept accordions with one-line
    summaries, and paints the palette swatches.
    """
    progress(0.1, desc="MiniCPM is reading your product…")
    out = run_stage1(image, brand, category, preset)
    progress(1.0, desc="Concepts ready")
    gr.Info("Five shot concepts drafted — review and edit them below.")
    pkg = out[0]
    locked = getattr(pkg.product_analysis, "canonical_description", "")
    acc_updates = [gr.update(label=_shot_summary(s), open=False) for s in pkg.shots]
    chip_updates = [_palette_chips(", ".join(s.color_palette)) for s in pkg.shots]
    return (*out, gr.update(visible=True), image,
            gr.update(value=locked, visible=True), *acc_updates, *chip_updates)


def advance_to_render(reached):
    """CTA handler: smooth-advance the workflow chrome from Direct to Render."""
    new = "export" if reached == "export" else "render"
    return gr.update(selected="render"), _stage_rail(new), new


def run_stage2_ui(pkg, preset, aspect, locked, *field_values, progress=gr.Progress()):
    """UI wrapper: preserve run_stage2() test contract while advancing the stage chrome."""
    _apply_locked(pkg, locked)
    images, updated_pkg, regen_state = run_stage2(
        pkg, preset, aspect, *field_values, progress=progress
    )
    gr.Info("Reel rendered — click any frame to inspect its concept.")
    return images, updated_pkg, regen_state, _stage_rail("export"), "export"


def sync_stage_rail(reached, evt: gr.SelectData):
    """Keep the stage rail honest when the user clicks tabs directly."""
    active = "direct" if evt.index == 0 else ("export" if reached == "export" else "render")
    return _stage_rail(active)


def regen_one(pkg, preset, aspect, shot_id, gallery, regen_state):
    if pkg is None or not gallery:
        raise gr.Error("Nothing to regenerate yet.")
    sid = int(shot_id)
    shot = next((s for s in pkg.shots if s.id == sid), None)
    if shot is None:
        raise gr.Error(f"Shot {sid} not found.")
    regen_state = dict(regen_state or {})
    regen_state[sid] = regen_state.get(sid, 0) + 1
    img = generate_frame(shot, preset, aspect, regen_counter=regen_state[sid],
                         product_desc=getattr(pkg.product_analysis,
                                              "canonical_description", ""))
    gallery = list(gallery)
    gallery[sid - 1] = (img, f"Shot {sid}: {shot.concept_name}")
    return gallery, regen_state


def regen_one_ui(pkg, preset, aspect, locked, shot_id, gallery, regen_state):
    """UI wrapper: preserve regen_one() test contract, add toast feedback."""
    _apply_locked(pkg, locked)
    gallery, regen_state = regen_one(pkg, preset, aspect, shot_id, gallery, regen_state)
    gr.Info(f"Shot {int(shot_id)} rerolled with a fresh seed.")
    return gallery, regen_state


def show_detail(pkg, evt: gr.SelectData):
    """FR-2.5: click a gallery frame -> full concept card + prompt.

    Also points the regenerate picker at the clicked shot, so
    'click frame -> Regenerate' needs no manual dropdown hunting.
    """
    if pkg is None:
        return "", gr.update()
    idx = evt.index if isinstance(evt.index, int) else evt.index[0]
    if idx is None or idx >= len(pkg.shots):
        return "", gr.update()
    s = pkg.shots[idx]
    detail = (f"### Shot {s.id}: {s.concept_name}\n"
              f"**Scene:** {s.scene}\n\n**Camera:** {s.camera_angle} · "
              f"**Lighting:** {s.lighting} · **Palette:** {', '.join(s.color_palette)}\n\n"
              f"**Props:** {s.props}\n\n**Marketing angle:** {s.marketing_angle}\n\n"
              f"**Prompt:** `{s.image_prompt}`")
    return detail, gr.update(value=s.id)

def export_zip(pkg, gallery, heroes, regen_state):
    """FR-3.1: concepts.md + concepts.json + 5 PNGs + selection_manifest.json."""
    if pkg is None or not gallery:
        raise gr.Error("Generate frames before exporting.")
    regen_state = regen_state or {}
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("concepts.json", pkg.to_json())
        md = ["# ShotCraft package\n"]
        for s in pkg.shots:
            md.append(f"## Shot {s.id}: {s.concept_name}\n{s.scene}\n\nPrompt: {s.image_prompt}\n")
        z.writestr("concepts.md", "\n".join(md))
        def _rc(sid):
            return int(regen_state.get(sid, regen_state.get(str(sid), 0)))
        manifest = {
            "hero_frames": heroes or [],
            "shots": [
                {
                    "id": s.id,
                    "concept_name": s.concept_name,
                    "file": f"shot_{s.id}.png",
                    "image_prompt": s.image_prompt,  # post-edit (Q-3: edits persist)
                    "regen_count": _rc(s.id),
                    "seed": seed_for(s.id, _rc(s.id)),
                    "is_hero": s.id in (heroes or []),
                }
                for s in pkg.shots
            ],
            "backend": API_URL,
        }
        z.writestr("selection_manifest.json", json.dumps(manifest, indent=2))
        for i, item in enumerate(gallery, start=1):
            img = item[0] if isinstance(item, (tuple, list)) else item
            ib = io.BytesIO(); img.save(ib, format="PNG")
            z.writestr(f"shot_{i}.png", ib.getvalue())
    fd, path = tempfile.mkstemp(prefix="shotcraft_pkg_", suffix=".zip")
    os.write(fd, buf.getvalue())
    os.close(fd)
    return path


def export_zip_ui(pkg, gallery, heroes, regen_state):
    """UI wrapper: preserve export_zip() test contract, reveal the download card."""
    path = export_zip(pkg, gallery, heroes, regen_state)
    gr.Info("Package ready — your ZIP is below.")
    return gr.update(value=path, visible=True)


with gr.Blocks(title="ShotCraft — AI Shot Director") as demo:
    gr.HTML(_hero_html())
    # Live status per page load (demo.load) — a build-time health() call would
    # freeze the banner at whatever was true when the Space booted.
    backend_status = gr.HTML(_status_checking_html())
    stage_rail = gr.HTML(_stage_rail("direct"), elem_id="stage-rail")
    pkg_state = gr.State(None)
    regen_state = gr.State({})
    reached_stage = gr.State("direct")

    with gr.Tabs(selected="direct", elem_id="shotcraft-tabs",
                 elem_classes=["workflow-tabs"]) as workflow_tabs:
        with gr.Tab("01 Direct", id="direct", elem_classes=["stage-tab"]):
            with gr.Group(elem_classes=["glass-panel"]):
                gr.Markdown("### Product briefing room")
                with gr.Row():
                    photo = gr.Image(label="Product photo", type="pil")
                    with gr.Column():
                        brand = gr.Textbox(
                            label="Brand description",
                            lines=2,
                            max_lines=4,
                            max_length=300,
                            placeholder="Materials, audience, vibe — e.g. Handmade ceramic mugs, "
                                        "Scandinavian style, eco-friendly, for slow-morning coffee people",
                            info="Up to 300 characters — used to ground every concept",
                        )
                        category = gr.Dropdown(
                            CATEGORIES, label="Category", value=CATEGORIES[0],
                            info="Guides scene and prop choices",
                        )
                        preset = gr.Dropdown(
                            STYLE_PRESETS, label="Style preset", value=STYLE_PRESETS[0],
                            info="Sets the overall art direction for every render",
                        )
                        go1 = gr.Button(
                            "Direct my shots",
                            variant="primary",
                            elem_classes=["shotcraft-primary"],
                        )
                demo_photo = os.path.join(EXAMPLES_DIR, "demo_product.png")
                if os.path.exists(demo_photo):  # FR-3.2 pre-loaded example
                    demo_btn = gr.Button(
                        "✨ Try the demo product",
                        size="sm",
                        elem_classes=["shotcraft-secondary"],
                    )
                    demo_btn.click(
                        lambda: (Image.open(demo_photo), _example_brand(), "Home", "Minimal"),
                        None, [photo, brand, category, preset],
                    )
            analysis_md = gr.Markdown(elem_classes=["analysis-card"])
            locked_box = gr.Textbox(
                label="🔒 Locked product look — prefixed to every render",
                visible=False,
                lines=2,
                max_lines=3,
                info="Correct it if the analysis got a color or detail wrong — "
                     "this exact description pins the product in all 5 shots",
                elem_classes=["locked-box"],
            )
            with gr.Row(elem_classes=["boards-toolbar"]):
                gr.Markdown("### Editable concept boards\nEvery field is live. "
                            "Edits carry into rendering and export.", scale=5)
                expand_all = gr.Button("Expand all", size="sm", scale=0, min_width=120,
                                       elem_classes=["shotcraft-secondary"])
                collapse_all = gr.Button("Collapse all", size="sm", scale=0, min_width=120,
                                         elem_classes=["shotcraft-secondary"])
            field_boxes = []
            accordions = []
            chip_htmls = []
            for i in range(5):
                with gr.Accordion(f"Shot {i + 1}", open=False,
                                  elem_classes=["concept-card"]) as acc:
                    with gr.Row():
                        nb = gr.Textbox(label="Concept name")
                        cb = gr.Textbox(label="Camera angle")
                        lb = gr.Textbox(label="Lighting")
                    sb = gr.Textbox(label="Scene", max_lines=2)
                    with gr.Row():
                        with gr.Column(scale=1):
                            pb = gr.Textbox(label="Palette (comma-separated hex)")
                            chips = gr.HTML()
                        rb = gr.Textbox(label="Props")
                        mb = gr.Textbox(label="Marketing angle")
                    ib = gr.Textbox(label="Image prompt (FLUX)", max_lines=3)
                    field_boxes.extend([nb, sb, cb, lb, pb, rb, mb, ib])
                    accordions.append(acc)
                    chip_htmls.append(chips)
                    pb.change(_palette_chips, pb, chips, show_progress="hidden")
            expand_all.click(lambda: [gr.update(open=True)] * 5, None, accordions)
            collapse_all.click(lambda: [gr.update(open=False)] * 5, None, accordions)
            continue_btn = gr.Button(
                "Continue to Render — generate 5 shots →",
                variant="primary",
                visible=False,
                elem_classes=["shotcraft-primary", "continue-cta"],
            )

        with gr.Tab("02 Render", id="render", elem_classes=["stage-tab"]):
            with gr.Group(elem_classes=["render-panel"]):
                gr.Markdown("### Render reel")
                with gr.Row():
                    aspect = gr.Radio(
                        ["1:1", "16:9"], value="1:1", label="Aspect ratio",
                        info="1:1 for product tiles · 16:9 for banners",
                    )
                    go2 = gr.Button(
                        "Re-render 5 shots",
                        variant="primary",
                        interactive=False,
                        elem_classes=["shotcraft-primary"],
                    )
                with gr.Row(equal_height=False):
                    with gr.Column(scale=1, min_width=190):
                        src_preview = gr.Image(
                            label="Source product",
                            type="pil",
                            interactive=False,
                            height=360,
                            buttons=[],
                            elem_classes=["source-preview"],
                        )
                    with gr.Column(scale=4):
                        gallery = gr.Gallery(
                            label="Storyboard reel",
                            columns=4,
                            height=360,
                            format="png",
                            buttons=["download", "fullscreen"],
                            elem_classes=["shotcraft-gallery"],
                        )
                detail_md = gr.Markdown()
                with gr.Row():
                    regen_id = gr.Dropdown(
                        [1, 2, 3, 4, 5], value=1, label="Shot #",
                        info="Click a frame above to pick it here",
                    )
                    regen_btn = gr.Button("Regenerate this shot", elem_classes=["shotcraft-secondary"])
                    heroes = gr.CheckboxGroup(
                        [1, 2, 3, 4, 5], label="Hero frames",
                        info="Lead frames for the campaign — recorded in the export manifest",
                    )
                    export_btn = gr.Button("Download package ZIP", elem_classes=["shotcraft-secondary"])
                zip_out = gr.File(label="Package", visible=False)

    demo.load(lambda: _backend_status_html(health()), None, backend_status)
    go1.click(run_stage1_ui, [photo, brand, category, preset],
              [pkg_state, analysis_md, *field_boxes, go2, continue_btn, src_preview,
               locked_box, *accordions, *chip_htmls])
    # One-click flow: the CTA switches to the Render tab, then immediately
    # kicks off the 5-shot render. go2 remains as a manual re-render after
    # concept edits (or if the auto-render failed on a cold backend).
    continue_btn.click(advance_to_render, [reached_stage],
                       [workflow_tabs, stage_rail, reached_stage],
                       js="() => window.scrollTo({top: 0, behavior: 'smooth'})"
                       ).then(run_stage2_ui, [pkg_state, preset, aspect, locked_box, *field_boxes],
                              [gallery, pkg_state, regen_state, stage_rail, reached_stage])
    workflow_tabs.select(sync_stage_rail, [reached_stage], [stage_rail])
    go2.click(run_stage2_ui, [pkg_state, preset, aspect, locked_box, *field_boxes],
              [gallery, pkg_state, regen_state, stage_rail, reached_stage])
    gallery.select(show_detail, [pkg_state], [detail_md, regen_id])
    regen_btn.click(regen_one_ui, [pkg_state, preset, aspect, locked_box, regen_id, gallery, regen_state],
                    [gallery, regen_state])
    export_btn.click(export_zip_ui, [pkg_state, gallery, heroes, regen_state], [zip_out])

if __name__ == "__main__":
    # ssr_mode=False: Gradio SSR (Node proxy) restores stale sessions across
    # Space rebuilds -> dead /tmp/gradio file refs crash Image preprocess.
    demo.launch(max_file_size="10mb", ssr_mode=False, theme=THEME,
                css=SHOTCRAFT_CSS, head=FORCE_THEME_HEAD)  # FR-1.1 upload cap