app.py

"""
Avatar Chatbot - HuggingFace Spaces Edition
Avatar profile system: each avatar has ref.png, persona.txt, idlevideos/
Chunked pipeline: LLM (once) → split sentences → TTS+FLOAT per chunk → stream to frontend
"""
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from pathlib import Path
import mimetypes
import os
import re
import logging
import subprocess
import time as _time
import uuid
import cv2
import asyncio

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# ---- Avatar system ----
AVATARS_DIR = Path("/app/avatars")
_current_avatar = None  # name of active avatar folder

TTS_DIR = Path("/tmp/tts_output")
LIPSYNC_DIR = Path("/tmp/lipsync_output")
REPO_ID = os.environ.get("SPACE_ID", "ffbeqwbe/AvatarChatbot")
FLOAT_ENABLED = os.environ.get("FLOAT_ENABLED", "true").lower() == "true"

for d in [TTS_DIR, LIPSYNC_DIR]:
    d.mkdir(parents=True, exist_ok=True)


def _ts():
    return _time.strftime("%H:%M:%S", _time.gmtime()) + f".{int(_time.time()*1000)%1000:03d}"


def _get_avatar_dir(name=None):
    """Get path to an avatar's directory."""
    if name is None:
        name = _current_avatar
    if name:
        return AVATARS_DIR / name
    return None


def _get_idle_videos_dir():
    """Get the current avatar's idle videos directory."""
    d = _get_avatar_dir()
    if d:
        return d / "idlevideos"
    return Path("/tmp/videos")


def _list_idle_videos():
    """List idle video files for the current avatar."""
    d = _get_idle_videos_dir()
    if not d.exists():
        return []
    return [
        {"name": f.stem, "filename": f.name}
        for f in sorted(d.iterdir())
        if f.suffix.lower() in (".mp4", ".webm", ".mkv", ".mov") and f.stat().st_size > 1024
    ]


import threading
_float_lock = threading.Lock()  # Only for FLOAT GPU inference
_sessions = {}
GREETING_DATA = None


def _discover_avatars():
    """Find all avatar profiles in the avatars directory."""
    if not AVATARS_DIR.exists():
        logger.warning(f"[AVATARS] Directory not found: {AVATARS_DIR}")
        return []
    avatars = []
    for d in sorted(AVATARS_DIR.iterdir()):
        if d.is_dir():
            has_ref = (d / "ref.png").exists() or (d / "ref.jpg").exists()
            has_persona = (d / "persona.txt").exists()
            has_videos = (d / "idlevideos").exists()
            avatars.append({
                "name": d.name,
                "has_ref": has_ref,
                "has_persona": has_persona,
                "has_videos": has_videos,
            })
    return avatars


def _parse_avatar_config(avatar_dir):
    """Parse config.txt from avatar folder. Returns dict of key=value pairs."""
    config = {}
    config_path = avatar_dir / "config.txt"
    if config_path.exists():
        for line in open(config_path).readlines():
            line = line.strip()
            if '=' in line and not line.startswith('#'):
                key, val = line.split('=', 1)
                config[key.strip().lower()] = val.strip()
    return config


def _resolve_voice(config: dict) -> str:
    """
    Pick the TTS voice from an avatar config.
    Precedence (first wins):
      1. cartesiavoicename → Cartesia voice UUID (e.g. "e07c00bc-4134-...")
      2. edgettsname       → raw edge-tts voice ID (e.g. "en-GB-RyanNeural")
      3. groqttsname       → legacy short name, mapped via VOICE_MAP in tts.py
      4. fallback default
    set_voice() in tts.py detects the format and routes to the right backend:
    UUID → Cartesia API, else → edge-tts (raw or via VOICE_MAP).
    """
    return (
        config.get("cartesiavoicename")
        or config.get("edgettsname")
        or config.get("groqttsname")
        or "diana"
    )


def _switch_avatar(avatar_name: str):
    """Switch to a different avatar profile."""
    global _current_avatar
    avatar_dir = AVATARS_DIR / avatar_name
    if not avatar_dir.exists():
        logger.error(f"[AVATARS] Avatar not found: {avatar_name}")
        return False

    _current_avatar = avatar_name
    logger.info(f"[AVATARS] Switching to avatar: {avatar_name}")

    # 0. Parse config.txt
    config = _parse_avatar_config(avatar_dir)

    # 1. Load persona into LLM
    persona_path = avatar_dir / "persona.txt"
    from llm import load_persona_from_file
    persona_text = load_persona_from_file(str(persona_path))
    logger.info(f"[AVATARS] Persona loaded: {persona_text[:80]}..." if persona_text else "[AVATARS] No persona.txt found")

    # 2. Set TTS voice from config (default: diana)
    tts_voice = _resolve_voice(config)
    from tts import set_voice
    set_voice(tts_voice)

    # 3. Update FLOAT reference image
    if FLOAT_ENABLED:
        try:
            from float_lipsync import get_lipsync
            lipsync = get_lipsync()
            if lipsync.ready:
                ref_path = avatar_dir / "ref.png"
                if not ref_path.exists():
                    ref_path = avatar_dir / "ref.jpg"
                if ref_path.exists():
                    lipsync.update_reference_image(str(ref_path))
                    logger.info(f"[AVATARS] Reference image updated: {ref_path}")

                    # Also update streamer
                    try:
                        from float_streamer import get_streamer
                        streamer = get_streamer()
                        if streamer.ready:
                            streamer.drain_buffer()
                            streamer.update_reference(lipsync.preprocessed_ref_image)
                            logger.info(f"[AVATARS] Streamer reference updated")
                    except Exception as e2:
                        logger.error(f"[AVATARS] Streamer update failed: {e2}")
                else:
                    logger.warning(f"[AVATARS] No ref image in {avatar_dir}")
        except Exception as e:
            logger.error(f"[AVATARS] Failed to update ref image: {e}")

    # 3. Idle videos — no action needed, _get_idle_videos_dir() uses _current_avatar dynamically

    logger.info(f"[AVATARS] ✓ Switched to {avatar_name}")
    return True


@app.on_event("startup")
async def startup():
    # Discover avatars and set default
    avatars = _discover_avatars()
    logger.info(f"[STARTUP] Found {len(avatars)} avatars: {[a['name'] for a in avatars]}")

    if avatars:
        # Use first avatar as default
        default = avatars[0]["name"]
        # Load persona for default avatar (before FLOAT init so ref image is set)
        global _current_avatar
        _current_avatar = default
        persona_path = AVATARS_DIR / default / "persona.txt"
        from llm import load_persona_from_file
        load_persona_from_file(str(persona_path))
        # Load config (TTS voice etc.)
        config = _parse_avatar_config(AVATARS_DIR / default)
        tts_voice = _resolve_voice(config)
        logger.info(f"[STARTUP] Default avatar: {default} | voice: {tts_voice}")

    # Run independent inits in parallel — saves ~1s
    init_tasks = [init_tts(), init_stt()]
    if FLOAT_ENABLED:
        init_tasks.append(init_float())
    await asyncio.gather(*init_tasks)

    # Apply default avatar's TTS voice (after TTS init)
    if avatars:
        config = _parse_avatar_config(AVATARS_DIR / _current_avatar)
        tts_voice = _resolve_voice(config)
        from tts import set_voice
        set_voice(tts_voice)

    # Streamer depends on FLOAT being ready
    if FLOAT_ENABLED:
        await init_streamer()

    # Fire greeting as background task — DON'T block "Application startup complete"
    # First wav2vec/FLOAT inference triggers CUDA kernel compilation (~10s cold start).
    # Running it in background means the app serves immediately; greeting plays ~10s later.
    asyncio.create_task(generate_greeting())
    logger.info(f"[{_ts()}] [STARTUP] ✓ App ready — greeting generating in background")


