rtc_handler.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
import asyncio
from video_handler import VideoHandler, AvatarState
from ai_handler import AIHandler
import json
from typing import Set
import logging
import os


from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

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


# Mount static files directory for serving HTML, JS, CSS files
# RecroTech widget demo için static dosyaları serve et
app.mount("/static", StaticFiles(directory="static"), name="static")

# Mount avatar_videos directory for serving avatar video files
# Avatar videoları /avatar_videos path'i üzerinden servis edilir
app.mount("/avatar_videos", StaticFiles(directory="avatar_videos"), name="avatar_videos")


class RTCHandler:
    def __init__(self, video_dir: str):
        """
        Initialize RTC handler for managing avatar video streams.

        Args:
            video_dir: Directory containing avatar videos
        """
        self.video_handler = VideoHandler(video_dir)
        self.active_connections: Set[WebSocket] = set()

    async def connect(self, websocket: WebSocket):
        """Accept and register a new WebSocket connection."""
        await websocket.accept()
        self.active_connections.add(websocket)
        logger.info(f"Client connected. Total connections: {len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket):
        """Remove a WebSocket connection."""
        self.active_connections.discard(websocket)
        logger.info(f"Client disconnected. Total connections: {len(self.active_connections)}")

    async def broadcast_frame(self, frame_bytes: bytes):
        """Broadcast a video frame to all connected clients."""
        disconnected = set()

        for connection in self.active_connections:
            try:
                await connection.send_bytes(frame_bytes)
            except Exception as e:
                logger.error(f"Error sending frame: {e}")
                disconnected.add(connection)

        # Remove disconnected clients
        for conn in disconnected:
            self.disconnect(conn)

    async def handle_message(self, websocket: WebSocket, message: str):
        """
        Handle incoming WebSocket messages to control avatar state.

        Expected message format: {"action": "standing|listening|thinking|talking"}
        """
        try:
            data = json.loads(message)
            action = data.get("action", "").lower()

            if action == "standing":
                self.video_handler.play_standing()
            elif action == "listening":
                self.video_handler.play_listening()
            elif action == "thinking":
                self.video_handler.play_thinking()
            elif action == "talking":
                self.video_handler.play_talking()
            else:
                await websocket.send_json({
                    "error": f"Unknown action: {action}",
                    "valid_actions": ["standing", "listening", "thinking", "talking"]
                })
                return

            # Confirm state change
            await websocket.send_json({
                "status": "success",
                "current_state": self.video_handler.get_current_state().value
            })

            logger.info(f"Avatar state changed to: {action}")

        except json.JSONDecodeError:
            await websocket.send_json({"error": "Invalid JSON message"})
        except Exception as e:
            logger.error(f"Error handling message: {e}")
            await websocket.send_json({"error": str(e)})


# Initialize handlers
rtc_handler = RTCHandler(video_dir="./avatar_videos")

system_prompt = """
You are a helpful AI assistant. Respond to user queries in max 2 sentences.
Be friendly, concise, and natural in your responses.
"""
ai_handler = AIHandler(system_prompt, rtc_handler=rtc_handler)


@app.get("/")
async def root():
    """Health check endpoint."""
    return {
        "status": "running",
        "current_state": rtc_handler.video_handler.get_current_state().value,
        "is_speaking": ai_handler.is_speaking
    }


@app.websocket("/ws/avatar")
async def websocket_endpoint(websocket: WebSocket):
    """
    WebSocket endpoint for avatar video streaming.

    Clients can send JSON messages to control avatar state:
    {"action": "standing"} - Switch to standing video
    {"action": "listening"} - Switch to listening video
    {"action": "thinking"} - Switch to thinking video
    {"action": "talking"} - Switch to talking video
    {"action": "writing"} - Switch to writing video
    """
    await rtc_handler.connect(websocket)

    # Start streaming task
    stream_task = asyncio.create_task(
        rtc_handler.video_handler.stream_frames(rtc_handler.broadcast_frame)
    )

    try:
        while True:
            # Receive messages from client
            message = await websocket.receive_text()
            await rtc_handler.handle_message(websocket, message)

    except WebSocketDisconnect:
        rtc_handler.disconnect(websocket)
        logger.info("Client disconnected normally")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
        rtc_handler.disconnect(websocket)
    finally:
        stream_task.cancel()
        try:
            await stream_task
        except asyncio.CancelledError:
            pass


@app.get("/stream/avatar")
async def http_stream():
    """
    HTTP endpoint for MJPEG video streaming.
    Can be used directly in HTML <img> tags.
    """

    async def generate():
        while True:
            frame = rtc_handler.video_handler.get_frame()
            if frame:
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
            await asyncio.sleep(0.033)  # ~30 FPS

    return StreamingResponse(
        generate(),
        media_type="multipart/x-mixed-replace; boundary=frame"
    )


@app.post("/avatar/state/{state}")
async def set_avatar_state(state: str):
    """
    REST endpoint to change avatar state.

    Args:
        state: One of "standing", "listening", "thinking", "talking"
    """
    state = state.lower()

    if state == "standing":
        rtc_handler.video_handler.play_standing()
    elif state == "listening":
        rtc_handler.video_handler.play_listening()
    elif state == "thinking":
        rtc_handler.video_handler.play_thinking()
    elif state == "talking":
        rtc_handler.video_handler.play_talking()
    elif state == "writing":
        rtc_handler.video_handler.play_writing()
    else:
        return {
            "error": f"Invalid state: {state}",
            "valid_states": ["standing", "listening", "thinking", "talking", "writing"]
        }, 400

    return {
        "status": "success",
        "current_state": rtc_handler.video_handler.get_current_state().value
    }


@app.post("/ai/speak")
async def speak_text(text: str):
    """
    Make the AI avatar speak the given text.
    Avatar will switch to talking state and play audio.
    """
    asyncio.create_task(ai_handler.speak_with_avatar(text))
    return {
        "status": "speaking",
        "text": text
    }


# Replace your /ai/chat endpoint with this debugged version:
# Replace your /ai/chat endpoint with this fixed version:

@app.post("/ai/chat")
async def chat(message: str):
    """
    Send a text message to AI and get TEXT response only.
    Avatar will animate: LISTENING -> THINKING -> TALKING -> STANDING
    """
    print(f"\n{'=' * 60}")
    print(f"💬 CHAT ENDPOINT CALLED")
    print(f"Message: {message}")
    print(f"{'=' * 60}")

    # 1. Set to listening
    print("🎭 Setting avatar to LISTENING state...")
    rtc_handler.video_handler.play_listening()
    await asyncio.sleep(0.5)  # Brief pause

    # 2. Set to thinking
    print("🎭 Setting avatar to THINKING state...")
    rtc_handler.video_handler.play_thinking()
    print(f"🎭 Current state: {rtc_handler.video_handler.get_current_state().value}")

    # Get AI response
    print("🤖 Getting AI response...")
    response = await ai_handler.send_message_get_answer(message)

    if response:
        # 3. Set to talking
        print("🎭 Setting avatar to TALKING state...")
        rtc_handler.video_handler.play_talking()

        # Estimate talking duration based on word count
        word_count = len(response.split())
        # Assume ~2.5 words per second speaking rate
        talking_duration = max(2.0, word_count / 2.5)  # Minimum 2 seconds

        print(f"🗣️ Talking for ~{talking_duration:.1f} seconds ({word_count} words)")

        # Schedule return to standing after talking
        async def return_to_standing():
            await asyncio.sleep(talking_duration)
            rtc_handler.video_handler.play_standing()
            print("🎭 Returned to STANDING state")

        asyncio.create_task(return_to_standing())

        return {
            "status": "success",
            "user_message": message,
            "ai_response": response,
            "talking_duration": talking_duration
        }
    else:
        print("❌ Failed to get AI response")
        rtc_handler.video_handler.play_standing()
        return {
            "status": "error",
            "message": "Failed to get AI response"
        }


@app.post("/ai/voice")
async def process_voice(audio: UploadFile = File(...)):
    """
    Upload voice audio and get spoken AI response.
    Complete workflow: STT -> AI -> TTS with avatar animations.
    Returns user transcript and AI response for chat display.
    """
    try:
        print(f"\n{'=' * 60}")
        print(f"🎤 VOICE ENDPOINT CALLED")
        print(f"{'=' * 60}")

        # Save uploaded audio temporarily
        temp_audio = f"temp_input_{audio.filename}"
        with open(temp_audio, "wb") as f:
            f.write(await audio.read())

        # Set to listening
        print("🎭 Setting avatar to LISTENING state...")
        rtc_handler.video_handler.play_listening()

        # Transcribe user speech (STT)
        print("🎤 Transcribing audio...")
        user_text = await ai_handler.transcribe_audio(temp_audio)

        if not user_text:
            rtc_handler.video_handler.play_standing()
            return {
                "status": "error",
                "message": "Ses metne dönüştürülemedi"
            }

        logger.info(f"🎤 Transcription: {user_text}")

        # Set to thinking
        print("🎭 Setting avatar to THINKING state...")
        rtc_handler.video_handler.play_thinking()
        print(f"🎭 Current state: {rtc_handler.video_handler.get_current_state().value}")

        # Get AI response
        print("🤖 Getting AI response...")
        ai_response = await ai_handler.send_message_get_answer(user_text)

        if not ai_response:
            rtc_handler.video_handler.play_standing()
            return {
                "status": "error",
                "message": "AI yanıt üretemedi",
                "stt_text": user_text
            }

        logger.info(f"🤖 AI Response: {ai_response}")

        # Generate audio for frontend
        print("🔊 Generating TTS audio...")
        audio_url = await ai_handler.generate_audio_for_frontend(ai_response)

        # Frontend handles audio playback and avatar state - no server-side playback
        print("✅ Voice processing complete")

        # Clean up temp file
        try:
            os.remove(temp_audio)
        except:
            pass

        return {
            "status": "success",
            "stt_text": user_text,
            "ai_response": ai_response,
            "audio_url": audio_url
        }

    except Exception as e:
        logger.error(f"Voice processing error: {e}")
        import traceback
        traceback.print_exc()
        rtc_handler.video_handler.play_standing()
        return {
            "status": "error",
            "message": str(e)
        }


if __name__ == "__main__":
    import uvicorn