server/main_fast.py

"""
MuseTalk Fast Speech-to-Speech Server
With pre-loaded models, avatar caching, and streaming chunks
"""

import os
import sys
import asyncio
import tempfile
import uuid
import time
import json
from pathlib import Path
from typing import Optional
import subprocess
import threading

sys.path.insert(0, str(Path(__file__).parent.parent))

from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import httpx
import base64
import cv2
import numpy as np
import queue

# H.264 encoder for efficient streaming
try:
    from server.h264_encoder import H264StreamEncoder
    H264_AVAILABLE = True
except ImportError:
    H264_AVAILABLE = False
    print("[WARNING] H.264 encoder not available, using JPEG only")

# API Keys
GROQ_API_KEY = "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"
ELEVENLABS_API_KEY = "sk_857e9e6f2412ddf3ff5334b736e4b571641d26225c0d8d62"
ELEVENLABS_VOICE_ID = "21m00Tcm4TlvDq8ikWAM"

# Paths
BASE_DIR = Path(__file__).parent.parent
RESULTS_DIR = BASE_DIR / "results" / "server"
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

app = FastAPI(title="MuseTalk Fast API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory=Path(__file__).parent / "static"), name="static")
app.mount("/avatar_videos", StaticFiles(directory=Path(__file__).parent / "avatar_videos"), name="avatar_videos")

# React app static files (built with Vite)
REACT_BUILD_DIR = Path(__file__).parent / "web" / "dist"
if REACT_BUILD_DIR.exists():
    app.mount("/app/assets", StaticFiles(directory=REACT_BUILD_DIR / "assets"), name="react_assets")

@app.middleware("http")
async def add_cache_headers(request, call_next):
    response = await call_next(request)
    if request.url.path.endswith(".mp4") or request.url.path.endswith(".jpg"):
        response.headers["Cache-Control"] = "public, max-age=3600, immutable"
    return response

# Session Metrics
class SessionMetrics:
    def __init__(self, request_id: str):
        self.request_id = request_id
        self.start_time = time.time()
        self.timings = {}
        self.frame_timings = []
        
    def mark(self, event: str):
        elapsed = time.time() - self.start_time
        self.timings[event] = round(elapsed * 1000)  # ms
        print(f"[METRICS] {event}: {elapsed:.3f}s")
        
    def frame_sent(self, index: int):
        self.frame_timings.append({"i": index, "t": round((time.time() - self.start_time) * 1000)})
    
    def to_dict(self):
        return {
            "request_id": self.request_id,
            "total_ms": round((time.time() - self.start_time) * 1000),
            "timings": self.timings,
            "frames": len(self.frame_timings)
        }


# Global state
engine = None
conversation_history = []
video_jobs = {}  # job_id -> {"status": str, "progress": int, "video_path": str}


class TextRequest(BaseModel):
    text: str


def convert_webm_to_wav(input_path: str, output_path: str) -> str:
    cmd = ["ffmpeg", "-y", "-v", "quiet", "-i", input_path, "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le", output_path]
    subprocess.run(cmd, capture_output=True)
    return output_path


DISTIL_WHISPER_URL = "http://localhost:8766"  # Distil-Whisper PT-BR (~50ms)

async def transcribe_audio_local(audio_path: str) -> str:
    """Transcribe using Distil-Whisper PT-BR (local, ~50ms)."""
    async with httpx.AsyncClient(timeout=30.0) as client:
        with open(audio_path, "rb") as f:
            response = await client.post(
                f"{DISTIL_WHISPER_URL}/inference",
                files={"file": (os.path.basename(audio_path), f, "audio/wav")},
                data={"response_format": "json", "language": "pt"}
            )
            if response.status_code != 200:
                raise HTTPException(status_code=500, detail=f"Distil-Whisper error: {response.text}")
            result = response.json()
            return result.get("text", "").strip()

async def transcribe_audio_groq(audio_path: str) -> str:
    """Fallback to Groq API if local Whisper fails."""
    async with httpx.AsyncClient(timeout=30.0) as client:
        with open(audio_path, "rb") as f:
            response = await client.post(
                "https://api.groq.com/openai/v1/audio/transcriptions",
                headers={"Authorization": f"Bearer {GROQ_API_KEY}"},
                files={"file": ("audio.wav", f, "audio/wav")},
                data={"model": "whisper-large-v3", "response_format": "text", "language": "pt"}
            )
            if response.status_code != 200:
                raise HTTPException(status_code=500, detail=f"Groq STT error: {response.text}")
            return response.text.strip()

async def transcribe_audio(audio_path: str) -> str:
    """Transcribe audio - tries local Whisper first, falls back to Groq."""
    try:
        return await transcribe_audio_local(audio_path)
    except Exception as e:
        print(f"[STT] Local Whisper failed ({e}), falling back to Groq")
        return await transcribe_audio_groq(audio_path)


VLLM_URL = "http://localhost:8000"
VLLM_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"

OLLAMA_URL = "http://localhost:11434"
OLLAMA_MODEL = "gemma3:1b-it-q4_K_M"

# System prompt for short, direct responses
LLM_SYSTEM_PROMPT = """Responda em portugues de forma muito curta e direta. Maximo 1 frase."""

async def chat_vllm(messages: list) -> str:
    """Chat using vLLM (fastest, ~65ms)."""
    import time
    start = time.time()

    # Build messages with system prompt
    llm_messages = [
        {"role": "system", "content": LLM_SYSTEM_PROMPT}
    ]
    for msg in messages:
        if msg.get("role") == "user":
            llm_messages.append({"role": "user", "content": msg.get("content", "")})

    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(
            f"{VLLM_URL}/v1/chat/completions",
            json={
                "model": VLLM_MODEL,
                "messages": llm_messages,
                "max_tokens": 50,
                "temperature": 0.7
            }
        )
        if response.status_code != 200:
            raise Exception(f"vLLM error: {response.text}")

        result = response.json()
        elapsed = (time.time() - start) * 1000
        content = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
        print(f"[LLM] vLLM local: {elapsed:.0f}ms")
        return content

async def chat_ollama(messages: list) -> str:
    """Chat using local Ollama/Gemma (fast, ~340ms)."""
    import time
    start = time.time()

    # Get user message
    user_msg = ""
    for msg in messages:
        if msg.get("role") == "user":
            user_msg = msg.get("content", "")

    # Simple prompt for Gemma
    prompt = f"Responda em portugues, muito curto (1 frase): {user_msg}"

    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": OLLAMA_MODEL,
                "prompt": prompt,
                "stream": False,
                "options": {"num_predict": 50, "temperature": 0.7, "stop": ["\n", "."]}
            }
        )
        if response.status_code != 200:
            raise Exception(f"Ollama error: {response.text}")

        result = response.json()
        elapsed = (time.time() - start) * 1000
        print(f"[LLM] Ollama local: {elapsed:.0f}ms")
        return result.get("response", "").strip()

