app.py

"""Gradio Space for Banuba AI Tasks API talking-avatar generation.

Flow:
1. OAuth2 client_credentials token.
2. Request presigned upload URLs for image/audio.
3. Upload media to object storage.
4. Create a `video.lipsync` task.
5. Poll task status with gr.Timer and restore the latest task from BrowserState.
"""

from __future__ import annotations

import hashlib
import mimetypes
import os
import tempfile
import threading
import time
import uuid
from pathlib import Path
from typing import Any, Dict, Optional

import gradio as gr
import requests

try:
    from posthog import Posthog
except Exception:  # pragma: no cover - analytics must never break the app
    Posthog = None  # type: ignore[assignment]


# -----------------------------
# Configuration
# -----------------------------

AUTH_BASE_URL = os.getenv("BANUBA_AUTH_BASE_URL", "https://ai.banuba.net/auth/v1").rstrip("/")
API_BASE_URL = os.getenv("BANUBA_API_BASE_URL", "https://ai.banuba.net/api/v1").rstrip("/")
CLIENT_ID = os.getenv("BANUBA_CLIENT_ID", "")
CLIENT_SECRET = os.getenv("BANUBA_CLIENT_SECRET", "")
BANUBA_SCOPE = os.getenv("BANUBA_SCOPE", "tasks:write tasks:read")

POSTHOG_PROJECT_API_KEY = os.getenv("POSTHOG_PROJECT_API_KEY") or os.getenv("POSTHOG_API_KEY")
POSTHOG_HOST = os.getenv("POSTHOG_HOST", "https://us.i.posthog.com")

# Google Analytics (gtag.js). The measurement ID is a public, client-side value.
GA_MEASUREMENT_ID = os.getenv("GA_MEASUREMENT_ID", "G-YN5BXMSZHC")

POLL_INTERVAL_SECONDS = float(os.getenv("POLL_INTERVAL_SECONDS", "5"))
POLL_MAX_INTERVAL_SECONDS = float(os.getenv("POLL_MAX_INTERVAL_SECONDS", "20"))
# Keep results inside the real system temp dir so Gradio is allowed to serve them.
# On macOS /tmp is a symlink outside $TMPDIR and Gradio's cache check rejects it.
RESULT_DIR = Path(os.getenv("TASK_RESULT_DIR") or (Path(tempfile.gettempdir()) / "banuba_talking_avatar_results"))
RESULT_DIR.mkdir(parents=True, exist_ok=True)

# Keep only the newest N result videos on disk so the ephemeral Space disk cannot
# fill up. Pruned files are re-downloaded from the API on demand if needed.
MAX_RESULT_FILES = int(os.getenv("MAX_RESULT_FILES", "50"))

EXAMPLES_DIR = Path(os.getenv("EXAMPLES_DIR", "assets/examples"))

DEFAULT_PROMPT = os.getenv(
    "BANUBA_DEFAULT_PROMPT",
    "Waist-up, direct-to-camera. I read the script confidently and friendly, "
    "with natural hand gestures timed to speech, brief pauses, subtle facial "
    "expressions. Even lighting, neutral background, steady camera",
)


def _optional_int_env(name: str, default: Optional[int]) -> Optional[int]:
    raw = os.getenv(name)
    if raw is None:
        return default
    raw = raw.strip()
    if raw == "" or raw.lower() in {"none", "null", "auto"}:
        return None
    return int(raw)


# Default max duration (seconds). Kept small to avoid slow/expensive long jobs.
DEFAULT_SECONDS = _optional_int_env("BANUBA_DEFAULT_SECONDS", 4)
DEFAULT_SEED = _optional_int_env("BANUBA_DEFAULT_SEED", 7)

SUCCESS_STATUSES = {"COMPLETED", "SUCCEEDED", "SUCCESS"}
FAILED_STATUSES = {"FAILED", "ERROR", "CANCELED", "CANCELLED"}

TOKEN_LOCK = threading.Lock()
TOKEN_CACHE: Dict[str, Any] = {"access_token": None, "expires_at": 0.0}

POSTHOG_CLIENT = None
POSTHOG_LOCK = threading.Lock()

SESSION_TASKS: Dict[str, str] = {}
TASK_RECORDS: Dict[str, Dict[str, Any]] = {}
TASK_EVENTS_SENT: Dict[str, set[str]] = {}
RECORDS_LOCK = threading.Lock()

# Cap the in-memory bookkeeping so a long-running Space cannot leak memory.
MAX_TRACKED_TASKS = int(os.getenv("MAX_TRACKED_TASKS", "500"))


# -----------------------------
# UI strings / CSS
# -----------------------------