async def init_tts():
    logger.info(f"[{_ts()}] [STARTUP] Initializing TTS...")
    t0 = _time.time()
    try:
        import tts as tts_module
        tts_module.initialize()
        logger.info(f"[{_ts()}] [STARTUP] ✓ TTS ready in {_time.time()-t0:.2f}s")
    except Exception as e:
        logger.error(f"[{_ts()}] [STARTUP] TTS init failed: {type(e).__name__}: {e}", exc_info=True)


async def generate_greeting():
    global GREETING_DATA
    greeting_text = "Hello! Feel free to ask me anything."
    logger.info(f"[{_ts()}] [GREETING] Generating warmup greeting...")
    t0 = _time.time()

    try:
        from tts import generate_audio
        # All blocking calls wrapped in to_thread — event loop stays responsive
        audio_path = await asyncio.to_thread(generate_audio, greeting_text)
        audio_url = None
        audio_duration = 0

        if audio_path and os.path.exists(audio_path):
            audio_url = f"/api/audio/{os.path.basename(audio_path)}"
            try:
                r = await asyncio.to_thread(
                    subprocess.run,
                    ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
                     "-of", "csv=p=0", audio_path],
                    capture_output=True, text=True, check=True
                )
                audio_duration = float(r.stdout.strip())
            except Exception:
                audio_duration = 5.0

        # Inject greeting into streamer — CUDA kernel compile happens here (~10s first call)
        if FLOAT_ENABLED and audio_path:
            try:
                from float_streamer import get_streamer
                streamer = get_streamer()
                if streamer.ready:
                    await asyncio.to_thread(streamer.inject_speech, audio_path, audio_url)
                    logger.info(f"[{_ts()}] [GREETING] Injected into streamer")
            except Exception as e:
                logger.warning(f"[GREETING] Streamer inject failed: {e}")

        GREETING_DATA = {
            "text": greeting_text,
            "audio_url": audio_url,
            "audio_duration": audio_duration,
        }
        logger.info(f"[{_ts()}] [GREETING] Ready in {_time.time()-t0:.2f}s")

    except Exception as e:
        logger.error(f"[{_ts()}] [GREETING] Failed: {e}", exc_info=True)
        GREETING_DATA = {"text": greeting_text, "audio_url": None, "audio_duration": 0}


async def init_stt():
    """Initialize Vosk speech-to-text."""
    logger.info(f"[{_ts()}] [STARTUP] Initializing STT...")
    try:
        import vosk_stt
        if vosk_stt.initialize():
            logger.info(f"[{_ts()}] [STARTUP] ✓ STT ready")
        else:
            logger.warning(f"[{_ts()}] [STARTUP] STT init failed (voice input disabled)")
    except Exception as e:
        logger.warning(f"[{_ts()}] [STARTUP] STT not available: {e}")


async def init_float():
    import torch
    logger.info(f"[STARTUP] CUDA available: {torch.cuda.is_available()}")
    if not torch.cuda.is_available():
        logger.warning("[STARTUP] No GPU - FLOAT disabled"); return
    try:
        from float_lipsync import get_lipsync
        from huggingface_hub import hf_hub_download, snapshot_download
        import shutil

        ckpt = Path("/app/checkpoints")
        ckpt.mkdir(parents=True, exist_ok=True)

        w2v = ckpt / "wav2vec2-base-960h"
        if not w2v.exists():
            logger.info("[STARTUP] Downloading wav2vec2-base-960h...")
            snapshot_download(repo_id="facebook/wav2vec2-base-960h", local_dir=str(w2v))

        emo = ckpt / "wav2vec-english-speech-emotion-recognition"
        if not emo.exists():
            logger.info("[STARTUP] Downloading emotion model...")
            snapshot_download(repo_id="r-f/wav2vec-english-speech-emotion-recognition", local_dir=str(emo))

        fp = ckpt / "float.pth"
        if not fp.exists() or fp.stat().st_size < 1024:
            logger.info("[STARTUP] Downloading float.pth...")
            dl = hf_hub_download(repo_id=REPO_ID, repo_type="space",
                                 filename="app/checkpoints/float.pth", local_dir="/tmp/hf_download")
            shutil.copy2(dl, fp)
            logger.info(f"[STARTUP] float.pth: {fp.stat().st_size/1e6:.1f} MB")

        # Use current avatar's ref.png for initial FLOAT setup
        ref_path = None
        avatar_dir = _get_avatar_dir()
        if avatar_dir:
            for name in ["ref.png", "ref.jpg"]:
                p = avatar_dir / name
                if p.exists():
                    ref_path = str(p)
                    break

        # Fallback to legacy location
        if not ref_path:
            assets = Path("/app/assets")
            assets.mkdir(parents=True, exist_ok=True)
            ref = assets / "ref.png"
            if not ref.exists() or ref.stat().st_size < 1024:
                for name in ["ref.png", "ref.jpg", "main2.png"]:
                    try:
                        dl = hf_hub_download(repo_id=REPO_ID, repo_type="space",
                                             filename=f"app/assets/{name}", local_dir="/tmp/hf_download")
                        shutil.copy2(dl, assets / name)
                        ref = assets / name
                        break
                    except Exception:
                        continue
                else:
                    logger.warning("[STARTUP] No ref image found"); return
            ref_path = str(ref)

        lipsync = get_lipsync()
        lipsync.initialize({
            "ref_path": ref_path, "ckpt_path": str(fp),
            "wav2vec_model_path": str(w2v), "audio2emotion_path": str(emo),
        })
        logger.info("[STARTUP] ✓ FLOAT ready!")
    except Exception as e:
        logger.error(f"[STARTUP] FLOAT init failed: {e}", exc_info=True)


async def init_streamer():
    """Initialize the FLOAT streaming engine after FLOAT model is loaded."""
    try:
        from float_lipsync import get_lipsync
        from float_streamer import get_streamer

        lipsync = get_lipsync()
        if not lipsync.ready:
            logger.warning("[STARTUP] FLOAT not ready — streamer disabled")
            return

        streamer = get_streamer()
        streamer.initialize(
            model=lipsync.model,
            device=lipsync.device,
            opt=lipsync.opt,
            ref_image_tensor=lipsync.preprocessed_ref_image,
            wav2vec_preprocessor=lipsync.wav2vec_preprocessor,
        )
        streamer.start()
        logger.info("[STARTUP] ✓ Streamer started!")
    except Exception as e:
        logger.error(f"[STARTUP] Streamer init failed: {e}", exc_info=True)


# ---- Avatar API endpoints ----

@app.get("/api/avatars")
async def list_avatars():
    """List all available avatar profiles."""
    avatars = _discover_avatars()
    return JSONResponse({
        "avatars": avatars,
        "current": _current_avatar,
    })