async def chat_groq(messages: list) -> str:
    """Fallback to Groq API."""
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(
            "https://api.groq.com/openai/v1/chat/completions",
            headers={"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"},
            json={"model": "llama-3.1-8b-instant", "messages": messages, "max_tokens": 500, "temperature": 0.7}
        )
        if response.status_code != 200:
            raise HTTPException(status_code=500, detail=f"Groq LLM error: {response.text}")
        return response.json()["choices"][0]["message"]["content"]

async def chat_llm(messages: list) -> str:
    """Chat - tries vLLM first, then Ollama, then Groq."""
    # Try vLLM first (fastest, ~65ms)
    try:
        return await chat_vllm(messages)
    except Exception as e:
        print(f"[LLM] vLLM failed ({e}), trying Ollama")

    # Try Ollama (fast, ~340ms)
    try:
        return await chat_ollama(messages)
    except Exception as e:
        print(f"[LLM] Ollama failed ({e}), falling back to Groq")

    # Fallback to Groq API
    return await chat_groq(messages)


def text_to_speech_espeak(text: str, output_path: str, target_duration_ms: int = None) -> str:
    """
    Ultra-fast TTS with espeak-ng (~15ms).
    Used to generate video while ElevenLabs loads in parallel.
    If target_duration_ms is provided, adjusts speed to match.
    """
    base_speed = 175

    if target_duration_ms:
        # First generate with base speed to calculate ratio
        temp_path = output_path + ".temp.wav"
        subprocess.run([
            "espeak-ng", "-v", "pt-br", "-s", str(base_speed), "-w", temp_path, text
        ], capture_output=True)

        # Get duration and calculate required speed
        result = subprocess.run([
            'ffprobe', '-v', 'error', '-show_entries', 'format=duration',
            '-of', 'default=noprint_wrappers=1:nokey=1', temp_path
        ], capture_output=True, text=True)
        base_dur_ms = int(float(result.stdout.strip()) * 1000)

        if base_dur_ms > 0:
            required_speed = int(base_speed * (base_dur_ms / target_duration_ms))
            required_speed = max(80, min(400, required_speed))
        else:
            required_speed = base_speed

        os.remove(temp_path)
    else:
        required_speed = base_speed

    # Generate final audio
    subprocess.run([
        "espeak-ng", "-v", "pt-br", "-s", str(required_speed), "-w", output_path, text
    ], capture_output=True)

    return output_path


async def text_to_speech_elevenlabs(text: str, output_path: str) -> str:
    """
    Text-to-speech with ElevenLabs.
    Uses eleven_flash_v2_5 model (fastest available).
    """
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(
            f"https://api.elevenlabs.io/v1/text-to-speech/{ELEVENLABS_VOICE_ID}/stream",
            headers={
                "xi-api-key": ELEVENLABS_API_KEY,
                "Content-Type": "application/json"
            },
            json={
                "text": text,
                "model_id": "eleven_flash_v2_5",
                "voice_settings": {"stability": 0.5, "similarity_boost": 0.75}
            }
        )
        if response.status_code != 200:
            raise HTTPException(status_code=500, detail=f"ElevenLabs error: {response.text}")

        with open(output_path, "wb") as f:
            f.write(response.content)

        return output_path


