app/websocket_manager.py

from fastapi import WebSocket, WebSocketDisconnect
from app.ai_responder import ai_answer
from app.session_manager import get_session
import json

SALES_NUMBER = "+91 7249477972"

class ConnectionManager:
    def __init__(self):
        self.active_connections = {}

    async def connect(self, websocket: WebSocket, session_id: str):
        await websocket.accept()
        self.active_connections[session_id] = websocket

    def disconnect(self, session_id: str):
        self.active_connections.pop(session_id, None)

    async def receive_audio(self, session_id: str):
        websocket = self.active_connections.get(session_id)  
        session = get_session(session_id)
        if not session:
            print(f"Session {session_id} not found, disconnecting...")
        print(f"current active connections: {self.active_connections}")
        if websocket:
            try:
                while True:
                    message = await websocket.receive_json()
                    print(f"Received message: {message}")
                    message_type = message.get("type")
                    print(f"Received message type: {message_type}")

                    if message_type == "setup": 
                        print(f"This is the setup message:- {message}")                  
                        print(f"Setup message received for session {session_id}.")

                    elif message_type == "prompt":
                        user_input = message["voicePrompt"]
                        print(f"Received transcription:- {user_input}")
                        ai_response = await ai_answer(user_input)
                        lower = ai_answer.lower()
                        if "sales" in lower:
                            dept = SALES_NUMBER
                        else:
                            dept = None
                        
                        if dept:
                            await websocket.send_text(json.dumps({
                                "type": "handoff",
                                "dialNumber": dept
                            }))
                            break  # end session so Twilio exits ConversationRelay
                        await websocket.send_json({
                            "type": "text",
                            "token": ai_response,
                            "last": True  # Assuming this is the last message in the conversation
                        })

                    elif message_type == "interrupt":
                        print(f"Interrupt message received for session {session_id}.")
                        # break
                        continue
                    
            except WebSocketDisconnect:
                self.disconnect(session_id)


manager = ConnectionManager()
async def websocket_endpoint(websocket: WebSocket, session_id: str):
    await manager.connect(websocket, session_id)
    await manager.receive_audio(session_id)