@app.post("/api/switch-avatar")
async def switch_avatar(request: Request):
    """Switch to a different avatar. Resets chat history."""
    body = await request.json()
    avatar_name = body.get("avatar", "").strip()
    if not avatar_name:
        raise HTTPException(400, "No avatar specified")

    success = _switch_avatar(avatar_name)
    if not success:
        raise HTTPException(404, f"Avatar not found: {avatar_name}")

    return JSONResponse({
        "status": "ok",
        "avatar": avatar_name,
        "videos": _list_idle_videos(),
    })


@app.get("/api/videos")
async def list_videos():
    vids = _list_idle_videos()
    return JSONResponse({"videos": vids})


# ---- Avatar creation ----
_creation_status = {"active": False, "progress": 0, "total": 0, "message": "", "avatar_name": "", "done": False, "error": None}


@app.post("/api/create-avatar")
async def create_avatar(request: Request):
    """Create a new avatar from uploaded image + settings. Kicks off idle video generation."""
    global _creation_status

    if _creation_status["active"]:
        raise HTTPException(409, "Avatar creation already in progress")

    import base64
    body = await request.json()
    name = body.get("name", "").strip()
    voice = body.get("voice", "diana").strip()
    persona = body.get("persona", "").strip()
    image_b64 = body.get("image", "")

    if not name:
        raise HTTPException(400, "Name is required")
    if not image_b64:
        raise HTTPException(400, "Image is required")

    safe_name = re.sub(r'[^\w\-]', '', name)
    if not safe_name:
        raise HTTPException(400, "Invalid name")

    avatar_dir = AVATARS_DIR / safe_name
    if avatar_dir.exists():
        raise HTTPException(409, f"Avatar '{safe_name}' already exists")

    avatar_dir.mkdir(parents=True, exist_ok=True)
    idle_dir = avatar_dir / "idlevideos"
    idle_dir.mkdir(exist_ok=True)

    try:
        if ',' in image_b64:
            image_b64 = image_b64.split(',', 1)[1]
        image_bytes = base64.b64decode(image_b64)
        ref_path = avatar_dir / "ref.png"

        import numpy as np
        nparr = np.frombuffer(image_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError("Could not decode image")
        cv2.imwrite(str(ref_path), img)
        logger.info(f"[CREATE] Saved ref image: {ref_path} ({img.shape})")
    except Exception as e:
        import shutil
        shutil.rmtree(str(avatar_dir), ignore_errors=True)
        raise HTTPException(400, f"Invalid image: {e}")

    with open(avatar_dir / "config.txt", "w") as f:
        # New avatars store the edge-tts voice ID directly.
        # Legacy avatars still use `groqttsname=` and are read via _resolve_voice().
        f.write(f"edgettsname={voice}\n")

    with open(avatar_dir / "persona.txt", "w") as f:
        f.write(persona if persona else f"You are {name}.")

    logger.info(f"[CREATE] Avatar profile created: {safe_name} | voice: {voice}")

    _creation_status = {
        "active": True, "progress": 0, "total": 6,
        "message": "Starting idle video generation...",
        "avatar_name": safe_name, "done": False, "error": None
    }

    import asyncio
    asyncio.get_event_loop().run_in_executor(None, _run_idle_generation, safe_name, str(ref_path), str(idle_dir))

    return JSONResponse({
        "status": "started",
        "avatar": safe_name,
        "message": "Avatar created. Generating idle videos (this takes a few minutes)..."
    })


def _run_idle_generation(avatar_name, ref_path, idle_dir):
    """Background task: generate idle clips and switch to new avatar when done."""
    global _creation_status
    try:
        from idle_generator import get_idle_generator
        generator = get_idle_generator()

        def on_progress(clip_idx, total, message):
            _creation_status["progress"] = clip_idx
            _creation_status["total"] = total
            _creation_status["message"] = message

        with _float_lock:
            generator.generate(
                ref_image_path=ref_path,
                output_dir=idle_dir,
                avatar_name=avatar_name,
                num_clips=6,
                progress_callback=on_progress,
            )

        _switch_avatar(avatar_name)

        _creation_status["done"] = True
        _creation_status["message"] = f"Avatar '{avatar_name}' ready!"
        logger.info(f"[CREATE] ✓ Avatar '{avatar_name}' fully created and activated")

    except Exception as e:
        _creation_status["error"] = str(e)
        _creation_status["message"] = f"Error: {e}"
        logger.error(f"[CREATE] Failed: {e}", exc_info=True)
    finally:
        _creation_status["active"] = False


@app.get("/api/create-avatar/status")
async def create_avatar_status():
    """Poll creation progress."""
    return JSONResponse(_creation_status)


@app.get("/api/stream/{filename}")
async def stream_video(filename: str):
    path = None
    # Search in current avatar's idle videos, lipsync output, and TTS output
    search_dirs = [_get_idle_videos_dir(), LIPSYNC_DIR, TTS_DIR]
    for d in search_dirs:
        if d is None:
            continue
        c = d / filename
        if c.exists() and c.is_file():
            path = c
            break
    if not path:
        raise HTTPException(404, "Not found")
    mime = mimetypes.guess_type(filename)[0] or "video/mp4"
    def gen():
        with open(path, "rb") as f:
            while chunk := f.read(1024*1024): yield chunk
    return StreamingResponse(gen(), media_type=mime,
        headers={"Content-Length": str(path.stat().st_size), "Accept-Ranges": "bytes"})


@app.get("/api/audio/{filename}")
async def stream_audio(filename: str):
    path = TTS_DIR / filename
    if not path.exists(): raise HTTPException(404, "Not found")
    mime = mimetypes.guess_type(filename)[0] or "audio/wav"
    def gen():
        with open(path, "rb") as f:
            while chunk := f.read(1024*1024): yield chunk
    return StreamingResponse(gen(), media_type=mime,
        headers={"Content-Length": str(path.stat().st_size)})


def split_sentences(text):
    """Split text into chunks for TTS+FLOAT processing.
    Each chunk should produce ~5-8s of audio so FLOAT finishes in ~3-4s."""
    parts = re.split(r'(?<=[.!?])\s+', text)
    parts = [s.strip() for s in parts if s.strip()]
    # Keep chunks short (~12 words) so FLOAT processes each in <5s
    MIN_WORDS = 12
    merged = []
    buffer = ""
    for p in parts:
        if buffer:
            buffer += " " + p
        else:
            buffer = p
        if len(buffer.split()) >= MIN_WORDS:
            merged.append(buffer)
            buffer = ""
    if buffer:
        if merged and len(buffer.split()) < 6:
            # Merge very short tail with previous chunk
            merged[-1] += " " + buffer
        else:
            merged.append(buffer)
    return merged if merged else [text]


@app.post("/api/chat")
async def chat(request: Request):
    t0 = _time.time()
    body = await request.json()
    user_text = body.get("text", "").strip()
    if not user_text: raise HTTPException(400, "No text")

    logger.info(f"[{_ts()}] [CHAT] User: {user_text}")

    t1 = _time.time()
    from llm import generate_response
    llm_result = generate_response(user_text)
    reply_text = llm_result['text']
    clean_text = llm_result.get('clean_text', reply_text)
    t2 = _time.time()
    logger.info(f"[{_ts()}] [CHAT] Reply: {reply_text}")
    logger.info(f"[{_ts()}] [TIMING] LLM: {t2-t1:.2f}s")

    sentences = split_sentences(reply_text)
    if not sentences:
        sentences = [reply_text]

    # Build clean sentences by stripping tags from each TTS sentence (ensures same count)
    clean_sentences = []
    for s in sentences:
        c = re.sub(r'\[[^\]]*\]\s*', '', s).strip()
        c = re.sub(r'\s{2,}', ' ', c)
        clean_sentences.append(c)
    
    logger.info(f"[{_ts()}] [CHAT] Split into {len(sentences)} chunks")

    session_id = str(uuid.uuid4())[:8]
    total = len(sentences)

    # Pre-generate ALL chunks in parallel
    import concurrent.futures

    results = [None] * total

    def generate_one(idx):
        sentence = sentences[idx]
        logger.info(f"[{_ts()}] [CHUNK {idx+1}/{total}] START | \"{sentence[:80]}\"")
        t_start = _time.time()

        # 1. TTS (Groq cloud API - no GPU, runs in parallel)
        from tts import generate_audio
        t_tts = _time.time()
        audio_path = generate_audio(sentence)
        t_tts_done = _time.time()

        audio_url = None
        audio_duration = 0
        audio_size = 0
        if audio_path and os.path.exists(audio_path):
            audio_size = os.path.getsize(audio_path)
            audio_url = f"/api/audio/{os.path.basename(audio_path)}"
            try:
                r = subprocess.run(["ffprobe","-v","quiet","-show_entries","format=duration",
                                    "-of","csv=p=0", audio_path], capture_output=True, text=True, check=True)
                audio_duration = float(r.stdout.strip())
            except Exception:
                audio_duration = 3.0

        logger.info(f"[{_ts()}] [CHUNK {idx+1}/{total}] TTS: {t_tts_done-t_tts:.2f}s | audio: {audio_duration:.1f}s, {audio_size/1024:.0f}KB")

        # 2. FLOAT lipsync
        lipsync_video_url = None
        if FLOAT_ENABLED and audio_path and os.path.exists(audio_path):
            try:
                from float_lipsync import get_lipsync
                lipsync = get_lipsync()
                if lipsync.ready:
                    with _float_lock:
                        t_float = _time.time()
                        lp = lipsync.generate(audio_path)
                        t_float_done = _time.time()
                    logger.info(f"[{_ts()}] [CHUNK {idx+1}/{total}] FLOAT: {t_float_done-t_float:.2f}s")
                    if lp and os.path.exists(lp):
                        lipsync_video_url = f"/api/stream/{os.path.basename(lp)}"
                        logger.info(f"[{_ts()}] [CHUNK {idx+1}/{total}] VIDEO: {os.path.getsize(lp)/1024:.0f}KB")
            except Exception as e:
                logger.error(f"[{_ts()}] [CHUNK {idx+1}/{total}] LIPSYNC FAILED: {e}", exc_info=True)

        total_time = _time.time() - t_start
        logger.info(f"[{_ts()}] [CHUNK {idx+1}/{total}] DONE | total: {total_time:.2f}s")

        results[idx] = {
            "audio_url": audio_url,
            "audio_duration": audio_duration,
            "lipsync_video_url": lipsync_video_url,
        }

    # Launch all chunks in parallel threads
    # TTS calls run truly in parallel; FLOAT serializes via _float_lock
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=total)
    futures = [executor.submit(generate_one, i) for i in range(total)]

    _sessions[session_id] = {
        "sentences": sentences,
        "clean_sentences": clean_sentences,
        "results": results,
        "futures": futures,
        "next_index": 0,
        "total": total,
        "done": False,
        "created": _time.time(),
    }

    logger.info(f"[{_ts()}] [TIMING] Chat response ready: {_time.time()-t0:.2f}s")

    return JSONResponse({
        "text": " ".join(clean_sentences),
        "tts_text": reply_text,
        "emotion": llm_result.get('emotion', 'neutral'),
        "session_id": session_id,
        "total_chunks": total,
    })