def generate_video_background(job_id: str, audio_path: str, resolution: int = 256, batch_size: int = 8):
    """Background task to generate video"""
    global engine, video_jobs

    def progress_callback(progress, message):
        video_jobs[job_id]["progress"] = progress
        video_jobs[job_id]["message"] = message

    try:
        video_jobs[job_id] = {"status": "processing", "progress": 0, "message": "Iniciando..."}

        output_path = str(RESULTS_DIR / f"video_{job_id}.mp4")
        engine.generate_video_fast(audio_path, output_path, resolution=resolution, batch_size=batch_size, callback=progress_callback)

        video_jobs[job_id]["status"] = "completed"
        video_jobs[job_id]["video_path"] = output_path
        video_jobs[job_id]["progress"] = 100

    except Exception as e:
        video_jobs[job_id]["status"] = "error"
        video_jobs[job_id]["error"] = str(e)


@app.on_event("startup")
async def startup_event():
    """Initialize engine on startup"""
    global engine
    print("\n" + "=" * 60)
    print("STARTING FAST MUSETALK SERVER")
    print("=" * 60)

    from server.fast_engine import initialize_engine
    engine = initialize_engine()

    print("\n✓ Server ready!")
    print("=" * 60 + "\n")


@app.get("/")
async def root():
    return FileResponse(Path(__file__).parent / "static" / "index_rtc.html")


@app.get("/websocket")
async def websocket_page():
    """Old WebSocket version (backup)"""
    return FileResponse(Path(__file__).parent / "static" / "index_video.html")


@app.get("/video")
async def video_page():
    return FileResponse(Path(__file__).parent / "static" / "index_video.html")


@app.get("/api/health")
async def health():
    return {"status": "ok", "engine_ready": engine is not None and engine.avatar_loaded}


@app.post("/api/chat")
async def chat(request: TextRequest):
    global conversation_history
    conversation_history.append({"role": "user", "content": request.text})
    messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:]
    response_text = await chat_llm(messages)
    conversation_history.append({"role": "assistant", "content": response_text})
    return {"text": response_text}


@app.post("/api/tts")
async def tts(request: TextRequest):
    output_path = str(RESULTS_DIR / f"tts_{uuid.uuid4().hex[:8]}.mp3")
    await text_to_speech_elevenlabs(request.text, output_path)
    return FileResponse(output_path, media_type="audio/mpeg")


@app.post("/api/transcribe")
async def transcribe(audio: UploadFile = File(...)):
    temp_audio = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    content = await audio.read()
    temp_audio.write(content)
    temp_audio.close()
    try:
        text = await transcribe_audio(temp_audio.name)
        return {"text": text}
    finally:
        os.unlink(temp_audio.name)


@app.post("/api/conversation")
async def conversation_fast(
    audio: UploadFile = File(...),
    resolution: int = Form(256),
    batch_size: int = Form(8),
    background_tasks: BackgroundTasks = None
):
    """
    Fast conversation: generates video with audio synced
    """
    start_time = time.time()
    job_id = uuid.uuid4().hex[:8]

    print(f"\n{'='*50}")
    print(f"NEW CONVERSATION - Job: {job_id}")
    print(f"Settings: resolution={resolution}, batch_size={batch_size}")
    print(f"{'='*50}")

    # Save and convert audio
    temp_webm = tempfile.NamedTemporaryFile(suffix=".webm", delete=False)
    content = await audio.read()
    temp_webm.write(content)
    temp_webm.close()

    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    temp_wav.close()
    convert_webm_to_wav(temp_webm.name, temp_wav.name)

    try:
        # STT
        stt_start = time.time()
        user_text = await transcribe_audio(temp_wav.name)
        stt_time = time.time() - stt_start
        print(f"STT ({stt_time:.2f}s): {user_text}")

        # LLM
        llm_start = time.time()
        global conversation_history
        conversation_history.append({"role": "user", "content": user_text})
        messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:]
        assistant_text = await chat_llm(messages)
        conversation_history.append({"role": "assistant", "content": assistant_text})
        llm_time = time.time() - llm_start
        print(f"LLM ({llm_time:.2f}s): {assistant_text}")

        # TTS
        tts_start = time.time()
        tts_path = str(RESULTS_DIR / f"tts_{job_id}.mp3")
        await text_to_speech_elevenlabs(assistant_text, tts_path)
        tts_time = time.time() - tts_start
        print(f"TTS ({tts_time:.2f}s)")

        total_time = time.time() - start_time
        print(f"Fast response: {total_time:.2f}s (resolution={resolution}, batch={batch_size})")

        # Start video generation in background
        tts_wav = str(RESULTS_DIR / f"tts_{job_id}.wav")
        subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", tts_path, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True)

        video_jobs[job_id] = {"status": "queued", "progress": 0}
        thread = threading.Thread(target=generate_video_background, args=(job_id, tts_wav, resolution, batch_size))
        thread.start()

        return {
            "job_id": job_id,
            "user_text": user_text,
            "assistant_text": assistant_text,
            "audio_url": f"/api/audio/{job_id}",
            "video_status": "generating",
            "timing": {"stt": round(stt_time, 2), "llm": round(llm_time, 2), "tts": round(tts_time, 2), "total": round(total_time, 2)}
        }
    finally:
        try:
            os.unlink(temp_webm.name)
            os.unlink(temp_wav.name)
        except:
            pass