# Load Google Analytics from the Blocks `js` hook (executed on page load) rather than
# from `head=`. Under ssr_mode=False Gradio stores `head` HTML in gradio_config and
# injects it client-side, which does NOT execute <script> tags — so a gtag.js <script>
# in `head` never runs. Here we queue the gtag commands, then append the library via
# createElement, which DOES execute external scripts.
GA_LOAD_JS = (
    f"""
() => {{
  if (window.__bnbGaLoaded) return;
  window.__bnbGaLoaded = true;
  window.dataLayer = window.dataLayer || [];
  function gtag(){{ window.dataLayer.push(arguments); }}
  window.gtag = gtag;
  gtag('js', new Date());
  gtag('config', '{GA_MEASUREMENT_ID}');
  const s = document.createElement('script');
  s.async = true;
  s.src = 'https://www.googletagmanager.com/gtag/js?id={GA_MEASUREMENT_ID}';
  document.head.appendChild(s);
}}
"""
    if GA_MEASUREMENT_ID
    else None
)

STATUS_EMPTY = "### Generation status\nUpload a reference image and reference audio, then click **Generate Talking Photo**."
STATUS_STARTING = "### Generation status\nStarting generation..."
STATUS_GENERATING = "### Generation status\nGenerating your talking photo. This may take a couple of minutes."
STATUS_DONE = "### Generation status\nDone - your video is ready."
STATUS_FAILED = "### Generation status\nGeneration failed. Please try again."
STATUS_SAMPLE = "### Generation status\nThis is a sample result for the selected example. Click **Generate Talking Photo** to make your own."

BTN_GENERATE = "Generate Talking Photo"
BTN_GENERATING = "Generating… please wait"

# Instant client-side feedback on click (no server round-trip), so the button locks
# immediately instead of only after the queued handler responds. Server state reconciles after.
GENERATE_CLICK_JS = """
() => {
  const root = document.getElementById('generate-btn');
  const btn = root && (root.tagName === 'BUTTON' ? root : root.querySelector('button'));
  if (btn) {
    btn.disabled = true;
    const label = btn.querySelector('span') || btn;
    label.textContent = 'Generating… please wait';
  }
}
"""

OUTPUT_EMPTY_HTML = """
<div class="output-state output-empty">
  <div class="output-icon">🎬</div>
  <div class="output-title">Output</div>
  <div class="output-text">Your generated talking photo will appear here.</div>
</div>
"""

OUTPUT_LOADING_HTML = """
<div class="output-state output-loading">
  <div class="loader"></div>
  <div class="output-title">Generating your talking photo</div>
  <div class="output-text">The task is running in the API. You can leave and come back; the Space will resume polling the latest task from this browser.</div>
</div>
"""

OUTPUT_FAILED_HTML = """
<div class="output-state output-failed">
  <div class="output-icon">⚠️</div>
  <div class="output-title">Generation failed</div>
  <div class="output-text">Please try again with another image/audio pair.</div>
</div>
"""

CUSTOM_CSS = """
main, .gradio-container, .fillable:not(.fill_width) {
  width: min(100%, 1180px) !important;
  max-width: 1180px !important;
  margin-left: auto !important;
  margin-right: auto !important;
}
.hero {
  padding: 18px 22px;
  border: 1px solid var(--border-color-primary);
  border-radius: 18px;
  background: linear-gradient(120deg, rgba(99,102,241,0.10), rgba(168,85,247,0.08));
  margin-bottom: 16px;
}
.hero h1 { margin: 0 0 6px 0; }
.hero p { margin: 0; opacity: 0.82; }
.output-state {
  min-height: 360px;
  border: 1px dashed var(--border-color-primary);
  border-radius: 16px;
  display: flex;
  flex-direction: column;
  align-items: center;
  justify-content: center;
  text-align: center;
  padding: 24px;
  background: var(--background-fill-secondary);
}
.output-icon { font-size: 40px; margin-bottom: 12px; }
.output-title { font-size: 18px; font-weight: 700; margin-bottom: 6px; }
.output-text { max-width: 520px; opacity: 0.72; }
.loader {
  width: 38px;
  height: 38px;
  border-radius: 50%;
  border: 4px solid var(--border-color-primary);
  border-top-color: var(--color-accent);
  animation: spin 1s linear infinite;
  margin-bottom: 16px;
}
@keyframes spin { to { transform: rotate(360deg); } }
.task-note {
  opacity: 0.70;
  font-size: 0.9rem;
  margin-top: 8px;
}
"""


# -----------------------------
# Helpers
# -----------------------------


def _now() -> float:
    return time.time()


def _btn(enabled: bool):
    """Generate-button update: disabled + 'Generating…' while busy, enabled + default label otherwise."""
    return gr.update(interactive=enabled, value=BTN_GENERATE if enabled else BTN_GENERATING)


def _request_session_hash(request: Optional[gr.Request]) -> str:
    try:
        return request.session_hash or "unknown-session"  # type: ignore[union-attr]
    except Exception:
        return "unknown-session"