@app.post("/api/chat/next")
async def chat_next(request: Request):
    t_req = _time.time()
    body = await request.json()
    session_id = body.get("session_id", "")
    session = _sessions.get(session_id)
    if not session: raise HTTPException(404, "Session not found")

    idx = session["next_index"]
    if idx >= session["total"]:
        session["done"] = True
        return JSONResponse({"done": True})

    sentence = session["sentences"][idx]
    session["next_index"] = idx + 1
    is_last = (idx + 1 >= session["total"])
    total = session["total"]

    logger.info(f"[{_ts()}] [REQ] /api/chat/next chunk {idx+1}/{total} received")

    # Wait for this specific chunk to finish (already running in background)
    import asyncio
    future = session["futures"][idx]
    await asyncio.to_thread(future.result)

    result = session["results"][idx]
    if result is None:
        result = {"audio_url": None, "audio_duration": 0, "lipsync_video_url": None}

    t_resp = _time.time()
    logger.info(f"[{_ts()}] [REQ] /api/chat/next chunk {idx+1}/{total} responding | request-to-response: {t_resp-t_req:.2f}s")

    if is_last:
        session["done"] = True
        if len(_sessions) > 10:
            oldest = sorted(_sessions.keys(), key=lambda k: _sessions[k]["created"])
            for k in oldest[:-10]:
                del _sessions[k]

    # Use clean sentence for UI display (tags and TTS normalization stripped)
    clean_sentence = session["clean_sentences"][idx]

    return JSONResponse({
        "done": False,
        "chunk_index": idx,
        "total_chunks": total,
        "is_last": is_last,
        "sentence": clean_sentence,
        "audio_url": result["audio_url"],
        "audio_duration": result["audio_duration"],
        "lipsync_video_url": result["lipsync_video_url"],
    })


@app.post("/api/reset")
async def reset_chat():
    from llm import reset_conversation
    reset_conversation()
    return JSONResponse({"status": "ok"})


@app.post("/api/transcribe")
async def transcribe(request: Request):
    """Transcribe audio blob with Vosk. Accepts base64 JSON or raw body."""
    try:
        content_type = request.headers.get("content-type", "")

        if "json" in content_type:
            body = await request.json()
            import base64
            audio_bytes = base64.b64decode(body.get("audio", ""))
            audio_ct = body.get("content_type", "audio/webm")
        else:
            audio_bytes = await request.body()
            audio_ct = content_type or "audio/webm"

        if not audio_bytes or len(audio_bytes) < 500:
            return JSONResponse({"text": ""})

        import asyncio
        import vosk_stt
        text = await asyncio.to_thread(vosk_stt.transcribe_audio, audio_bytes, audio_ct)
        return JSONResponse({"text": text})

    except Exception as e:
        logger.error(f"[{_ts()}] [STT] Transcribe error: {e}", exc_info=True)
        return JSONResponse({"text": "", "error": str(e)})


@app.get("/api/greeting")
async def get_greeting():
    if GREETING_DATA:
        return JSONResponse(GREETING_DATA)
    return JSONResponse({"text": None})


# ---- WebSocket frame streaming ----