@app.get("/api/audio/{job_id}")
async def get_audio(job_id: str):
    audio_path = RESULTS_DIR / f"tts_{job_id}.mp3"
    if not audio_path.exists():
        raise HTTPException(status_code=404, detail="Audio not found")
    return FileResponse(audio_path, media_type="audio/mpeg")


@app.get("/api/video-status/{job_id}")
async def video_status(job_id: str):
    """Check video generation status"""
    if job_id not in video_jobs:
        raise HTTPException(status_code=404, detail="Job not found")
    return video_jobs[job_id]


@app.get("/api/video/{job_id}")
async def get_video(job_id: str):
    """Get generated video"""
    if job_id not in video_jobs:
        raise HTTPException(status_code=404, detail="Job not found")

    job = video_jobs[job_id]
    if job["status"] != "completed":
        raise HTTPException(status_code=202, detail=f"Video still generating: {job.get('progress', 0)}%")

    return FileResponse(job["video_path"], media_type="video/mp4")


@app.get("/api/video-stream/{job_id}")
async def video_stream(job_id: str):
    """SSE stream for video generation progress"""
    async def event_generator():
        while True:
            if job_id not in video_jobs:
                yield f"data: {json.dumps({'error': 'Job not found'})}\n\n"
                break

            job = video_jobs[job_id]
            yield f"data: {json.dumps(job)}\n\n"

            if job["status"] in ["completed", "error"]:
                break

            await asyncio.sleep(0.5)

    return StreamingResponse(event_generator(), media_type="text/event-stream")


@app.post("/api/conversation-with-video")
async def conversation_with_video(audio: UploadFile = File(...)):
    """
    Full conversation with video - waits for video to complete
    """
    start_time = time.time()
    job_id = uuid.uuid4().hex[:8]

    # Save and convert audio
    temp_webm = tempfile.NamedTemporaryFile(suffix=".webm", delete=False)
    content = await audio.read()
    temp_webm.write(content)
    temp_webm.close()

    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    temp_wav.close()
    convert_webm_to_wav(temp_webm.name, temp_wav.name)

    try:
        # STT
        stt_start = time.time()
        user_text = await transcribe_audio(temp_wav.name)
        stt_time = time.time() - stt_start

        # LLM
        llm_start = time.time()
        global conversation_history
        conversation_history.append({"role": "user", "content": user_text})
        messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:]
        assistant_text = await chat_llm(messages)
        conversation_history.append({"role": "assistant", "content": assistant_text})
        llm_time = time.time() - llm_start

        # TTS
        tts_start = time.time()
        tts_path = str(RESULTS_DIR / f"tts_{job_id}.mp3")
        await text_to_speech_elevenlabs(assistant_text, tts_path)
        tts_time = time.time() - tts_start

        # Convert to WAV for MuseTalk
        tts_wav = str(RESULTS_DIR / f"tts_{job_id}.wav")
        subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", tts_path, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True)

        # Generate video (blocking)
        video_start = time.time()
        video_path = str(RESULTS_DIR / f"video_{job_id}.mp4")
        engine.generate_video_fast(tts_wav, video_path)
        video_time = time.time() - video_start

        total_time = time.time() - start_time

        return {
            "job_id": job_id,
            "user_text": user_text,
            "assistant_text": assistant_text,
            "audio_url": f"/api/audio/{job_id}",
            "video_url": f"/api/video-file/{job_id}",
            "timing": {
                "stt": round(stt_time, 2),
                "llm": round(llm_time, 2),
                "tts": round(tts_time, 2),
                "video": round(video_time, 2),
                "total": round(total_time, 2)
            }
        }
    finally:
        try:
            os.unlink(temp_webm.name)
            os.unlink(temp_wav.name)
        except:
            pass


@app.get("/api/video-file/{job_id}")
async def get_video_file(job_id: str):
    video_path = RESULTS_DIR / f"video_{job_id}.mp4"
    if not video_path.exists():
        raise HTTPException(status_code=404, detail="Video not found")
    return FileResponse(video_path, media_type="video/mp4")


@app.post("/api/clear-history")
async def clear_history():
    global conversation_history
    conversation_history = []
    return {"status": "ok"}


@app.get("/streaming")
async def streaming_page():
    """Streaming video page"""
    return FileResponse(Path(__file__).parent / "static" / "index_streaming.html")


# WebSocket streaming state
streaming_sessions = {}