def _ensure_browser_state(state: Optional[dict]) -> dict:
    state = dict(state or {})
    if not state.get("distinct_id"):
        state["distinct_id"] = f"anon_{uuid.uuid4().hex}"
    return state


def _safe_int(value: Optional[int]) -> Optional[int]:
    if value is None:
        return None
    return int(value)


def _guess_content_type(path: str, fallback: str = "application/octet-stream") -> str:
    guessed, _ = mimetypes.guess_type(path)
    return guessed or fallback


def _file_properties(path: Optional[str]) -> Dict[str, Any]:
    if not path:
        return {}
    p = Path(path)
    props: Dict[str, Any] = {
        "filename": p.name,
        "content_type": _guess_content_type(str(p)),
    }
    try:
        props["size_bytes"] = p.stat().st_size
    except OSError:
        pass
    return props


def _prune_result_dir(keep: int = MAX_RESULT_FILES) -> None:
    """Best-effort retention: drop stale temp files and keep only the newest `keep` videos."""
    try:
        entries = [p for p in RESULT_DIR.iterdir() if p.is_file()]
    except OSError:
        return
    media = []
    for p in entries:
        if p.suffix == ".tmp":
            try:
                p.unlink()
            except OSError:
                pass
            continue
        media.append(p)
    media.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    for p in media[keep:]:
        try:
            p.unlink()
        except OSError:
            pass


def _status_with_task(base: str, task_id: Optional[str], remote_status: Optional[str] = None) -> str:
    parts = [base]
    if task_id:
        parts.append(f"\n<div class='task-note'>Task ID: <code>{task_id}</code></div>")
    if remote_status:
        parts.append(f"\n<div class='task-note'>Status: <code>{remote_status}</code></div>")
    return "".join(parts)


def _set_session_task(request: Optional[gr.Request], task_id: Optional[str]) -> None:
    if not request or not task_id:
        return
    SESSION_TASKS[_request_session_hash(request)] = task_id


def _evict_old_tasks_locked() -> None:
    """Drop the oldest task bookkeeping once we exceed the cap. Caller holds RECORDS_LOCK."""
    while len(TASK_RECORDS) > MAX_TRACKED_TASKS:
        oldest = next(iter(TASK_RECORDS))
        TASK_RECORDS.pop(oldest, None)
        TASK_EVENTS_SENT.pop(oldest, None)


def _record_task(task_id: str, **updates: Any) -> None:
    with RECORDS_LOCK:
        record = TASK_RECORDS.setdefault(task_id, {})
        record.update(updates)
        record.setdefault("created_local_at", _now())
        record["updated_local_at"] = _now()
        _evict_old_tasks_locked()


def _get_task_record(task_id: str) -> Dict[str, Any]:
    with RECORDS_LOCK:
        return dict(TASK_RECORDS.get(task_id, {}))


def _event_sent_once(task_id: str, event: str) -> bool:
    with RECORDS_LOCK:
        sent = TASK_EVENTS_SENT.setdefault(task_id, set())
        if event in sent:
            return False
        sent.add(event)
        return True


def _is_terminal(status: Optional[str]) -> bool:
    normalized = (status or "").upper()
    return normalized in SUCCESS_STATUSES or normalized in FAILED_STATUSES


def _is_success(status: Optional[str]) -> bool:
    return (status or "").upper() in SUCCESS_STATUSES


def _is_failed(status: Optional[str]) -> bool:
    return (status or "").upper() in FAILED_STATUSES


def _next_poll_interval(state: dict, *, reset: bool = False) -> float:
    """Steady polling interval so the API status updates regularly (no sparse backoff)."""
    if reset:
        state["poll_attempt"] = 0
    return max(1.0, min(POLL_INTERVAL_SECONDS, POLL_MAX_INTERVAL_SECONDS))


def _timer(value: Optional[float]):
    """Timer update: stop when value is None, otherwise run at that interval.

    A bare None returned to a gr.Timer output does NOT stop it (gradio treats it as
    no-op), so the timer must be toggled explicitly via `active`.
    """
    if value is None:
        return gr.Timer(active=False)
    return gr.Timer(value=value, active=True)


# -----------------------------
# Analytics
# -----------------------------


def _get_posthog_client():
    global POSTHOG_CLIENT
    if not POSTHOG_PROJECT_API_KEY or Posthog is None:
        return None
    with POSTHOG_LOCK:
        if POSTHOG_CLIENT is None:
            POSTHOG_CLIENT = Posthog(
                POSTHOG_PROJECT_API_KEY,
                host=POSTHOG_HOST,
                sync_mode=False,
                timeout=2,
            )
        return POSTHOG_CLIENT


