server/main.py

"""
MuseTalk Speech-to-Speech Avatar Server
Streaming architecture with Groq (STT + LLM) and ElevenLabs (TTS)
"""

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

# Add parent directory to path for MuseTalk imports
sys.path.insert(0, str(Path(__file__).parent.parent))


def convert_webm_to_wav(input_path: str, output_path: str) -> str:
    """Convert webm audio to wav using ffmpeg"""
    cmd = [
        "ffmpeg", "-y", "-i", input_path,
        "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le",
        output_path
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"FFmpeg error: {result.stderr}")
    return output_path

from fastapi import FastAPI, UploadFile, File, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from sse_starlette.sse import EventSourceResponse
import asyncio
from pydantic import BaseModel
import uvicorn
import httpx
import numpy as np

# API Keys
GROQ_API_KEY = "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"
ELEVENLABS_API_KEY = "sk_857e9e6f2412ddf3ff5334b736e4b571641d26225c0d8d62"
ELEVENLABS_VOICE_ID = "21m00Tcm4TlvDq8ikWAM"  # Rachel voice, change as needed

# MuseTalk paths
MUSETALK_DIR = Path(__file__).parent.parent
AVATAR_VIDEO = MUSETALK_DIR / "data" / "video" / "avatar.mp4"
RESULTS_DIR = MUSETALK_DIR / "results" / "server"
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

app = FastAPI(title="MuseTalk Speech-to-Speech API")

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

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

# Global MuseTalk model (loaded once)
musetalk_model = None
conversation_history = []


class ChatMessage(BaseModel):
    role: str
    content: str


class TextRequest(BaseModel):
    text: str


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


async def chat_groq(messages: list) -> str:
    """Chat with Groq Llama 8B"""
    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_groq_streaming(messages: list):
    """Stream chat response from Groq Llama 8B"""
    async with httpx.AsyncClient(timeout=60.0) as client:
        async with client.stream(
            "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,
                "stream": True
            }
        ) as response:
            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    data = line[6:]
                    if data == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data)
                        if chunk["choices"][0].get("delta", {}).get("content"):
                            yield chunk["choices"][0]["delta"]["content"]
                    except json.JSONDecodeError:
                        pass


async def text_to_speech_elevenlabs(text: str, output_path: str) -> str:
    """Convert text to speech using ElevenLabs"""
    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}",
            headers={
                "xi-api-key": ELEVENLABS_API_KEY,
                "Content-Type": "application/json"
            },
            json={
                "text": text,
                "model_id": "eleven_multilingual_v2",
                "voice_settings": {
                    "stability": 0.5,
                    "similarity_boost": 0.75
                }
            }
        )
        if response.status_code != 200:
            raise HTTPException(status_code=500, detail=f"ElevenLabs TTS error: {response.text}")

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


async def text_to_speech_elevenlabs_streaming(text: str, output_path: str) -> str:
    """Stream text to speech using ElevenLabs for faster first byte"""
    async with httpx.AsyncClient(timeout=60.0) as client:
        async with client.stream(
            "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_multilingual_v2",
                "voice_settings": {
                    "stability": 0.5,
                    "similarity_boost": 0.75
                }
            }
        ) as response:
            with open(output_path, "wb") as f:
                async for chunk in response.aiter_bytes():
                    f.write(chunk)
    return output_path