@app.websocket("/ws/stream")
async def websocket_stream(websocket: WebSocket):
    """
    WebSocket endpoint for real-time frame streaming.
    """
    await websocket.accept()
    session_id = uuid.uuid4().hex[:8]
    print(f"\n[WebSocket] Client connected: {session_id}")

    try:
        while True:
            data = await websocket.receive_json()

            if data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})
                continue

            if data.get("type") == "start":
                # Unique request ID for this specific request
                request_id = uuid.uuid4().hex[:8]
                metrics = SessionMetrics(request_id)

                audio_base64 = data.get("audio_base64")
                resolution = data.get("resolution", 256)
                batch_size = data.get("batch_size", 8)
                codec = data.get("codec", "jpeg").lower()  # "jpeg" or "h264"

                # Validate codec
                if codec == "h264" and not H264_AVAILABLE:
                    codec = "jpeg"
                    await websocket.send_json({
                        "type": "warning",
                        "message": "H.264 not available, using JPEG"
                    })

                print(f"\n{'='*50}")
                print(f"[WebSocket] New request: {request_id}")
                print(f"[WebSocket] Settings: resolution={resolution}, batch={batch_size}, codec={codec}")
                print(f"{'='*50}")

                # Decode and save audio
                audio_bytes = base64.b64decode(audio_base64)
                temp_audio = tempfile.NamedTemporaryFile(suffix=".webm", delete=False)
                temp_audio.write(audio_bytes)
                temp_audio.close()

                temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
                temp_wav.close()
                convert_webm_to_wav(temp_audio.name, temp_wav.name)

                try:
                    # STT
                    await websocket.send_json({"type": "status", "message": "Transcrevendo..."})
                    user_text = await transcribe_audio(temp_wav.name)
                    await websocket.send_json({"type": "transcription", "text": user_text})
                    print(f"[STT] {user_text}")
                    metrics.mark("stt_done")

                    # LLM
                    await websocket.send_json({"type": "status", "message": "Pensando..."})
                    global conversation_history
                    conversation_history.append({"role": "user", "content": user_text})
                    messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:]
                    assistant_text = await chat_llm(messages)
                    conversation_history.append({"role": "assistant", "content": assistant_text})
                    await websocket.send_json({"type": "response", "text": assistant_text})
                    print(f"[LLM] {assistant_text}")
                    metrics.mark("llm_done")

                    # TTS - PARALLEL PIPELINE
                    # 1. Generate espeak immediately (fast ~15ms) for video timing
                    # 2. Start ElevenLabs in parallel for high quality audio
                    await websocket.send_json({"type": "status", "message": "Gerando voz..."})

                    espeak_wav = str(RESULTS_DIR / f"stream_espeak_{request_id}.wav")
                    tts_path = str(RESULTS_DIR / f"stream_tts_{request_id}.mp3")

                    # Fast espeak for video generation
                    espeak_start = time.time()
                    text_to_speech_espeak(assistant_text, espeak_wav)
                    espeak_time = (time.time() - espeak_start) * 1000
                    metrics.mark("espeak_done")
                    print(f"[TTS] espeak: {espeak_time:.0f}ms")

                    # Start ElevenLabs in parallel (don't await yet)
                    elevenlabs_start = time.time()
                    elevenlabs_task = asyncio.create_task(
                        text_to_speech_elevenlabs(assistant_text, tts_path)
                    )

                    # Convert espeak for MuseTalk (16kHz)
                    tts_wav = str(RESULTS_DIR / f"stream_tts_{request_id}.wav")
                    subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", espeak_wav, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True)
                    metrics.mark("tts_done")

                    # Send audio URL (will serve ElevenLabs when ready, espeak as fallback)
                    await websocket.send_json({"type": "audio", "url": f"/api/stream-audio/{request_id}"})

                    # Stream frames (using espeak timing, video starts immediately!)
                    await websocket.send_json({"type": "status", "message": "Gerando video..."})
                    print(f"[VIDEO] Starting frame generation (codec={codec})...")

                    frame_count = 0
                    total_frames = 0
                    gen_start = time.time()
                    h264_encoder = None
                    total_bytes_sent = 0

                    for frame_data in engine.generate_frames_streaming(tts_wav, resolution=resolution, batch_size=batch_size):
                        if frame_data["type"] == "info":
                            total_frames = frame_data["total_frames"]
                            fps = frame_data["fps"]

                            # Initialize H.264 encoder if needed
                            if codec == "h264":
                                h264_encoder = H264StreamEncoder(
                                    width=resolution,
                                    height=resolution,
                                    fps=fps
                                )
                                # Send codec config (SPS/PPS) for decoder init
                                codec_config = h264_encoder.start()
                                if codec_config:
                                    await websocket.send_bytes(b'\x00' + codec_config)  # 0x00 = config
                                    print(f"[H264] Sent codec config: {len(codec_config)} bytes")

                            await websocket.send_json({
                                "type": "info",
                                "total_frames": total_frames,
                                "fps": fps,
                                "codec": codec
                            })
                            print(f"[VIDEO] Total frames: {total_frames}, FPS: {fps}")

                        elif frame_data["type"] == "frame":
                            frame = frame_data["frame"]

                            if codec == "h264" and h264_encoder:
                                # H.264 encoding (delta frames)
                                h264_data, meta = h264_encoder.encode(frame)
                                if h264_data:
                                    # Send as binary: 0x01=keyframe, 0x02=delta
                                    frame_type = b'\x01' if meta['keyframe'] else b'\x02'
                                    await websocket.send_bytes(frame_type + h264_data)
                                    total_bytes_sent += len(h264_data)
                            else:
                                # JPEG encoding (full frames)
                                _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
                                frame_base64 = base64.b64encode(buffer).decode('utf-8')
                                total_bytes_sent += len(buffer)

                                await websocket.send_json({
                                    "type": "frame",
                                    "frame": frame_base64,
                                    "index": frame_data["index"],
                                    "total": total_frames
                                })

                            frame_count += 1

                            # Yield to event loop every frame to ensure immediate send
                            await asyncio.sleep(0)

                            # Log first 5 frames to show timing
                            if frame_count == 1: metrics.mark("first_frame")
                            if frame_count <= 5:
                                elapsed = time.time() - gen_start
                                if codec == "h264":
                                    frame_type_str = "I" if meta.get('keyframe') else "P"
                                else:
                                    frame_type_str = "JPEG"
                                print(f"[VIDEO] Frame {frame_count} ({frame_type_str}) sent at {elapsed:.2f}s")
                            # Then log every batch
                            elif frame_count % batch_size == 0:
                                elapsed = time.time() - gen_start
                                print(f"[VIDEO] Batch complete: {frame_count}/{total_frames} frames at {elapsed:.2f}s")

                    # Cleanup H.264 encoder
                    if h264_encoder:
                        h264_encoder.stop()

                    gen_time = time.time() - gen_start
                    avg_frame_size = total_bytes_sent / max(1, frame_count)
                    print(f"[VIDEO] Complete: {frame_count} frames in {gen_time:.2f}s ({frame_count/gen_time:.1f} fps)")
                    print(f"[VIDEO] Bandwidth: {total_bytes_sent/1024:.1f}KB total, {avg_frame_size:.0f}B/frame avg ({codec})")

                    # Ensure ElevenLabs audio is ready for playback
                    try:
                        await elevenlabs_task
                        metrics.mark("elevenlabs_done")
                        elevenlabs_time = (time.time() - elevenlabs_start) * 1000
                        print(f"[TTS] ElevenLabs ready: {elevenlabs_time:.0f}ms")
                    except Exception as e:
                        print(f"[TTS] ElevenLabs error (using espeak fallback): {e}")

                    metrics.mark("all_done")

                    await websocket.send_json({
                        "type": "done",
                        "total_frames": frame_count,
                        "gen_time": round(gen_time, 2),
                        "codec": codec,
                        "total_bytes": total_bytes_sent,
                        "avg_frame_bytes": round(avg_frame_size),
                        "metrics": metrics.to_dict()
                    })

                finally:
                    try:
                        os.unlink(temp_audio.name)
                        os.unlink(temp_wav.name)
                    except:
                        pass

            elif data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})

    except WebSocketDisconnect:
        print(f"[WebSocket] Client disconnected: {session_id}")
    except Exception as e:
        print(f"[WebSocket] Error: {e}")
        import traceback
        traceback.print_exc()
        try:
            await websocket.send_json({"type": "error", "message": str(e)})
        except:
            pass