def track(event: str, state: Optional[dict] = None, properties: Optional[dict] = None) -> None:
    """Best-effort PostHog tracking. Analytics errors are intentionally ignored."""
    state = _ensure_browser_state(state)
    client = _get_posthog_client()
    if client is None:
        return

    props = {
        "app": "banuba_talking_avatar_space",
        "source": "hf_space",
        **(properties or {}),
    }
    try:
        client.capture(event, distinct_id=state["distinct_id"], properties=props)
    except Exception as exc:  # pragma: no cover
        print(f"[analytics] failed to capture {event}: {exc}", flush=True)


def track_task_once(event: str, task_id: Optional[str], state: Optional[dict], properties: Optional[dict] = None) -> None:
    if not task_id:
        track(event, state, properties)
        return
    if _event_sent_once(task_id, event):
        track(event, state, {"task_id": task_id, **(properties or {})})


# -----------------------------
# Banuba API client
# -----------------------------


class BanubaAPIError(RuntimeError):
    pass


class BanubaClient:
    def __init__(self) -> None:
        self.session = requests.Session()

    def _require_credentials(self) -> None:
        if not CLIENT_ID or not CLIENT_SECRET:
            raise BanubaAPIError(
                "Missing BANUBA_CLIENT_ID or BANUBA_CLIENT_SECRET in Space secrets."
            )

    def access_token(self) -> str:
        self._require_credentials()
        with TOKEN_LOCK:
            if TOKEN_CACHE["access_token"] and TOKEN_CACHE["expires_at"] > _now() + 60:
                return str(TOKEN_CACHE["access_token"])

            response = self.session.post(
                f"{AUTH_BASE_URL}/token",
                headers={"Content-Type": "application/x-www-form-urlencoded"},
                data={
                    "grant_type": "client_credentials",
                    "scope": BANUBA_SCOPE,
                    "client_id": CLIENT_ID,
                    "client_secret": CLIENT_SECRET,
                },
                timeout=30,
            )
            if not response.ok:
                raise BanubaAPIError(f"Token request failed: {response.status_code} {response.text[:500]}")
            payload = response.json()
            access_token = payload.get("access_token")
            if not access_token:
                raise BanubaAPIError("Token response did not include access_token.")
            expires_in = int(payload.get("expires_in", 3600))
            TOKEN_CACHE["access_token"] = access_token
            TOKEN_CACHE["expires_at"] = _now() + max(expires_in - 30, 30)
            return str(access_token)

    def _headers(self) -> Dict[str, str]:
        return {"Authorization": f"Bearer {self.access_token()}"}

    def request_upload(self, file_path: str) -> Dict[str, str]:
        content_type = _guess_content_type(file_path)
        body = {"filename": Path(file_path).name, "content_type": content_type}
        response = self.session.post(
            f"{API_BASE_URL}/upload",
            headers={**self._headers(), "Content-Type": "application/json"},
            json=body,
            timeout=30,
        )
        if not response.ok:
            raise BanubaAPIError(f"Upload location request failed: {response.status_code} {response.text[:500]}")
        payload = response.json()
        if not payload.get("url") or not payload.get("upload_url"):
            raise BanubaAPIError("Upload location response did not include url/upload_url.")
        return {"url": payload["url"], "upload_url": payload["upload_url"], "content_type": content_type}

    def upload_file_to_presigned_url(self, file_path: str, upload_url: str, content_type: str) -> None:
        with open(file_path, "rb") as file_obj:
            response = self.session.put(
                upload_url,
                data=file_obj,
                headers={"Content-Type": content_type} if content_type else None,
                timeout=600,
            )
        if not response.ok:
            raise BanubaAPIError(f"Presigned upload failed: {response.status_code} {response.text[:500]}")

    def upload_asset(self, file_path: str) -> str:
        location = self.request_upload(file_path)
        self.upload_file_to_presigned_url(file_path, location["upload_url"], location["content_type"])
        return location["url"]

    def create_lipsync_task(
        self,
        image_url: str,
        audio_url: str,
        *,
        seconds: Optional[int] = DEFAULT_SECONDS,
        seed: Optional[int] = DEFAULT_SEED,
        prompt: Optional[str] = DEFAULT_PROMPT,
    ) -> Dict[str, Any]:
        task_input: Dict[str, Any] = {
            "assets": {
                "image": {"url": image_url},
                "audio": {"url": audio_url},
            }
        }
        if seconds is not None:
            task_input["seconds"] = _safe_int(seconds)
        if seed is not None:
            task_input["seed"] = _safe_int(seed)
        if prompt:
            task_input["prompt"] = prompt

        body = {"type": "video.lipsync", "input": task_input}
        response = self.session.post(
            f"{API_BASE_URL}/tasks",
            headers={**self._headers(), "Content-Type": "application/json"},
            json=body,
            timeout=30,
        )
        if not response.ok:
            raise BanubaAPIError(f"Create task failed: {response.status_code} {response.text[:500]}")
        payload = response.json()
        if not payload.get("id"):
            raise BanubaAPIError("Create task response did not include id.")
        return payload

    def get_task(self, task_id: str) -> Dict[str, Any]:
        response = self.session.get(
            f"{API_BASE_URL}/tasks/{task_id}",
            headers=self._headers(),
            timeout=30,
        )
        if not response.ok:
            raise BanubaAPIError(f"Task status request failed: {response.status_code} {response.text[:500]}")
        return response.json()

    def download_result(self, result_url: str, task_id: str, content_type: Optional[str] = None) -> str:
        ext = ".mp4"
        if content_type and "webm" in content_type:
            ext = ".webm"
        elif content_type and "quicktime" in content_type:
            ext = ".mov"
        out_path = RESULT_DIR / f"{task_id}{ext}"
        if out_path.exists() and out_path.stat().st_size > 0:
            return str(out_path)

        with self.session.get(result_url, stream=True, timeout=600) as response:
            if not response.ok:
                raise BanubaAPIError(f"Result download failed: {response.status_code} {response.text[:500]}")
            tmp_path = out_path.with_suffix(out_path.suffix + ".tmp")
            with open(tmp_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        f.write(chunk)
            tmp_path.replace(out_path)
        _prune_result_dir()
        return str(out_path)


BANUBA = BanubaClient()


# -----------------------------
# Examples
# -----------------------------


def _discover_examples() -> list[list[str]]:
    image_exts = {".jpg", ".jpeg", ".png", ".webp"}
    audio_exts = [".wav", ".mp3", ".m4a", ".aac", ".ogg", ".flac"]
    root = Path(EXAMPLES_DIR)
    if not root.exists():
        return []

    examples: list[list[str]] = []
    for image_path in sorted(root.rglob("*")):
        if image_path.suffix.lower() not in image_exts:
            continue
        for ext in audio_exts:
            audio_path = image_path.with_suffix(ext)
            if audio_path.exists():
                examples.append([str(image_path), str(audio_path)])
                break
    return examples


def _path_is_example(path: Optional[str]) -> bool:
    if not path:
        return False
    try:
        Path(path).resolve().relative_to(Path(EXAMPLES_DIR).resolve())
        return True
    except Exception:
        return False


def _example_key(image_path: Optional[str], audio_path: Optional[str]) -> Optional[str]:
    if not image_path or not audio_path:
        return None
    if not (_path_is_example(image_path) or _path_is_example(audio_path)):
        return None
    raw = f"{Path(image_path).name}|{Path(audio_path).name}"
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]