def generate_lipsync_video(audio_path: str, output_path: str) -> str:
    """Generate lip-sync video using MuseTalk"""
    import yaml

    # Create temporary config with the audio path
    temp_config = RESULTS_DIR / f"config_{uuid.uuid4().hex[:8]}.yaml"
    config_data = {
        "task_0": {
            "video_path": str(AVATAR_VIDEO),
            "audio_path": audio_path
        }
    }
    with open(temp_config, "w") as f:
        yaml.dump(config_data, f)

    # Use MuseTalk inference script
    cmd = [
        "python3", "-m", "scripts.inference",
        "--inference_config", str(temp_config),
        "--result_dir", str(RESULTS_DIR),
        "--ffmpeg_path", "/usr/bin",
        "--vae_type", "sd-vae-ft-mse",
        "--unet_config", str(MUSETALK_DIR / "models" / "musetalk" / "musetalk.json"),
        "--batch_size", "8"
    ]

    env = os.environ.copy()
    env["FFMPEG_PATH"] = "/usr/bin"

    result = subprocess.run(
        cmd,
        cwd=str(MUSETALK_DIR),
        capture_output=True,
        text=True,
        env=env
    )

    # Cleanup temp config
    try:
        temp_config.unlink()
    except:
        pass

    if result.returncode != 0:
        print(f"MuseTalk error: {result.stderr}")
        print(f"MuseTalk stdout: {result.stdout}")
        raise HTTPException(status_code=500, detail=f"MuseTalk error: {result.stderr}")

    # Find the generated video
    video_files = list(RESULTS_DIR.glob("*.mp4"))
    if video_files:
        latest = max(video_files, key=lambda p: p.stat().st_mtime)
        return str(latest)

    raise HTTPException(status_code=500, detail="No video generated")


@app.get("/")
async def root():
    """Serve main HTML page"""
    return FileResponse(Path(__file__).parent / "static" / "index.html")


@app.post("/api/transcribe")
async def transcribe(audio: UploadFile = File(...)):
    """Transcribe audio to text using Groq Whisper"""
    # Save uploaded audio
    temp_audio = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    content = await audio.read()
    temp_audio.write(content)
    temp_audio.close()

    try:
        text = await transcribe_audio_groq(temp_audio.name)
        return {"text": text}
    finally:
        os.unlink(temp_audio.name)


@app.post("/api/chat")
async def chat(request: TextRequest):
    """Chat with LLM and return text response"""
    global conversation_history

    # Add user message to history
    conversation_history.append({"role": "user", "content": request.text})

    # System prompt for avatar
    messages = [
        {
            "role": "system",
            "content": "Voce e uma assistente virtual amigavel e prestativa. Responda de forma concisa e natural, como em uma conversa. Mantenha respostas curtas (1-3 frases) para uma experiencia de conversa fluida."
        }
    ] + conversation_history[-10:]  # Keep last 10 messages for context

    # Get response
    response_text = await chat_groq(messages)

    # Add assistant response to history
    conversation_history.append({"role": "assistant", "content": response_text})

    return {"text": response_text}


@app.post("/api/tts")
async def tts(request: TextRequest):
    """Convert text to speech and return audio file"""
    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/generate-video")
async def generate_video(audio: UploadFile = File(...)):
    """Generate lip-sync video from audio"""
    # Save uploaded audio
    temp_audio = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    content = await audio.read()
    temp_audio.write(content)
    temp_audio.close()

    try:
        output_path = str(RESULTS_DIR / f"video_{uuid.uuid4().hex[:8]}.mp4")
        video_path = generate_lipsync_video(temp_audio.name, output_path)
        return FileResponse(video_path, media_type="video/mp4")
    finally:
        os.unlink(temp_audio.name)


# Store for pending video jobs
pending_videos = {}


def generate_video_background(job_id: str, tts_path: str):
    """Generate lip-sync video in background thread"""
    try:
        pending_videos[job_id] = {"status": "generating", "path": None, "progress": 0}

        # Convert mp3 to wav for MuseTalk
        wav_path = tts_path.replace('.mp3', '.wav')
        subprocess.run([
            "ffmpeg", "-y", "-v", "quiet",
            "-i", tts_path,
            "-ar", "16000", "-ac", "1",
            wav_path
        ], capture_output=True)

        # Generate video
        video_path = generate_lipsync_video(wav_path, "")

        pending_videos[job_id] = {"status": "completed", "path": video_path, "progress": 100}
        print(f"Video {job_id} completed: {video_path}")

    except Exception as e:
        print(f"Video generation error: {e}")
        pending_videos[job_id] = {"status": "error", "error": str(e)}