@app.get("/api/stream-audio/{session_id}")
async def get_stream_audio(session_id: str):
    """Get TTS audio for streaming session (ElevenLabs preferred, espeak fallback)"""
    # Try ElevenLabs first (high quality)
    elevenlabs_path = RESULTS_DIR / f"stream_tts_{session_id}.mp3"
    if elevenlabs_path.exists():
        return FileResponse(elevenlabs_path, media_type="audio/mpeg")

    # Fallback to espeak (if ElevenLabs not ready yet)
    espeak_path = RESULTS_DIR / f"stream_espeak_{session_id}.wav"
    if espeak_path.exists():
        return FileResponse(espeak_path, media_type="audio/wav")

    raise HTTPException(status_code=404, detail="Audio not found")


@app.get("/api/idle-frames")
async def get_idle_frames():
    """Get idle animation frames from idle.mp4 (for avatar listening state)"""
    global engine
    if engine is None:
        raise HTTPException(status_code=503, detail="Engine not initialized")

    idle_frames, idle_fps = engine.get_idle_frames()
    if not idle_frames:
        raise HTTPException(status_code=404, detail="No idle frames available")

    # Encode frames as base64 JPEG (RGB to BGR for cv2)
    frames_b64 = []
    for frame in idle_frames:
        # Convert RGB to BGR for cv2.imencode
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        _, buffer = cv2.imencode('.jpg', frame_bgr, [cv2.IMWRITE_JPEG_QUALITY, 85])
        frames_b64.append(base64.b64encode(buffer).decode('utf-8'))

    return {
        "frames": frames_b64,
        "fps": idle_fps,
        "count": len(frames_b64)
    }


# ============================================
# WebRTC Streaming
# ============================================
from server.webrtc_stream import webrtc_manager

class RTCOfferRequest(BaseModel):
    sdp: str
    type: str
    session_id: str