@app.websocket("/ws/video")
async def ws_video(websocket: WebSocket):
    """Stream JPEG frames to the client at 25fps via WebSocket."""
    await websocket.accept()
    import asyncio

    try:
        from float_streamer import get_streamer
        streamer = get_streamer()
        if not streamer.ready:
            await websocket.send_json({"type": "error", "message": "Streamer not ready"})
            await websocket.close()
            return

        logger.info(f"[{_ts()}] [WS] Client connected")
        frame_interval = 1.0 / 25.0  # 40ms per frame at 25fps

        while True:
            t_start = _time.time()

            # Pull either a JPEG byte string (frame) or a dict (speech_start event)
            item = streamer.get_frame(timeout=0.02)

            if isinstance(item, dict):
                # Audio trigger — send instantly, skip the 25fps delay
                await websocket.send_json(item)
                continue
            elif item is not None:
                # Video frame
                await websocket.send_bytes(item)

            # Pace to 25fps — wait remainder of frame interval
            elapsed = _time.time() - t_start
            wait = frame_interval - elapsed
            if wait > 0:
                await asyncio.sleep(wait)

    except WebSocketDisconnect:
        logger.info(f"[{_ts()}] [WS] Client disconnected")
    except Exception as e:
        logger.error(f"[{_ts()}] [WS] Error: {e}", exc_info=True)


@app.post("/api/chat/stream")
async def chat_stream(request: Request):
    """
    Streaming-mode chat: generates TTS and injects into the FLOAT streamer.
    No lipsync video files — frames come via WebSocket.
    """
    t0 = _time.time()
    body = await request.json()
    user_text = body.get("text", "").strip()
    if not user_text:
        raise HTTPException(400, "No text")

    logger.info(f"[{_ts()}] [CHAT-STREAM] User: {user_text}")

    # 1. LLM — run in worker thread so the event loop stays free to send WS frames
    t1 = _time.time()
    from llm import generate_response
    llm_result = await asyncio.to_thread(generate_response, user_text)
    reply_text = llm_result['text']
    clean_text = llm_result.get('clean_text', reply_text)
    t2 = _time.time()
    logger.info(f"[{_ts()}] [CHAT-STREAM] Reply: {clean_text[:80]}...")
    logger.info(f"[{_ts()}] [CHAT-STREAM] LLM: {t2-t1:.2f}s")

    # 2. TTS (full response as single audio) — also in worker thread
    from tts import generate_audio
    t3 = _time.time()
    audio_path = await asyncio.to_thread(generate_audio, reply_text)
    t4 = _time.time()

    audio_url = None
    audio_duration = 0
    if audio_path and os.path.exists(audio_path):
        audio_url = f"/api/audio/{os.path.basename(audio_path)}"
        try:
            # ffprobe is blocking too — run it off-loop
            r = await asyncio.to_thread(
                subprocess.run,
                ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
                 "-of", "csv=p=0", audio_path],
                capture_output=True, text=True, check=True,
            )
            audio_duration = float(r.stdout.strip())
        except Exception:
            audio_duration = 5.0
    logger.info(f"[{_ts()}] [CHAT-STREAM] TTS: {t4-t3:.2f}s | {audio_duration:.1f}s audio")

    # 3. Inject into streamer (non-blocking — queues audio features)
    try:
        from float_streamer import get_streamer
        streamer = get_streamer()
        if streamer.ready:
            # inject_speech is fast (~10-50ms) but does include the buffer-flush; off-load
            # for safety so the event loop never stalls mid-frame.
            await asyncio.to_thread(streamer.inject_speech, audio_path, audio_url)
            logger.info(f"[{_ts()}] [CHAT-STREAM] Audio injected into streamer")
    except Exception as e:
        logger.error(f"[{_ts()}] [CHAT-STREAM] Streamer inject failed: {e}")

    logger.info(f"[{_ts()}] [CHAT-STREAM] Total: {_time.time()-t0:.2f}s")

    return JSONResponse({
        "text": clean_text,
        "audio_url": audio_url,
        "audio_duration": audio_duration,
    })


@app.get("/", response_class=HTMLResponse)
async def index():
    return HTML_PAGE