@app.post("/api/conversation")
async def full_conversation(audio: UploadFile = File(...)):
    """
    Fast speech-to-speech pipeline:
    Returns audio immediately, video generates in background
    """
    start_time = time.time()
    job_id = uuid.uuid4().hex[:8]

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

    # Convert to wav for Groq
    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    temp_wav.close()
    convert_webm_to_wav(temp_webm.name, temp_wav.name)

    try:
        # 1. Transcribe audio (STT) - ~0.2s with Groq
        stt_start = time.time()
        user_text = await transcribe_audio_groq(temp_wav.name)
        stt_time = time.time() - stt_start
        print(f"STT ({stt_time:.2f}s): {user_text}")

        # 2. Get LLM response - ~0.2s with Groq
        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 e natural (1-3 frases)."
            }
        ] + conversation_history[-10:]

        assistant_text = await chat_groq(messages)
        conversation_history.append({"role": "assistant", "content": assistant_text})
        llm_time = time.time() - llm_start
        print(f"LLM ({llm_time:.2f}s): {assistant_text}")

        # 3. Generate TTS audio - ~1s with ElevenLabs
        tts_start = time.time()
        tts_path = str(RESULTS_DIR / f"tts_{job_id}.mp3")
        await text_to_speech_elevenlabs_streaming(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 (without video)")

        # 4. Start video generation in background
        thread = threading.Thread(target=generate_video_background, args=(job_id, tts_path))
        thread.start()

        # Return immediately with audio
        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):
    """Serve TTS audio"""
    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 pending_videos:
        return {"status": "not_found"}
    return pending_videos[job_id]


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

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

    video_path = job["path"]
    if not video_path or not Path(video_path).exists():
        raise HTTPException(status_code=404, detail="Video file not found")

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


@app.post("/api/conversation-full")
async def full_conversation_with_video(audio: UploadFile = File(...)):
    """
    Full speech-to-speech pipeline with video (slower but complete)
    """
    start_time = time.time()

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

    # Convert to wav for Groq
    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    temp_wav.close()
    convert_webm_to_wav(temp_webm.name, temp_wav.name)

    try:
        # 1. Transcribe audio (STT)
        stt_start = time.time()
        user_text = await transcribe_audio_groq(temp_wav.name)
        print(f"STT ({time.time() - stt_start:.2f}s): {user_text}")

        # 2. Get LLM response
        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 e natural (1-3 frases)."
            }
        ] + conversation_history[-10:]

        assistant_text = await chat_groq(messages)
        conversation_history.append({"role": "assistant", "content": assistant_text})
        print(f"LLM ({time.time() - llm_start:.2f}s): {assistant_text}")

        # 3. Generate TTS audio
        tts_start = time.time()
        tts_path = str(RESULTS_DIR / f"tts_{uuid.uuid4().hex[:8]}.mp3")
        await text_to_speech_elevenlabs_streaming(assistant_text, tts_path)
        print(f"TTS ({time.time() - tts_start:.2f}s)")

        # 4. Generate lip-sync video
        lipsync_start = time.time()
        video_path = generate_lipsync_video(tts_path, "")
        print(f"LipSync ({time.time() - lipsync_start:.2f}s)")

        total_time = time.time() - start_time
        print(f"Total pipeline: {total_time:.2f}s")

        return {
            "user_text": user_text,
            "assistant_text": assistant_text,
            "video_url": f"/api/video/{Path(video_path).name}",
            "timing": {
                "stt": round(time.time() - stt_start, 2),
                "llm": round(time.time() - llm_start, 2),
                "tts": round(time.time() - tts_start, 2),
                "lipsync": round(time.time() - lipsync_start, 2),
                "total": round(total_time, 2)
            }
        }
    finally:
        os.unlink(temp_webm.name)
        os.unlink(temp_wav.name)


@app.get("/api/video/{filename}")
async def get_video(filename: str):
    """Serve generated video"""
    video_path = RESULTS_DIR / filename
    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():
    """Clear conversation history"""
    global conversation_history
    conversation_history = []
    return {"status": "ok"}


@app.get("/api/health")
async def health():
    """Health check"""
    return {"status": "ok", "avatar": str(AVATAR_VIDEO.exists())}


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)