def _example_output_path(image_path: Optional[str]) -> Optional[str]:
    """Pre-rendered sample video for an example, matched by filename stem.

    Gradio may pass a cached *copy* of the example image (not the original path),
    so we resolve the sample by stem inside EXAMPLES_DIR rather than as a sibling.
    """
    if not image_path:
        return None
    stem = Path(image_path).stem
    root = Path(EXAMPLES_DIR)
    for ext in (".mp4", ".webm", ".mov"):
        sibling = Path(image_path).with_suffix(ext)
        if sibling.exists():
            return str(sibling)
        if root.exists():
            for candidate in root.rglob(f"{stem}{ext}"):
                if candidate.is_file():
                    return str(candidate)
    return None


EXAMPLES = _discover_examples()


# -----------------------------
# Gradio event handlers
# -----------------------------


def on_space_load(browser_state: Optional[dict], request: gr.Request):
    state = _ensure_browser_state(browser_state)
    task_id = state.get("task_id")
    track("space_viewed", state, {"has_saved_task": bool(task_id)})

    if task_id:
        _set_session_task(request, task_id)
        status, output_html, video_update, button_update, timer_value = _poll_task_core(task_id, state)
        return status, output_html, video_update, task_id, state, button_update, _timer(timer_value)

    return (
        STATUS_EMPTY,
        OUTPUT_EMPTY_HTML,
        gr.update(value=None, visible=False),
        "",
        state,
        _btn(True),
        _timer(None),
    )


def on_reference_image_uploaded(image_path: Optional[str], browser_state: Optional[dict]):
    state = _ensure_browser_state(browser_state)
    if image_path and not _path_is_example(image_path):
        track("reference_image_uploaded", state, _file_properties(image_path))
    return state


def on_reference_audio_uploaded(audio_path: Optional[str], browser_state: Optional[dict]):
    state = _ensure_browser_state(browser_state)
    if audio_path and not _path_is_example(audio_path):
        track("reference_audio_uploaded", state, _file_properties(audio_path))
    return state


def maybe_track_example_selected(
    image_path: Optional[str],
    audio_path: Optional[str],
    browser_state: Optional[dict],
):
    state = _ensure_browser_state(browser_state)
    key = _example_key(image_path, audio_path)
    if key and state.get("last_example_selected") != key:
        state["last_example_selected"] = key
        track(
            "example_selected",
            state,
            {
                "example_key": key,
                "image_filename": Path(image_path).name if image_path else None,
                "audio_filename": Path(audio_path).name if audio_path else None,
            },
        )
    return state