HTML_PAGE = """\
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1.0"/>
<title>Avatar Chatbot</title>
<style>
*{margin:0;padding:0;box-sizing:border-box}
body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',sans-serif;background:#0a0f1c;color:#e2e8f0;height:100vh;display:flex;flex-direction:column;overflow:hidden}
.app-container{display:flex;flex-direction:column;height:100vh;max-width:900px;margin:0 auto;width:100%;padding:1rem;gap:.75rem}
.top-bar{display:flex;align-items:center;gap:.75rem;flex-shrink:0}
.top-bar label{font-size:.85rem;color:#94a3b8}
.top-bar select{background:#1e293b;color:#e2e8f0;border:1px solid rgba(59,130,246,.3);border-radius:6px;padding:6px 10px;font-size:.85rem;outline:none;cursor:pointer}
.top-bar select:focus{border-color:#3b82f6}
.video-section{position:relative;width:100%;min-height:300px;max-height:45vh;background:#000;border-radius:12px;overflow:hidden;flex-shrink:0;display:flex;align-items:center;justify-content:center}
#avatarCanvas{width:100%;height:100%;object-fit:contain}
.chat-section{flex:1;display:flex;flex-direction:column;min-height:0;background:rgba(15,23,42,.6);border:1px solid rgba(59,130,246,.15);border-radius:12px;overflow:hidden}
.chat-log{flex:1;overflow-y:auto;padding:1rem;display:flex;flex-direction:column;gap:.5rem}
.chat-log::-webkit-scrollbar{width:6px}.chat-log::-webkit-scrollbar-track{background:transparent}.chat-log::-webkit-scrollbar-thumb{background:#334155;border-radius:3px}
.msg{display:flex;width:100%}.msg.user{justify-content:flex-end}.msg.bot{justify-content:flex-start}
.bubble{max-width:85%;padding:10px 14px;border-radius:12px;word-wrap:break-word;font-size:.9rem;line-height:1.4}
.bubble.user{background:linear-gradient(135deg,#1e40af,#3b82f6);color:#fff;border-bottom-right-radius:4px}
.bubble.bot{background:rgba(30,41,59,.8);border:1px solid rgba(59,130,246,.2);color:#e2e8f0;border-bottom-left-radius:4px}
.bubble.bot .label{color:#60a5fa;font-weight:600;font-size:.8rem;margin-bottom:2px}
.bubble.user .label{color:rgba(255,255,255,.7);font-weight:600;font-size:.8rem;margin-bottom:2px}
.bubble.typing{color:#64748b;font-style:italic}
.bubble.error{background:rgba(127,29,29,.5);border-color:rgba(239,68,68,.3);color:#fca5a5}
@keyframes dots{0%,20%{content:''}40%{content:'.'}60%{content:'..'}80%,100%{content:'...'}}
.typing-dots::after{content:'';animation:dots 1.5s infinite}
.input-bar{display:flex;gap:.5rem;padding:.75rem 1rem;background:rgba(15,23,42,.8);border-top:1px solid rgba(59,130,246,.1)}
.input-bar input{flex:1;background:rgba(30,41,59,.8);border:1px solid rgba(59,130,246,.25);border-radius:8px;padding:10px 14px;color:#e2e8f0;font-size:.9rem;outline:none;transition:border-color .2s}
.input-bar input:focus{border-color:#3b82f6}.input-bar input::placeholder{color:#64748b}
.input-bar button{padding:10px 20px;border:none;border-radius:8px;font-size:.9rem;cursor:pointer;transition:all .15s;font-weight:500}
.btn-send{background:#3b82f6;color:#fff}.btn-send:hover{background:#2563eb}.btn-send:disabled{background:#334155;color:#64748b;cursor:not-allowed}
.btn-reset{background:rgba(100,116,139,.3);color:#94a3b8;border:1px solid rgba(100,116,139,.3)}.btn-reset:hover{background:rgba(100,116,139,.5)}
.btn-mic{background:rgba(100,116,139,.3);color:#94a3b8;border:1px solid rgba(100,116,139,.3);font-size:1.1rem;padding:10px 14px}.btn-mic:hover{background:rgba(100,116,139,.5)}
.btn-mic.recording{background:rgba(239,68,68,.3);color:#fca5a5;border-color:rgba(239,68,68,.5);animation:pulse-rec 1.5s infinite}
@keyframes pulse-rec{0%,100%{box-shadow:0 0 0 0 rgba(239,68,68,.3)}50%{box-shadow:0 0 0 8px rgba(239,68,68,0)}}
.btn-clear{background:rgba(100,116,139,.2);color:#64748b;border:1px solid rgba(100,116,139,.2);padding:10px 12px;font-size:.85rem;display:none}.btn-clear:hover{background:rgba(100,116,139,.4);color:#94a3b8}
.btn-topbar{background:rgba(59,130,246,.15);color:#93c5fd;border:1px solid rgba(59,130,246,.3);border-radius:6px;padding:5px 10px;font-size:.8rem;cursor:pointer;transition:background .15s}
.btn-topbar:hover{background:rgba(59,130,246,.3)}
.modal-backdrop{position:fixed;inset:0;background:rgba(0,0,0,.65);display:none;align-items:center;justify-content:center;z-index:100;padding:1rem}
.modal-backdrop.show{display:flex}
.modal{background:#0f172a;border:1px solid rgba(59,130,246,.3);border-radius:12px;padding:1.25rem;width:100%;max-width:480px;max-height:90vh;overflow-y:auto}
.modal h2{font-size:1.1rem;margin-bottom:1rem}
.modal .field{margin-bottom:.85rem}
.modal label{display:block;font-size:.78rem;color:#94a3b8;margin-bottom:4px}
.modal input[type=text],.modal textarea{width:100%;background:rgba(30,41,59,.8);border:1px solid rgba(59,130,246,.25);border-radius:6px;padding:8px 10px;color:#e2e8f0;font-size:.85rem;outline:none;font-family:inherit}
.modal input[type=text]:focus,.modal textarea:focus{border-color:#3b82f6}
.modal textarea{min-height:80px;resize:vertical}
.modal .hint{font-size:.7rem;color:#64748b;margin-top:3px}
.modal .hint a{color:#60a5fa;text-decoration:none}.modal .hint a:hover{text-decoration:underline}
.modal input[type=file]{font-size:.78rem;color:#94a3b8}
.preview-img{display:none;max-width:120px;max-height:120px;border-radius:8px;margin-top:8px;border:1px solid rgba(59,130,246,.3)}
.modal-actions{display:flex;justify-content:flex-end;gap:.5rem;margin-top:1rem}
.modal-actions button{padding:8px 16px;border:none;border-radius:6px;font-size:.85rem;cursor:pointer}
.btn-cancel{background:rgba(100,116,139,.3);color:#cbd5e1}
.btn-cancel:hover{background:rgba(100,116,139,.5)}
.btn-create{background:#3b82f6;color:#fff}
.btn-create:hover{background:#2563eb}
.btn-create:disabled{background:#334155;color:#64748b;cursor:wait}
.progress{display:none;margin-top:1rem;padding:.75rem;background:rgba(30,41,59,.6);border-radius:8px;font-size:.8rem}
.progress.show{display:block}
.progress-bar{height:6px;background:rgba(100,116,139,.3);border-radius:3px;overflow:hidden;margin-top:6px}
.progress-fill{height:100%;background:#3b82f6;width:0%;transition:width .3s}
.progress.error{color:#fca5a5}
.progress.done{color:#86efac}
</style>
</head>
<body>
<div class="app-container">
  <div class="top-bar">
    <label for="avatarSelect">Avatar:</label>
    <select id="avatarSelect"></select>
    <button class="btn-topbar" id="btnNewAvatar" title="Create a new avatar">+ New</button>
  </div>
  <div class="video-section">
    <canvas id="avatarCanvas" width="512" height="512"></canvas>
  </div>
  <div class="chat-section">
    <div class="chat-log" id="chatLog"></div>
    <div class="input-bar">
      <input type="text" id="chatInput" placeholder="Say something..." autocomplete="off"/>
      <button class="btn-clear" id="btnClear" title="Clear input">&#10005;</button>
      <button class="btn-send" id="btnSend">Send</button>
      <button class="btn-mic" id="btnMic" title="Voice input">&#127908;</button>
      <button class="btn-reset" id="btnReset" title="Reset conversation">&#8635;</button>
    </div>
  </div>
</div>
<div class="modal-backdrop" id="createModal">
  <div class="modal">
    <h2>Create New Avatar</h2>
    <div class="field">
      <label>Name</label>
      <input type="text" id="newName" placeholder="e.g. Maya" maxlength="40"/>
    </div>
    <div class="field">
      <label>Reference image (face photo)</label>
      <input type="file" id="newImage" accept="image/png,image/jpeg"/>
      <img class="preview-img" id="imagePreview"/>
    </div>
    <div class="field">
      <label>Voice (edge-tts ID)</label>
      <input type="text" id="newVoice" value="en-US-AvaMultilingualNeural" placeholder="e.g. en-US-AriaNeural"/>
      <div class="hint">Preview voices at <a href="https://tts.travisvn.com/" target="_blank">tts.travisvn.com</a> or <a href="https://geeksta.net/tools/tts-samples/" target="_blank">geeksta.net</a> &mdash; paste the voice ID here</div>
    </div>
    <div class="field">
      <label>Persona (optional)</label>
      <textarea id="newPersona" placeholder="A friendly AI assistant who likes talking about books..."></textarea>
    </div>
    <div class="progress" id="createProgress"></div>
    <div class="modal-actions">
      <button class="btn-cancel" id="btnCancelCreate">Cancel</button>
      <button class="btn-create" id="btnConfirmCreate">Create</button>
    </div>
  </div>
</div>
<script>
// ======================================================================
// STATE
// ======================================================================
let isProcessing=false, msgCounter=0, currentAvatarName='Avatar', isRecording=false;
const canvas=document.getElementById('avatarCanvas');
const ctx=canvas.getContext('2d');
const chatLog=document.getElementById('chatLog');
const chatInput=document.getElementById('chatInput');
const btnSend=document.getElementById('btnSend');
const btnReset=document.getElementById('btnReset');
const avatarSelect=document.getElementById('avatarSelect');
const btnMic=document.getElementById('btnMic');
const btnClear=document.getElementById('btnClear');

// ======================================================================
// WEBSOCKET FRAME STREAMING
// ======================================================================
let ws=null, frameImg=new Image();

function drawFrame(data){
    const blob=new Blob([data], {type:'image/jpeg'});
    const url=URL.createObjectURL(blob);
    frameImg.onload=()=>{
        ctx.drawImage(frameImg, 0, 0, canvas.width, canvas.height);
        URL.revokeObjectURL(url);
    };
    frameImg.src=url;
}

function connectWebSocket(){
    const proto=location.protocol==='https:'?'wss:':'ws:';
    ws=new WebSocket(proto+'//'+location.host+'/ws/video');
    ws.binaryType='arraybuffer';

    ws.onopen=()=>{ console.log('[WS] Connected'); };

    ws.onmessage=(e)=>{
        if(typeof e.data==='string'){
            const msg=JSON.parse(e.data);
            if(msg.type==='speech_start' && msg.audio_url){
                // Preload audio then play — frames draw continuously
                const audio=new Audio(msg.audio_url);
                audio.addEventListener('canplaythrough', ()=>{
                    audio.play().catch(err=>console.warn('[AUDIO]',err));
                }, {once:true});
                audio.load();
            }
        } else {
            // Always draw frames immediately — server paces at 25fps
            drawFrame(e.data);
        }
    };

    ws.onclose=()=>{
        console.log('[WS] Disconnected, reconnecting in 2s...');
        setTimeout(connectWebSocket, 2000);
    };

    ws.onerror=(e)=>{
        console.error('[WS] Error:', e);
    };
}

// ======================================================================
// CHAT
// ======================================================================
function streamText(el, text, totalMs){
    // Reveal text into `el` progressively over `totalMs`, driven by requestAnimationFrame.
    // Naturally smooth at any text length / duration. Auto-scrolls the parent chatLog as it grows.
    return new Promise((resolve)=>{
        if(!text){ resolve(); return; }
        el.textContent='';
        const startTime=performance.now();
        const tick=(now)=>{
            const progress=Math.min(1, (now-startTime)/totalMs);
            const charsShown=Math.floor(text.length*progress);
            el.textContent=text.slice(0, charsShown);
            chatLog.scrollTop=chatLog.scrollHeight;
            if(progress<1){
                requestAnimationFrame(tick);
            } else {
                el.textContent=text;
                chatLog.scrollTop=chatLog.scrollHeight;
                resolve();
            }
        };
        requestAnimationFrame(tick);
    });
}

function addMessage(role, text, extra){
    const id='msg_'+(msgCounter++);
    const w=document.createElement('div'); w.className='msg '+role;
    const b=document.createElement('div'); b.className='bubble '+role+(extra?' '+extra:''); b.id=id;
    const l=document.createElement('div'); l.className='label'; l.textContent=role==='user'?'You':currentAvatarName;
    const c=document.createElement('div'); c.className='content'; c.textContent=text;
    b.appendChild(l); b.appendChild(c); w.appendChild(b);
    chatLog.appendChild(w); chatLog.scrollTop=chatLog.scrollHeight;
    return {id:id, contentEl:c, bubbleEl:b};
}

function addTypingIndicator(){
    const r=addMessage('bot','','typing');
    r.contentEl.className='content typing-dots'; r.contentEl.textContent='thinking';
    return r;
}

async function sendMessage(){
    const text=chatInput.value.trim();
    if(!text||isProcessing) return;
    isProcessing=true; btnSend.disabled=true; chatInput.value=''; btnClear.style.display='none';
    addMessage('user', text);
    const typing=addTypingIndicator();

    try{
        // Single request — LLM + TTS + inject into streamer
        const res=await fetch('/api/chat/stream',{
            method:'POST', headers:{'Content-Type':'application/json'},
            body:JSON.stringify({text:text})
        });
        if(!res.ok) throw new Error('HTTP '+res.status);
        const data=await res.json();

        // Reveal text gradually, paced to roughly match the speech duration.
        // The audio plays via WS speech_start shortly after the response arrives,
        // so the text and speech end at about the same moment.
        typing.bubbleEl.classList.remove('typing');
        typing.contentEl.classList.remove('typing-dots');
        const streamDurationMs = (data.audio_duration || 3) * 1000;
        await streamText(typing.contentEl, data.text, streamDurationMs);

        // Tail wait so the next prompt isn't accepted before audio actually finishes
        // (text streams over audio_duration; audio playback begins ~1-2s after that
        // window starts because of FLOAT's ODE solve, so we wait the residual here).
        await new Promise(r=>setTimeout(r, 1500));

    }catch(err){
        typing.bubbleEl.classList.remove('typing');
        typing.bubbleEl.classList.add('error');
        typing.contentEl.classList.remove('typing-dots');
        typing.contentEl.textContent='Error: '+err.message;
    }
    isProcessing=false; btnSend.disabled=false; chatInput.focus();
}

btnSend.addEventListener('click', sendMessage);
chatInput.addEventListener('keydown', e=>{ if(e.key==='Enter'&&!e.shiftKey){ e.preventDefault(); if(!isRecording) sendMessage(); }});

// ---- Avatar switching ----
avatarSelect.addEventListener('change', async ()=>{
    const name=avatarSelect.value;
    if(!name || isProcessing) return;
    btnSend.disabled=true;
    try{
        await fetch('/api/switch-avatar',{
            method:'POST', headers:{'Content-Type':'application/json'},
            body:JSON.stringify({avatar:name})
        });
        currentAvatarName=name;
        chatLog.innerHTML='';
        addMessage('bot','Hi, I am '+name+'. How can I help you?');
    }catch(e){
        console.error('Switch failed:', e);
    }
    btnSend.disabled=false;
});

btnReset.addEventListener('click', async()=>{
    await fetch('/api/reset',{method:'POST'});
    chatLog.innerHTML='';
    addMessage('bot','Conversation reset. How can I help you?');
});

// ---- Clear button ----
chatInput.addEventListener('input',()=>{
    btnClear.style.display=chatInput.value?'block':'none';
});
btnClear.addEventListener('click',()=>{
    chatInput.value='';
    btnClear.style.display='none';
    chatInput.focus();
});

// ---- Voice Input (Mic) ----
let mediaRecorder=null, previewTimer=null, allChunks=[], micStream=null, micMimeType='';

async function transcribeBlob(blob){
    const reader=new FileReader();
    return new Promise((resolve)=>{
        reader.onloadend=async()=>{
            const b64=reader.result.split(',')[1];
            try{
                const res=await fetch('/api/transcribe',{
                    method:'POST',
                    headers:{'Content-Type':'application/json'},
                    body:JSON.stringify({audio:b64, content_type:micMimeType||'audio/webm'})
                });
                const data=await res.json();
                resolve(data.text||'');
            }catch(e){ resolve(''); }
        };
        reader.readAsDataURL(blob);
    });
}

btnMic.addEventListener('click', async()=>{
    if(isRecording){
        if(previewTimer) clearInterval(previewTimer);
        previewTimer=null;
        mediaRecorder.stop();
        return;
    }
    try{
        micStream=await navigator.mediaDevices.getUserMedia({audio:true});
        micMimeType=MediaRecorder.isTypeSupported('audio/webm;codecs=opus')?'audio/webm;codecs=opus':
                    MediaRecorder.isTypeSupported('audio/webm')?'audio/webm':'';
        mediaRecorder=new MediaRecorder(micStream, micMimeType?{mimeType:micMimeType}:{});
        allChunks=[];
        mediaRecorder.ondataavailable=(e)=>{ if(e.data.size>0) allChunks.push(e.data); };
        mediaRecorder.onstop=async()=>{
            isRecording=false;
            btnMic.classList.remove('recording');
            btnMic.textContent=String.fromCodePoint(0x1F3A4);
            chatInput.placeholder='Transcribing...';
            micStream.getTracks().forEach(t=>t.stop());
            if(allChunks.length===0){ chatInput.placeholder='Say something...'; btnSend.disabled=false; return; }
            const fullBlob=new Blob(allChunks, {type:micMimeType||'audio/webm'});
            const text=await transcribeBlob(fullBlob);
            chatInput.value=text;
            chatInput.placeholder='Say something...';
            btnClear.style.display=text?'block':'none';
            btnSend.disabled=false;
            chatInput.focus();
        };
        mediaRecorder.start(500);
        isRecording=true;
        btnMic.classList.add('recording');
        btnMic.textContent=String.fromCodePoint(0x23F9);
        chatInput.value='';
        chatInput.placeholder='Listening...';
        btnSend.disabled=true;
        btnClear.style.display='none';
        let previewBusy=false;
        previewTimer=setInterval(async()=>{
            if(previewBusy||allChunks.length<2) return;
            previewBusy=true;
            try{
                const previewBlob=new Blob(allChunks.slice(), {type:micMimeType||'audio/webm'});
                if(previewBlob.size>2000){
                    const text=await transcribeBlob(previewBlob);
                    if(isRecording && text) chatInput.value=text;
                }
            }catch(e){}
            previewBusy=false;
        }, 2000);
    }catch(err){
        console.error('Mic access error:',err);
        alert('Could not access microphone. Please allow mic permission.');
    }
});

// ======================================================================
// CREATE AVATAR MODAL
// ======================================================================
const createModal=document.getElementById('createModal');
const btnNewAvatar=document.getElementById('btnNewAvatar');
const btnCancelCreate=document.getElementById('btnCancelCreate');
const btnConfirmCreate=document.getElementById('btnConfirmCreate');
const newImage=document.getElementById('newImage');
const imagePreview=document.getElementById('imagePreview');
const createProgress=document.getElementById('createProgress');

let imageDataUrl=null, statusPollTimer=null;

function openCreateModal(){
  document.getElementById('newName').value='';
  document.getElementById('newVoice').value='en-US-AvaMultilingualNeural';
  document.getElementById('newPersona').value='';
  newImage.value='';
  imagePreview.style.display='none';
  imageDataUrl=null;
  createProgress.className='progress';
  createProgress.textContent='';
  btnConfirmCreate.disabled=false;
  btnConfirmCreate.textContent='Create';
  createModal.classList.add('show');
}
function closeCreateModal(){
  createModal.classList.remove('show');
  if(statusPollTimer){ clearInterval(statusPollTimer); statusPollTimer=null; }
}

btnNewAvatar.addEventListener('click', openCreateModal);
btnCancelCreate.addEventListener('click', closeCreateModal);
createModal.addEventListener('click', (e)=>{ if(e.target===createModal) closeCreateModal(); });

newImage.addEventListener('change', (e)=>{
  const file=e.target.files[0];
  if(!file) return;
  const reader=new FileReader();
  reader.onload=(ev)=>{
    imageDataUrl=ev.target.result;
    imagePreview.src=imageDataUrl;
    imagePreview.style.display='block';
  };
  reader.readAsDataURL(file);
});

btnConfirmCreate.addEventListener('click', async()=>{
  const name=document.getElementById('newName').value.trim();
  const voice=document.getElementById('newVoice').value.trim();
  const persona=document.getElementById('newPersona').value.trim();
  if(!name){ alert('Please enter a name'); return; }
  if(!imageDataUrl){ alert('Please choose a reference image'); return; }
  btnConfirmCreate.disabled=true;
  btnConfirmCreate.textContent='Starting...';
  createProgress.className='progress show';
  createProgress.textContent='Uploading...';
  try{
    const res=await fetch('/api/create-avatar', {
      method:'POST', headers:{'Content-Type':'application/json'},
      body:JSON.stringify({name, voice, persona, image:imageDataUrl})
    });
    if(!res.ok){
      const err=await res.json().catch(()=>({detail:'Unknown error'}));
      throw new Error(err.detail||('HTTP '+res.status));
    }
    const data=await res.json();
    createProgress.innerHTML='Generating idle videos... <div class="progress-bar"><div class="progress-fill" id="pfill"></div></div>';
    pollCreateStatus(data.avatar);
  }catch(e){
    createProgress.className='progress show error';
    createProgress.textContent='Error: '+e.message;
    btnConfirmCreate.disabled=false;
    btnConfirmCreate.textContent='Create';
  }
});

function pollCreateStatus(avatarName){
  if(statusPollTimer) clearInterval(statusPollTimer);
  statusPollTimer=setInterval(async()=>{
    try{
      const res=await fetch('/api/create-avatar/status');
      const s=await res.json();
      const fill=document.getElementById('pfill');
      if(fill && s.total>0){
        fill.style.width=Math.round(s.progress/s.total*100)+'%';
      }
      const msg=document.createElement('div');
      msg.textContent=s.message||'Working...';
      const existing=createProgress.querySelector('.status-msg');
      if(existing) existing.remove();
      msg.className='status-msg';
      createProgress.appendChild(msg);
      if(s.error){
        clearInterval(statusPollTimer); statusPollTimer=null;
        createProgress.className='progress show error';
        createProgress.textContent='Error: '+s.error;
        btnConfirmCreate.disabled=false;
        btnConfirmCreate.textContent='Create';
      } else if(s.done){
        clearInterval(statusPollTimer); statusPollTimer=null;
        createProgress.className='progress show done';
        createProgress.textContent='Done! Switching to '+avatarName+'...';
        // Refresh avatar list and select the new one
        const ar=await fetch('/api/avatars');
        const ad=await ar.json();
        avatarSelect.innerHTML='';
        for(const a of ad.avatars){
          const opt=document.createElement('option');
          opt.value=a.name; opt.textContent=a.name;
          if(a.name===avatarName) opt.selected=true;
          avatarSelect.appendChild(opt);
        }
        currentAvatarName=avatarName;
        setTimeout(closeCreateModal, 1500);
      }
    }catch(e){ console.warn('[CREATE] Poll error:', e); }
  }, 1500);
}

// ======================================================================
// INIT
// ======================================================================
async function init(){
    try{
        // Load avatars
        const avatarRes=await fetch('/api/avatars');
        const avatarData=await avatarRes.json();
        avatarSelect.innerHTML='';
        for(const a of avatarData.avatars){
            const opt=document.createElement('option');
            opt.value=a.name; opt.textContent=a.name;
            if(a.name===avatarData.current) opt.selected=true;
            avatarSelect.appendChild(opt);
        }
        if(avatarData.current) currentAvatarName=avatarData.current;

        // Show greeting
        const greetRes=await fetch('/api/greeting');
        const greeting=await greetRes.json();
        addMessage('bot', greeting.text || 'Hello! Feel free to ask me anything.');

        // Connect WebSocket for live frame streaming
        connectWebSocket();

    }catch(e){ console.error('Init error:', e); }
}
init();
</script>
</body>
</html>
"""