@app.post("/api/rtc/offer")
async def rtc_offer(request: RTCOfferRequest):
    """Handle WebRTC offer and return answer."""
    try:
        answer = await webrtc_manager.handle_offer(
            session_id=request.session_id,
            sdp=request.sdp,
            type=request.type,
            fps=30
        )
        return answer
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/rtc/close/{session_id}")
async def rtc_close(session_id: str):
    """Close WebRTC connection."""
    await webrtc_manager.close_connection(session_id)
    return {"status": "closed"}

@app.get("/rtc")
async def rtc_page():
    """Serve WebRTC test page."""
    return FileResponse(Path(__file__).parent / "static" / "index_rtc.html")

@app.get("/ws")
async def ws_page():
    """Serve WebSocket-only streaming page (proxy-friendly, no WebRTC)."""
    return FileResponse(Path(__file__).parent / "static" / "index_ws.html")

@app.get("/h264")
async def h264_page():
    """Serve H.264 WebCodecs streaming page (efficient delta frames)."""
    return FileResponse(Path(__file__).parent / "static" / "index_h264.html")

@app.get("/auto")
async def auto_page():
    """Auto-detect best streaming mode (WebRTC if UDP available, else H.264/JPEG)."""
    return FileResponse(Path(__file__).parent / "static" / "index_auto.html")

@app.get("/app/{full_path:path}")
async def react_app(full_path: str = ""):
    """Serve React SPA - all routes return index.html for client-side routing."""
    react_index = Path(__file__).parent / "web" / "dist" / "index.html"
    if react_index.exists():
        return FileResponse(react_index)
    return FileResponse(Path(__file__).parent / "static" / "index_auto.html")

@app.get("/app")
async def react_app_root():
    """React app root."""
    return await react_app("")