def show_example_output(image_path: Optional[str], audio_path: Optional[str]):
    """When an example is clicked, preview its pre-rendered sample video (if one is shipped)."""
    sample = _example_output_path(image_path)
    if sample:
        return gr.update(value=sample, visible=True), "", STATUS_SAMPLE
    # No sample shipped for this example: leave the output area untouched.
    return gr.update(), gr.update(), gr.update()


def start_generation(
    image_path: Optional[str],
    audio_path: Optional[str],
    prompt: Optional[str],
    max_seconds: Optional[float],
    browser_state: Optional[dict],
    request: gr.Request,
):
    state = _ensure_browser_state(browser_state)
    prompt = (prompt or "").strip() or DEFAULT_PROMPT
    try:
        seconds = int(max_seconds)
    except (TypeError, ValueError):
        seconds = DEFAULT_SECONDS or 4
    seconds = max(1, min(seconds, 10))

    yield (
        STATUS_STARTING,
        OUTPUT_LOADING_HTML,
        gr.update(value=None, visible=False),
        state.get("task_id", ""),
        state,
        _btn(False),
        _timer(None),
    )

    if not image_path or not audio_path:
        track("generation_failed", state, {"error": "missing_required_assets"})
        yield (
            STATUS_FAILED + "\n\nBoth reference image and reference audio are required.",
            OUTPUT_FAILED_HTML,
            gr.update(value=None, visible=False),
            state.get("task_id", ""),
            state,
            _btn(True),
            _timer(None),
        )
        return

    try:
        image_asset_url = BANUBA.upload_asset(image_path)
        audio_asset_url = BANUBA.upload_asset(audio_path)
        task = BANUBA.create_lipsync_task(
            image_asset_url,
            audio_asset_url,
            seconds=seconds,
            seed=DEFAULT_SEED,
            prompt=prompt,
        )
        task_id = str(task["id"])
        state.update(
            {
                "task_id": task_id,
                "task_started_at": _now(),
                "task_status": "PENDING",
                "poll_attempt": 0,
            }
        )
        _set_session_task(request, task_id)
        _record_task(
            task_id,
            status="PENDING",
            started_at=state["task_started_at"],
            seconds=seconds,
            seed=DEFAULT_SEED,
            distinct_id=state["distinct_id"],
            prompt_hash=hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:16] if prompt else None,
        )
        track_task_once(
            "generation_started",
            task_id,
            state,
            {
                "seconds": seconds,
                "seed": DEFAULT_SEED,
                "prompt_chars": len(prompt or ""),
                "prompt_is_default": prompt == DEFAULT_PROMPT,
                "image_content_type": _guess_content_type(image_path),
                "audio_content_type": _guess_content_type(audio_path),
            },
        )
        yield (
            _status_with_task(STATUS_GENERATING, task_id, "IN_PROGRESS"),
            OUTPUT_LOADING_HTML,
            gr.update(value=None, visible=False),
            task_id,
            state,
            _btn(False),
            _timer(_next_poll_interval(state, reset=True)),
        )
    except Exception as exc:
        error_message = str(exc)
        track("generation_failed", state, {"error": error_message[:500]})
        print(f"[generation] failed: {error_message}", flush=True)
        yield (
            STATUS_FAILED + f"\n\n`{error_message[:300]}`",
            OUTPUT_FAILED_HTML,
            gr.update(value=None, visible=False),
            state.get("task_id", ""),
            state,
            _btn(True),
            _timer(None),
        )