@app.websocket("/ws/rtc")
async def websocket_rtc(websocket: WebSocket):
    """
    WebSocket endpoint for WebRTC signaling + audio processing.
    Video is sent via WebRTC, signaling via WebSocket.
    """
    print(f"\n[WebRTC-WS] Accepting connection...")
    try:
        await websocket.accept()
    except Exception as e:
        print(f"[WebRTC-WS] Accept error: {e}")
        raise
    session_id = uuid.uuid4().hex[:8]
    print(f"[WebRTC-WS] Client connected: {session_id}")

    try:
        while True:
            data = await websocket.receive_json()

            if data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})
                continue

            if data.get("type") == "offer":
                # Handle WebRTC offer
                answer = await webrtc_manager.handle_offer(
                    session_id=session_id,
                    sdp=data["sdp"],
                    type="offer",
                    fps=30
                )
                await websocket.send_json({
                    "type": "answer",
                    "sdp": answer["sdp"]
                })
                print(f"[WebRTC-WS] Sent answer to {session_id}")

            elif data.get("type") == "start":
                # Process audio and stream video via WebRTC
                request_id = uuid.uuid4().hex[:8]
                metrics = SessionMetrics(request_id)

                audio_base64 = data.get("audio_base64")
                resolution = data.get("resolution", 256)
                batch_size = data.get("batch_size", 16)

                print(f"\n{'='*50}")
                print(f"[WebRTC] New request: {request_id}")
                print(f"[WebRTC] Settings: resolution={resolution}, batch={batch_size}")
                print(f"{'='*50}")

                # Decode audio
                audio_bytes = base64.b64decode(audio_base64)
                temp_audio = tempfile.NamedTemporaryFile(suffix=".webm", delete=False)
                temp_audio.write(audio_bytes)
                temp_audio.close()

                temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
                temp_wav.close()
                convert_webm_to_wav(temp_audio.name, temp_wav.name)

                try:
                    # STT
                    await websocket.send_json({"type": "status", "message": "Transcrevendo..."})
                    user_text = await transcribe_audio(temp_wav.name)
                    metrics.mark("stt_done")
                    await websocket.send_json({"type": "transcription", "text": user_text})
                    print(f"[STT] {user_text}")

                    # LLM
                    await websocket.send_json({"type": "status", "message": "Pensando..."})
                    global conversation_history
                    conversation_history.append({"role": "user", "content": user_text})
                    messages = [{"role": "system", "content": "Voce e uma assistente virtual amigavel. Responda de forma concisa (1-3 frases)."}] + conversation_history[-10:]
                    assistant_text = await chat_llm(messages)
                    conversation_history.append({"role": "assistant", "content": assistant_text})
                    metrics.mark("llm_done")
                    await websocket.send_json({"type": "response", "text": assistant_text})
                    print(f"[LLM] {assistant_text}")

                    # TTS - PARALLEL PIPELINE
                    # 1. Generate espeak immediately (fast ~15ms) for video timing
                    # 2. Start ElevenLabs in parallel for high quality audio
                    await websocket.send_json({"type": "status", "message": "Gerando voz..."})

                    espeak_wav = str(RESULTS_DIR / f"stream_espeak_{request_id}.wav")
                    tts_path = str(RESULTS_DIR / f"stream_tts_{request_id}.mp3")

                    # Fast espeak for video generation
                    espeak_start = time.time()
                    text_to_speech_espeak(assistant_text, espeak_wav)
                    espeak_time = (time.time() - espeak_start) * 1000
                    metrics.mark("espeak_done")
                    print(f"[TTS] espeak: {espeak_time:.0f}ms")

                    # Start ElevenLabs in parallel (don't await yet)
                    elevenlabs_start = time.time()
                    elevenlabs_task = asyncio.create_task(
                        text_to_speech_elevenlabs(assistant_text, tts_path)
                    )

                    # Convert espeak for MuseTalk (16kHz)
                    tts_wav = str(RESULTS_DIR / f"stream_tts_{request_id}.wav")
                    subprocess.run(["ffmpeg", "-y", "-v", "quiet", "-i", espeak_wav, "-ar", "16000", "-ac", "1", tts_wav], capture_output=True)
                    metrics.mark("tts_done")

                    # Send audio URL (will serve ElevenLabs when ready)
                    await websocket.send_json({"type": "audio", "url": f"/api/stream-audio/{request_id}"})

                    # Check WebRTC connection or use WebSocket fallback
                    use_websocket_video = data.get("ws_video", False)  # Client can request WS video
                    rtc_connected = webrtc_manager.is_connected(session_id)

                    if not use_websocket_video and not rtc_connected:
                        await websocket.send_json({"type": "status", "message": "Conectando WebRTC..."})
                        print(f"[WebRTC] Waiting for connection to be established...")
                        rtc_connected = await webrtc_manager.wait_for_connection(session_id, timeout=5.0)

                    if not rtc_connected:
                        use_websocket_video = True
                        await websocket.send_json({"type": "ws_video_mode", "enabled": True})
                        print(f"[WS-VIDEO] WebRTC not available, using WebSocket video fallback")

                    # Stream frames
                    await websocket.send_json({"type": "status", "message": "Gerando video..."})
                    mode = "WS-VIDEO" if use_websocket_video else "WebRTC"
                    print(f"[{mode}] Starting frame generation...")

                    frame_count = 0
                    total_frames = 0
                    gen_start = time.time()

                    for frame_data in engine.generate_frames_streaming(tts_wav, resolution=resolution, batch_size=batch_size):
                        if frame_data["type"] == "info":
                            total_frames = frame_data["total_frames"]
                            fps = frame_data["fps"]
                            await websocket.send_json({
                                "type": "info",
                                "total_frames": total_frames,
                                "fps": fps
                            })
                            print(f"[{mode}] Total frames: {total_frames}, FPS: {fps}")

                        elif frame_data["type"] == "frame":
                            frame = frame_data["frame"]

                            if use_websocket_video:
                                # Send frame as JPEG base64 via WebSocket
                                _, jpeg_data = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
                                frame_b64 = base64.b64encode(jpeg_data).decode('utf-8')
                                await websocket.send_json({
                                    "type": "video_frame",
                                    "frame": frame_b64,
                                    "index": frame_count
                                })
                            else:
                                # Send frame via WebRTC
                                await webrtc_manager.send_frame(session_id, frame)

                            frame_count += 1

                            if frame_count == 1:
                                metrics.mark("first_frame")

                            # Small delay to match video FPS
                            await asyncio.sleep(1.0 / 30.0)

                            if frame_count <= 5:
                                elapsed = time.time() - gen_start
                                print(f"[{mode}] Frame {frame_count} sent at {elapsed:.2f}s")
                            elif frame_count % batch_size == 0:
                                elapsed = time.time() - gen_start
                                print(f"[{mode}] Batch complete: {frame_count}/{total_frames} frames at {elapsed:.2f}s")

                    gen_time = time.time() - gen_start
                    print(f"[{mode}] Complete: {frame_count} frames in {gen_time:.2f}s ({frame_count/gen_time:.1f} fps)")

                    # Ensure ElevenLabs audio is ready for playback
                    try:
                        await elevenlabs_task
                        metrics.mark("elevenlabs_done")
                        elevenlabs_time = (time.time() - elevenlabs_start) * 1000
                        print(f"[TTS] ElevenLabs ready: {elevenlabs_time:.0f}ms")
                    except Exception as e:
                        print(f"[TTS] ElevenLabs error (using espeak fallback): {e}")

                    metrics.mark("all_done")

                    await websocket.send_json({
                        "type": "done",
                        "total_frames": frame_count,
                        "gen_time": round(gen_time, 2),
                        "metrics": metrics.to_dict()
                    })

                finally:
                    try:
                        os.unlink(temp_audio.name)
                        os.unlink(temp_wav.name)
                    except:
                        pass

    except WebSocketDisconnect:
        print(f"[WebRTC-WS] Client disconnected: {session_id}")
        await webrtc_manager.close_connection(session_id)
    except Exception as e:
        print(f"[WebRTC-WS] Error: {e}")
        await webrtc_manager.close_connection(session_id)


# ============================================
# Main entry point
# ============================================
if __name__ == "__main__":
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8000
    uvicorn.run(app, host="0.0.0.0", port=port)