def _poll_task_core(task_id: Optional[str], state: dict):
    if not task_id:
        return (
            STATUS_EMPTY,
            OUTPUT_EMPTY_HTML,
            gr.update(value=None, visible=False),
            _btn(True),
            None,
        )

    try:
        task = BANUBA.get_task(task_id)
        status = str(task.get("status", "UNKNOWN")).upper()
        state["task_status"] = status
        _record_task(task_id, status=status, last_api_payload=task)

        if _is_success(status):
            result = task.get("result") or {}
            result_url = result.get("result_file_url")
            content_type = result.get("content_type")
            if not result_url:
                raise BanubaAPIError("Task is completed but result.result_file_url is missing.")
            video_path = BANUBA.download_result(result_url, task_id, content_type)
            state["result_video_path"] = video_path
            state["task_status"] = status
            _record_task(
                task_id,
                status=status,
                result_video_path=video_path,
                content_type=content_type,
                completed_at=_now(),
            )

            started_at = state.get("task_started_at") or _get_task_record(task_id).get("started_at")
            elapsed = round(_now() - float(started_at), 2) if started_at else None
            track_task_once(
                "generation_completed",
                task_id,
                state,
                {"status": status, "elapsed_seconds": elapsed, "content_type": content_type},
            )
            return (
                _status_with_task(STATUS_DONE, task_id, status),
                "",
                gr.update(value=video_path, visible=True),
                _btn(True),
                None,
            )

        if _is_failed(status):
            _record_task(task_id, status=status, failed_at=_now())
            error_detail = None
            if isinstance(task, dict):
                error_detail = task.get("error") or task.get("error_message") or task.get("message")
            track_task_once(
                "generation_failed",
                task_id,
                state,
                {"status": status, "error": str(error_detail)[:500] if error_detail else None},
            )
            return (
                _status_with_task(STATUS_FAILED, task_id, status),
                OUTPUT_FAILED_HTML,
                gr.update(value=None, visible=False),
                _btn(True),
                None,
            )

        return (
            _status_with_task(STATUS_GENERATING, task_id, "IN_PROGRESS"),
            OUTPUT_LOADING_HTML,
            gr.update(value=None, visible=False),
            _btn(False),
            _next_poll_interval(state),
        )
    except Exception as exc:
        error_message = str(exc)
        print(f"[poll] failed for task {task_id}: {error_message}", flush=True)
        # Polling errors are usually transient; keep polling unless the API returns a terminal failure.
        return (
            _status_with_task(STATUS_GENERATING + f"\n\nTemporary polling issue: `{error_message[:250]}`", task_id),
            OUTPUT_LOADING_HTML,
            gr.update(value=None, visible=False),
            _btn(False),
            _next_poll_interval(state),
        )


def poll_generation(task_id: Optional[str], browser_state: Optional[dict], request: gr.Request):
    state = _ensure_browser_state(browser_state)
    task_id = task_id or state.get("task_id")
    if task_id:
        state["task_id"] = task_id
        _set_session_task(request, task_id)
    status, output_html, video_update, button_update, timer_value = _poll_task_core(task_id, state)
    return status, output_html, video_update, task_id or "", state, button_update, _timer(timer_value)


def on_result_viewed(task_id: Optional[str], browser_state: Optional[dict]):
    """Fires when the user actually plays the generated video (distinct from completion)."""
    state = _ensure_browser_state(browser_state)
    task_id = task_id or state.get("task_id")
    content_type = _get_task_record(task_id).get("content_type") if task_id else None
    track_task_once("result_viewed", task_id, state, {"content_type": content_type})
    return state


def mark_abandoned_on_unload(request: gr.Request):
    session_hash = _request_session_hash(request)
    task_id = SESSION_TASKS.pop(session_hash, None)
    if not task_id:
        return
    record = _get_task_record(task_id)
    status = str(record.get("status", "UNKNOWN")).upper()
    if not _is_terminal(status):
        # BrowserState is unavailable during unload, so reuse the distinct_id captured at
        # generation start to keep the funnel joined; fall back to the session hash.
        distinct_id = record.get("distinct_id") or f"session_{session_hash}"
        fallback_state = {"distinct_id": distinct_id}
        track_task_once("generation_abandoned", task_id, fallback_state, {"status_at_unload": status})


# -----------------------------
# UI
# -----------------------------


def make_browser_state_component():
    browser_state_cls = getattr(gr, "BrowserState", gr.State)
    if browser_state_cls is gr.State:
        return browser_state_cls({})
    secret = os.getenv("BROWSER_STATE_SECRET", "banuba-talking-avatar-v1")
    try:
        return browser_state_cls({}, storage_key="banuba_talking_avatar_state_v1", secret=secret)
    except TypeError:
        # Older BrowserState versions accept only the default value.
        return browser_state_cls({})


with gr.Blocks(title="AI Talking Photo", analytics_enabled=False) as demo:
    client_state = make_browser_state_component()
    task_id_state = gr.State("")
    poll_timer = gr.Timer(value=None, active=True)

    gr.HTML(
        """
        <div class="hero">
          <h1>AI Talking Photo</h1>
          <p>Banuba AI Talking Photo turns a single photo into a lifelike talking avatar with natural lip-sync, facial expressions, and motion — studio-quality presentations, lessons, and videos without cameras, actors, or editing. Supports any language.</p>
          <p class="hero-note">This is a free demo: output is <strong>480p</strong> and capped at <strong>1–10 seconds</strong> (default 4s). Upload an image and audio, set the maximum duration, and generate. For full-length (up to ~1 hour) and higher-resolution output, contact <a href="mailto:support@banuba.com">support@banuba.com</a> or <a href="https://www.banuba.com" target="_blank" rel="noopener">banuba.com</a>.</p>
        </div>
        """
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            reference_image = gr.Image(
                label="Reference image input",
                type="filepath",
                sources=["upload"],
                image_mode="RGB",
                height=320,
            )
            reference_audio = gr.Audio(
                label="Reference audio input",
                type="filepath",
                sources=["upload"],
            )

            prompt_input = gr.Textbox(
                label="Prompt (optional)",
                value=DEFAULT_PROMPT,
                lines=3,
                info="Guides delivery, framing, gestures, lighting and background. Leave as-is to use the default.",
            )

            seconds_input = gr.Number(
                label="Maximum duration (seconds)",
                value=DEFAULT_SECONDS if DEFAULT_SECONDS is not None else 4,
                minimum=1,
                maximum=10,
                step=1,
                precision=0,
                info="Length of the generated video (1–10s). Trim the audio in the player above if needed.",
            )

            generate_btn = gr.Button(BTN_GENERATE, variant="primary", size="lg", elem_id="generate-btn")

        with gr.Column(scale=1):
            generation_status = gr.Markdown(STATUS_EMPTY)
            output_state = gr.HTML(OUTPUT_EMPTY_HTML)
            output_video = gr.Video(
                label="Output",
                format="mp4",
                autoplay=False,
                visible=False,
                height=420,
            )

            gr.Markdown("### Examples")
            if EXAMPLES:
                gr.Examples(
                    examples=EXAMPLES,
                    inputs=[reference_image, reference_audio],
                    outputs=[output_video, output_state, generation_status],
                    fn=show_example_output,
                    run_on_click=True,
                    cache_examples=False,
                    examples_per_page=6,
                )
            else:
                gr.Markdown(
                    "_No packaged examples yet. Add paired files like `assets/examples/example_1.png` "
                    "and `assets/examples/example_1.wav` to show clickable examples here._"
                )

    if GA_LOAD_JS:
        # Run the Google Analytics loader on page load. demo.load(None, js=...) is the
        # documented, battle-tested way to execute JS on load (same trick used to toggle
        # dark mode); Blocks(js=...) proved unreliable to auto-fire under ssr_mode=False.
        demo.load(None, inputs=None, outputs=None, js=GA_LOAD_JS)

    demo.load(
        on_space_load,
        inputs=[client_state],
        outputs=[generation_status, output_state, output_video, task_id_state, client_state, generate_btn, poll_timer],
        api_visibility="private",
    )

    if hasattr(reference_image, "upload"):
        reference_image.upload(
            on_reference_image_uploaded,
            inputs=[reference_image, client_state],
            outputs=[client_state],
            api_visibility="private",
        )
    else:
        reference_image.change(
            on_reference_image_uploaded,
            inputs=[reference_image, client_state],
            outputs=[client_state],
            api_visibility="private",
        )

    if hasattr(reference_audio, "upload"):
        reference_audio.upload(
            on_reference_audio_uploaded,
            inputs=[reference_audio, client_state],
            outputs=[client_state],
            api_visibility="private",
        )
    else:
        reference_audio.change(
            on_reference_audio_uploaded,
            inputs=[reference_audio, client_state],
            outputs=[client_state],
            api_visibility="private",
        )

    reference_image.change(
        maybe_track_example_selected,
        inputs=[reference_image, reference_audio, client_state],
        outputs=[client_state],
        api_visibility="private",
    )
    reference_audio.change(
        maybe_track_example_selected,
        inputs=[reference_image, reference_audio, client_state],
        outputs=[client_state],
        api_visibility="private",
    )

    # Disable the button the instant it is clicked (queue=False runs immediately,
    # before start_generation is enqueued) so rapid clicks cannot queue duplicate jobs.
    generate_btn.click(
        lambda: (_btn(False), STATUS_STARTING, OUTPUT_LOADING_HTML),
        outputs=[generate_btn, generation_status, output_state],
        js=GENERATE_CLICK_JS,
        queue=False,
        api_visibility="private",
    ).then(
        start_generation,
        inputs=[reference_image, reference_audio, prompt_input, seconds_input, client_state],
        outputs=[generation_status, output_state, output_video, task_id_state, client_state, generate_btn, poll_timer],
        api_name="start_generation",
        api_description="Upload image/audio assets, create a Banuba video.lipsync task, and start polling.",
        show_progress="minimal",
    )

    poll_timer.tick(
        poll_generation,
        inputs=[task_id_state, client_state],
        outputs=[generation_status, output_state, output_video, task_id_state, client_state, generate_btn, poll_timer],
        api_name="poll_generation",
        api_description="Poll a Banuba task id and return the current generation state/result video.",
        show_progress="hidden",
        concurrency_limit=None,
    )

    if hasattr(output_video, "play"):
        output_video.play(
            on_result_viewed,
            inputs=[task_id_state, client_state],
            outputs=[client_state],
            api_visibility="private",
        )

    demo.unload(mark_abandoned_on_unload)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=16, max_size=64).launch(
        show_error=True,
        css=CUSTOM_CSS,
        allowed_paths=[str(RESULT_DIR), str(Path(EXAMPLES_DIR).resolve())],
        # Experimental SSR creates duplicate per-session SSE lifecycles, which trips a
        # gradio-internal KeyError on stream teardown. Disable it for stability.
        ssr_mode=False,
    )