app.py

#!/usr/bin/env python3
#
#  san_integration_script.py (v5 - Provider Format Fix)
#  ======================================================
#  Description:
#  - Establishes a real-time, two-way audio bridge between a SAN system
#    and the Millis AI platform.
#  - Dynamically detects the audio format from the SAN `start` event.
#  - Forwards inbound audio to Millis AI at 16kHz for processing.
#  - Receives the AI's audio response at 16kHz.
#  - Streams the audio back to the SAN system using the exact format
#    it originally specified.
#
#  Changes in this version:
#  - Fixed the `reverse-media` event payload to match the provider's
#    expected format (simplified JSON, lowercase 'callid').
#  - Fixed `ImportError` by changing `starlette.websockets.State` to
#    `starlette.websockets.WebSocketState`.
#  - Updated the final connection check to use `WebSocketState.DISCONNECTED`.
#  - Added comprehensive logging for debugging audio flow.
#  - Fixed agent ID extraction to use platform_agent_id from extraParams.
#  - Added MongoDB integration for credential management.
#  -------------------------------------------------------------------

import os
import json
import base64
import asyncio
import logging
from datetime import datetime
from typing import Optional, Dict, Any
from collections import deque
from asyncio_throttle import Throttler

# Load environment variables from .env file
from dotenv import load_dotenv
load_dotenv()

# Third-party libraries
import numpy as np
from scipy import signal as scipy_signal
import websockets
from websockets.connection import State as WsState
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import uvicorn
from motor.motor_asyncio import AsyncIOMotorClient

# Import WebSocketState instead of State
from starlette.websockets import WebSocketState

# ---------- Logging Configuration -----------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("san-integration-app")

# ---------- Environment & Configuration -----------------------------------
AGENT_ID = "dummy"
PUBLIC_KEY = os.getenv("MILLIS_PUBLIC_KEY")
MILLIS_WS_URI = "wss://api-west.millis.ai:8080/millis"

# ---------- MongoDB Configuration -----------------------------------------
MONGODB_CONNECTION_STRING = os.getenv("MONGODB_CONNECTION_STRING", "mongodb://43.204.206.231:27017")
MONGODB_DATABASE_NAME = os.getenv("MONGODB_DATABASE_NAME", "masterDB")
MONGODB_COLLECTION = os.getenv("MONGODB_COLLECTION", "call_metadata")

# ---------- FastAPI Application -------------------------------------------
app = FastAPI()

def validate_environment():
    """Validate that all required environment variables are set."""
    required_vars = {
        "MILLIS_AGENT_ID": AGENT_ID,
        "MILLIS_PUBLIC_KEY": PUBLIC_KEY,
        "MONGODB_CONNECTION_STRING": MONGODB_CONNECTION_STRING,
        "MONGODB_DATABASE_NAME": MONGODB_DATABASE_NAME,
        "MONGODB_COLLECTION": MONGODB_COLLECTION
    }
    
    missing_vars = [var for var, value in required_vars.items() if not value]
    
    if missing_vars:
        error_msg = f"Missing required environment variables: {', '.join(missing_vars)}"
        logger.error(error_msg)
        logger.error("Please set these variables in your .env file or environment")
        raise ValueError(error_msg)

@app.on_event("startup")
async def startup_event():
    """Initialize MongoDB connection on startup."""
    logger.info("=== APPLICATION STARTUP ===")
    
    # Validate environment variables first
    validate_environment()
    
    success = await connect_mongodb()
    if not success:
        logger.error("Failed to connect to MongoDB during startup")
        logger.warning("Application will continue running but MongoDB features will be disabled")
        logger.info("To fix this issue:")
        logger.info("1. Ensure MongoDB is installed and running")
        logger.info("2. Set MONGODB_CONNECTION_STRING environment variable")
        logger.info("3. For local development: mongodb://localhost:27017")
    else:
        logger.info("MongoDB connection established successfully")
    logger.info("=== STARTUP COMPLETE ===")

@app.on_event("shutdown")
async def shutdown_event():
    """Close MongoDB connection on shutdown."""
    logger.info("=== APPLICATION SHUTDOWN ===")
    await close_mongodb()
    logger.info("=== SHUTDOWN COMPLETE ===")

# MongoDB client instance
mongodb_client: Optional[AsyncIOMotorClient] = None
mongodb_db = None

async def connect_mongodb():
    """Initialize MongoDB connection."""
    global mongodb_client, mongodb_db
    try:
        logger.info(f"Connecting to MongoDB at {MONGODB_CONNECTION_STRING}")
        
        if not MONGODB_CONNECTION_STRING or MONGODB_CONNECTION_STRING == "MONGODB_DATABASE_NAME":
            logger.error("Invalid MongoDB connection string. Please set MONGODB_CONNECTION_STRING environment variable.")
            logger.error("Example: mongodb://localhost:27017 or mongodb://username:password@host:port/database")
            return False
            
        mongodb_client = AsyncIOMotorClient(MONGODB_CONNECTION_STRING)
        mongodb_db = mongodb_client[MONGODB_DATABASE_NAME]
        
        # Test the connection
        await mongodb_client.admin.command('ping')
        logger.info(f"Successfully connected to MongoDB database: {MONGODB_DATABASE_NAME}")
        return True
    except Exception as e:
        logger.error(f"Failed to connect to MongoDB: {e}")
        logger.error("Please ensure MongoDB is running and the connection string is correct.")
        logger.error("For local development, try: mongodb://localhost:27017")
        return False

async def close_mongodb():
    """Close MongoDB connection."""
    global mongodb_client
    if mongodb_client:
        mongodb_client.close()
        logger.info("MongoDB connection closed")

def filter_metadata_for_millis(document: dict) -> dict:
    """
    Filter MongoDB document to exclude unwanted fields before sending to Millis AI.
    
    Args:
        document: The MongoDB document containing call data
        
    Returns:
        Filtered dictionary with only the fields that should be sent to Millis AI
    """
    # Fields to exclude from metadata sent to Millis AI
    excluded_fields = {
        "public_key", "call_id", "platform_agent_id", "stored_at", "agent_id"
    }
    
    filtered_metadata = {}
    
    for key, value in document.items():
        if key not in excluded_fields:
            # Convert value to string if it's not None
            if value is not None:
                filtered_metadata[key] = str(value)
            else:
                filtered_metadata[key] = ""
    
    logger.info(f"Filtered metadata - Original fields: {list(document.keys())}")
    logger.info(f"Filtered metadata - Excluded fields: {list(excluded_fields)}")
    logger.info(f"Filtered metadata - Final fields: {list(filtered_metadata.keys())}")
    
    return filtered_metadata

async def fetch_call_credentials(call_id: str) -> Dict[str, Any]:
    """
    Fetch call credentials and metadata from MongoDB.
    
    Args:
        call_id: The call ID to look up
        
    Returns:
        Dictionary containing platform_agent_id, public_key, metadata, and agent_id
    """
    if mongodb_db is None:
        logger.error("MongoDB not connected")
        return {}
    
    try:
        logger.info(f"Fetching credentials for call_id: {call_id}")
        logger.info(f"MongoDB collection: {MONGODB_COLLECTION}")
        collection = mongodb_db[MONGODB_COLLECTION]
        
        # Query for the call metadata - try multiple field names
        possible_queries = [
            {"call_id": call_id},
            {"metadata.call_id": call_id},
            {"callId": call_id},
            {"id": call_id}
        ]
        
        document = None
        for i, query in enumerate(possible_queries):
            logger.info(f"MongoDB query attempt {i+1}: {query}")
            document = await collection.find_one(query)
            if document:
                logger.info(f"Found document with query: {query}")
                break
        
        if not document:
            logger.info(f"No document found for call_id: {call_id}")
            logger.info("Available documents in collection:")
            all_docs = await collection.find({}).to_list(length=10)
            for doc in all_docs:
                if "call_id" in doc:
                    logger.info(f"  - call_id: {doc['call_id']}")
                if "metadata" in doc and "call_id" in doc["metadata"]:
                    logger.info(f"  - metadata.call_id: {doc['metadata']['call_id']}")
            
            # Use the first available document with metadata as fallback
            for doc in all_docs:
                if "metadata" in doc and doc["metadata"]:
                    logger.info(f"Using fallback document with call_id: {doc.get('call_id', 'unknown')}")
                    document = doc
                    break
        
        if not document:
            logger.warning(f"No credentials found for call_id: {call_id}")
            logger.info("MongoDB query returned: None")
            return {}
        
        # Log document with masked sensitive data
        safe_document = document.copy()
        if "public_key" in safe_document:
            safe_document["public_key"] = safe_document["public_key"][:10] + "..." if safe_document["public_key"] else "None"
        logger.info(f"MongoDB document found: {safe_document}")
        
        # Filter the document to exclude unwanted fields for Millis AI
        filtered_metadata = filter_metadata_for_millis(document)
        
        # Extract the required fields
        credentials = {
            "platform_agent_id": document.get("platform_agent_id"),
            "public_key": document.get("public_key"),
            "metadata": filtered_metadata
        }
        
        logger.info(f"Retrieved credentials for call_id {call_id}: {credentials}")
        logger.info(f"Metadata keys: {list(credentials.get('metadata', {}).keys())}")
        return credentials
        
    except Exception as e:
        logger.error(f"Error fetching credentials for call_id {call_id}: {e}")
        return {}

async def fetch_call_credentials_with_fallback(call_id: str, fallback_msg: dict) -> Dict[str, Any]:
    """
    Fetch call credentials from MongoDB with fallback to message data.
    
    Args:
        call_id: The call ID to look up
        fallback_msg: The original message to extract fallback credentials from
        
    Returns:
        Dictionary containing platform_agent_id, public_key, and agent_id
    """
    # Try MongoDB first
    logger.info(f"=== MONGODB FALLBACK FOR CALL {call_id} ===")
    credentials = await fetch_call_credentials(call_id)
    
    if credentials and credentials.get("platform_agent_id") and credentials.get("public_key"):
        logger.info(f"Using MongoDB credentials for call_id: {call_id}")
        return credentials
    
    # Fallback to message data (safety net)
    logger.warning(f"MongoDB credentials incomplete for call_id {call_id}, using message fallback")
    
    extra_params = fallback_msg.get("extraParams", {})
    custom_field = fallback_msg.get("custom_field", {})
    
    # Extract agent_id from fallback sources
    agent_id = (extra_params.get("platform_agent_id") or 
                custom_field.get("agentId") or 
                fallback_msg.get("agentId"))
    
    # Extract public_key from fallback sources
    public_key = (extra_params.get("publicKey") or extra_params.get("public_key") or 
                  custom_field.get("publicKey") or custom_field.get("public_key") or 
                  fallback_msg.get("publicKey") or fallback_msg.get("public_key"))
    
    fallback_credentials = {
        "platform_agent_id": agent_id,
        "public_key": public_key,
        "metadata": {}
    }
    
    logger.warning(f"Using fallback credentials for call_id {call_id}: {fallback_credentials}")
    return fallback_credentials

# ---------------------------------------------------------------------------#
#                          REAL-TIME AUDIO PROCESSOR                         #
# ---------------------------------------------------------------------------#
class RealTimeAudioProcessor:
    """
    Manages a single live call, bridging audio between the SAN system and Millis AI.
    """
    PHONE_RATE = 8000
    MILLIS_RATE = 16000
    CHUNK_MS = 40
    BYTES_PER_SAMPLE = 2

    MILLIS_CHUNK_SIZE = int(MILLIS_RATE * CHUNK_MS / 1000 * BYTES_PER_SAMPLE)
    PHONE_CHUNK_SIZE = int(PHONE_RATE * CHUNK_MS / 1000 * BYTES_PER_SAMPLE)

    # Throttler config
    BURST_SIZE = 10
    THROTTLE_RATE = 70
    THROTTLE_PERIOD = 1.0

    def __init__(self, agent_id: str, public_key: str, metadata: dict = None):
        self.agent_id = agent_id
        self.public_key = public_key
        self.metadata = metadata or {}
        self.ws: Optional[websockets.WebSocketClientProtocol] = None
        self.connected = False

        self.inbound = bytearray()
        self.outbound = asyncio.Queue()
        self._gen_id = 0
        self.in_lock = asyncio.Lock()
        self.out_lock = asyncio.Lock()

        self.is_paused = False
        self.stream_id: Optional[str] = None
        self.call_id: Optional[str] = None
        self.media_format: dict = {
            "encoding": "PCM", "sampleRate": self.PHONE_RATE, "channels": 1
        }
        self._packet_counter = 0
        self._ignore_outgoing = False
        self._throttler = Throttler(rate_limit=self.THROTTLE_RATE,
                                    period=self.THROTTLE_PERIOD)
        self._burst_sent = 0

    async def connect(self) -> bool:
        logger.info(f"Connecting to Millis AI for call {self.call_id}...")
        logger.info(f"Agent ID: {self.agent_id}")
        logger.info(f"Public Key: {self.public_key[:10]}...")
        logger.info(f"Millis URI: {MILLIS_WS_URI}")
        
        try:
            logger.info("Establishing WebSocket connection...")
            self.ws = await websockets.connect(MILLIS_WS_URI, open_timeout=10)
            logger.info("WebSocket connection established")
            
            initiate_payload = {
                "method": "initiate",
                "data": {
                    "agent": {
                        "agent_id": self.agent_id
                    },
                    "public_key": self.public_key,
                    "metadata": self.metadata,
                    "include_metadata_in_prompt": True
                }
            }
            logger.info(f"Millis connection metadata: {self.metadata}")
            logger.info(f"Metadata keys being sent: {list(self.metadata.keys())}")
            logger.info(f"Sample metadata values:")
            for key, value in list(self.metadata.items())[:5]:  # Show first 5 items
                logger.info(f"  {key}: {value}")
            logger.info(f"Total metadata fields: {len(self.metadata)}")
            logger.info(f"include_metadata_in_prompt: {initiate_payload['data']['include_metadata_in_prompt']}")
            logger.info(f"Sending initiate payload: {json.dumps(initiate_payload, indent=2)}")
            
            await self.ws.send(json.dumps(initiate_payload))
            logger.info("Initiate payload sent, waiting for response...")
            
            logger.info("Waiting for Millis AI response...")
            msg = await asyncio.wait_for(self.ws.recv(), timeout=10)
            logger.info(f"Received response: {msg}")
            
            try:
                parsed_msg = json.loads(msg)
                logger.info(f"Parsed response: {json.dumps(parsed_msg, indent=2)}")
                method = parsed_msg.get("method")
                logger.info(f"Response method: {method}")
                
                if method != "onready":
                    logger.error(f"Expected 'onready' method, got '{method}'")
                    raise RuntimeError(f"Millis AI did not send 'onready' confirmation. Got: {method}")
                
                self.connected = True
                logger.info("Successfully connected to Millis AI.")
                return True
                
            except json.JSONDecodeError as e:
                logger.error(f"Failed to parse response as JSON: {e}")
                logger.error(f"Raw response: {msg}")
                raise
                
        except asyncio.TimeoutError:
            logger.error("Connection timeout - Millis AI did not respond within 10 seconds")
            logger.error(f"Agent ID: {self.agent_id}")
            logger.error(f"Public Key: {self.public_key[:10]}...")
            logger.error("Please check:")
            logger.error("1. Network connectivity to Millis AI")
            logger.error("2. Agent ID and Public Key are valid")
            logger.error("3. Millis AI service is running")
            self.connected = False
            return False
        except websockets.exceptions.InvalidURI:
            logger.error(f"Invalid WebSocket URI: {MILLIS_WS_URI}")
            self.connected = False
            return False
        except websockets.exceptions.ConnectionClosed:
            logger.error("WebSocket connection was closed unexpectedly")
            self.connected = False
            return False
        except Exception as e:
            logger.error(f"Millis AI connection failed: {type(e).__name__}: {e}")
            self.connected = False
            return False

    async def disconnect(self):
        if self.ws and self.ws.state != WsState.CLOSED:
            await self.ws.close()
        self.connected = False
        self.ws = None
        logger.info("Disconnected from Millis AI.")

    @staticmethod
    def _resample(data: bytes, from_rate: int, to_rate: int) -> bytes:
        if not data: return b""
        arr = np.frombuffer(data, dtype=np.int16)
        if arr.size == 0: return b""
        new_len = int(arr.size * to_rate / from_rate)
        resampled = scipy_signal.resample(arr, new_len).astype(np.int16)
        return resampled.tobytes()

    async def _pump_inbound_to_millis(self):
        logger.info(f"Starting inbound audio pump for call {self.call_id}")
        processed_chunks = 0
        while self.connected:
            chunk8 = None
            async with self.in_lock:
                if len(self.inbound) >= self.PHONE_CHUNK_SIZE:
                    chunk8 = self.inbound[:self.PHONE_CHUNK_SIZE]
                    del self.inbound[:self.PHONE_CHUNK_SIZE]
            if not chunk8:
                await asyncio.sleep(0.005)
                continue
            try:
                chunk16 = self._resample(chunk8, self.PHONE_RATE, self.MILLIS_RATE)
                processed_chunks += 1
                # Only log every 100 chunks to reduce noise
                if processed_chunks % 100 == 0:
                    logger.info(f"Processed {processed_chunks} audio chunks to Millis AI")
                await self.ws.send(chunk16)
                self._packet_counter += 1
                if self._packet_counter >= 1_000:
                    logger.info("Sending ping to Millis AI")
                    await self.ws.send(json.dumps({"method": "ping"}))
                    self._packet_counter = 0
            except Exception as e:
                logger.error(f"Error in _pump_inbound_to_millis: {e}")
                self.connected = False

    async def _pump_millis_to_outbound(self):
        logger.info(f"Starting outbound audio pump for call {self.call_id}")
        received_chunks = 0
        while self.connected and self.ws and self.ws.state == WsState.OPEN:
            try:
                msg = await self.ws.recv()

                if not isinstance(msg, bytes):
                    try:
                        evt = json.loads(msg)
                    except json.JSONDecodeError:
                        logger.warning(f"Unparseable JSON from Millis: {msg}")
                        continue

                    method = evt.get("method")
                    data   = evt.get("data", "")

                    # log non-trivial events
                    if method not in ("ping", "pong"):
                        logger.info(f"Millis event: {method} – {data}")

                    # human barge-in
                    if method in ("ai_action", "clear"):
                        qsize = self.outbound.qsize()
                        logger.info(f"[clear] gen={self._gen_id} flushing outbound queue ({qsize} entries)")
                        self._gen_id += 1
                        self._burst_sent = 0
                        self.outbound = asyncio.Queue()
                        async with self.in_lock:
                            self.inbound.clear()
                        self.is_paused = True
                        self._ignore_outgoing = True
                        logger.info("Cleared buffers on user barge-in")
                        continue

                    # AI is about to speak again
                    if method == "start_answering":
                        self._ignore_outgoing = False
                        self.is_paused = False
                        logger.info("AI resumed – now forwarding new bytes")
                        continue

                    # pause/unpause (optional, keep if you use TTS pause)
                    if method == "pause":
                        self.is_paused = True
                        logger.info("Audio paused")
                        continue
                    if method == "unpause":
                        self.is_paused = False
                        logger.info("Audio unpaused")
                        continue

                    # anything else — skip into next recv
                    continue

                # —— 2) It’s raw audio bytes —— 
                # drop if we’re paused or ignoring
                if self.is_paused or self._ignore_outgoing:
                    logger.info(f"[out←Millis] dropping raw‐audio chunk (paused or ignore_outgoing)")
                    continue

                # buffer it
                received_chunks += 1
                if received_chunks % 100 == 0:
                    logger.info(f"Received {received_chunks} audio chunks from Millis AI")

                await self.outbound.put((self._gen_id, msg))

            except websockets.exceptions.ConnectionClosed:
                logger.warning("Millis AI closed the connection.")
                self.connected = False

            except Exception as e:
                logger.warning(f"Error reading from Millis AI: {type(e).__name__}: {e}")
                self.connected = False
    
    async def _really_send(self, client_ws: WebSocket, payload: dict):
        await client_ws.send_json(payload)

    async def _pump_outbound_to_carrier(self, client_ws: WebSocket):
        logger.info(f"Starting carrier outbound pump for call {self.call_id}")
        sent_packets = 0
        while self.connected:
            try:
                gen, chunk16 = await asyncio.wait_for(self.outbound.get(), timeout=0.5)
            except asyncio.TimeoutError:
                logger.info("[out→Carrier] get() timed out waiting for next packet")
                continue
            if gen != self._gen_id:
                logger.info(f"[out→Carrier] dropping stale gen={gen} (current={self._gen_id})")
                continue
            if self.is_paused or self._ignore_outgoing:
                logger.info("[out→Carrier] dropping pkt because paused/ignore_outgoing")
                continue

            if self.is_paused or self._ignore_outgoing:
                continue
            try:
                target_rate = self.media_format.get("sampleRate", self.PHONE_RATE)
                chunk_resampled = self._resample(chunk16, self.MILLIS_RATE, target_rate)
                payload = base64.b64encode(chunk_resampled).decode()
                sent_packets += 1
                # Only log every 200 packets to reduce noise
                if sent_packets % 200 == 0:
                    logger.info(f"Sent {sent_packets} audio packets to SAN")

                reverse_media_payload = {
                    "event": "reverse-media",
                    "callid": self.call_id,  # Changed from "callId" to "callid"
                    "payload": payload,
                    # Removed "streamId" and "mediaFormat" fields
                }

                if self._burst_sent < self.BURST_SIZE:
                    # send immediately
                    await self._really_send(client_ws, reverse_media_payload)
                    logger.info(f"[out→Carrier] sending burst pkt to SAN payload len={len(payload)}")
                    self._burst_sent += 1
                else:
                    # throttle further sends
                    async with self._throttler:
                        logger.info(f"[out→Carrier] sending burst pkt to SAN payload len={len(payload)}")
                        await self._really_send(client_ws, reverse_media_payload)

            except Exception as e:
                logger.error(f"Error in _pump_outbound_to_carrier: {e}")
                break

    async def start(self, client_ws: WebSocket) -> list[asyncio.Task]:
        logger.info(f"Starting RealTimeAudioProcessor for call {self.call_id}")
        if not await self.connect(): 
            logger.error("Failed to connect to Millis AI")
            return []
        
        logger.info("Creating audio processing tasks")
        tasks = [
            asyncio.create_task(self._pump_millis_to_outbound()),
            asyncio.create_task(self._pump_inbound_to_millis()),
            asyncio.create_task(self._pump_outbound_to_carrier(client_ws)),
        ]
        logger.info(f"Created {len(tasks)} tasks")
        return tasks

async def stop_processor(proc: Optional[RealTimeAudioProcessor], tasks: list[asyncio.Task]):
    if not proc: return
    logger.info(f"Stopping processor for call {proc.call_id}")
    for t in tasks:
        if not t.done(): t.cancel()
    await proc.disconnect()

# ---------------------------------------------------------------------------#
#                          FASTAPI /media ENDPOINT                           #
# ---------------------------------------------------------------------------#
@app.websocket("/media")
async def media_socket(ws: WebSocket):
    await ws.accept()
    logger.info("SAN system WebSocket accepted.")

    processor: Optional[RealTimeAudioProcessor] = None
    tasks: list[asyncio.Task] = []
    active_call_id: Optional[str] = None

    try:
        while True:
            raw = await ws.receive_text()
            
            msg = json.loads(raw)
            event = msg.get("event")
            
            # Enhanced logging for start events only
            if event == "start":
                logger.info("=== START EVENT ===")
                logger.info(f"Call ID: {msg.get('callId')}, Stream ID: {msg.get('streamId')}")
                logger.info(f"Agent ID: {msg.get('extraParams', {}).get('platform_agent_id')}")
                logger.info("=== END START EVENT ===")
            elif event == "media":
                payload_b64 = msg.get("payload")
                if payload_b64:
                    pcm = base64.b64decode(payload_b64)
                    # Reduce logging frequency for media chunks
                    logger.debug(f"Received media chunk: {len(pcm)} bytes for call")
                else:
                    logger.warning("Media event received but no payload found")
            else:
                logger.info(f"Received event: {event}")

            if event == "start":
                logger.info("=== START EVENT PROCESSING ===")
                new_call_id = msg.get("callId")
                stream_id = msg.get("streamId")
                logger.info(f"Start event details - callId: {new_call_id}, streamId: {stream_id}")

                if processor and new_call_id != active_call_id:
                    logger.info(f"New call detected ({active_call_id} -> {new_call_id}). Stopping old processor.")
                    await stop_processor(processor, tasks)
                    processor, tasks = None, []

                if processor is None:
                    logger.info(f"Starting processor for call: {new_call_id}")
                    
                    # Extract credentials directly from message first
                    logger.info(f"=== EXTRACTING CREDENTIALS FOR CALL {new_call_id} ===")
                    extra_params = msg.get("extraParams", {})
                    
                    # Get credentials from message
                    agent_id = extra_params.get("platform_agent_id")
                    
                    # Always use environment variable for public_key (primary source)
                    public_key = PUBLIC_KEY
                    logger.info(f"Using environment PUBLIC_KEY: {public_key[:10] + '...' if public_key else 'None'}")
                    
                    # Use environment variables as fallback for missing agent_id
                    if not agent_id:
                        agent_id = AGENT_ID
                        logger.info(f"Using environment AGENT_ID: {agent_id}")
                    
                    # If agent_id still missing, try MongoDB as final fallback
                    if not agent_id:
                        logger.warning("Agent ID missing, trying MongoDB fallback")
                        credentials = await fetch_call_credentials_with_fallback(new_call_id, msg)
                        agent_id = credentials.get("platform_agent_id") or agent_id
                        # Keep using environment public_key, don't override
                    
                    logger.info(f"Using credentials - agent_id: {agent_id}, public_key: {public_key[:10] + '...' if public_key else 'None'}")
                    
                    if not agent_id:
                        logger.error(f"No agent_id found for call_id: {new_call_id}")
                        logger.error("Closing WebSocket due to missing agent_id")
                        await ws.close(code=1008, reason="Missing agent_id")
                        return
                    
                    if not public_key:
                        logger.error("No public_key found in environment variables")
                        logger.error("Closing WebSocket due to missing public_key")
                        await ws.close(code=1008, reason="Missing public_key")
                        return
                    
                    # Try to get metadata from MongoDB using call_id from extraParams
                    logger.info("=== FETCHING METADATA ===")
                    extra_params = msg.get("extraParams", {})
                    mongodb_call_id = extra_params.get("call_id")
                    
                    if mongodb_call_id:
                        logger.info(f"Found call_id in extraParams: {mongodb_call_id}")
                        mongodb_credentials = await fetch_call_credentials(mongodb_call_id)
                        metadata = mongodb_credentials.get("metadata", {})
                        
                        if metadata:
                            logger.info(f"Using MongoDB metadata with {len(metadata)} fields")
                            logger.info(f"MongoDB metadata keys: {list(metadata.keys())}")
                            # Merge with extraParams to include platform_agent_id
                            metadata.update(extra_params)
                            logger.info(f"Metadata after merging with extraParams: {dict(list(metadata.items())[:3])}")
                        else:
                            logger.info("No MongoDB metadata found, using extraParams")
                            metadata = extra_params
                    else:
                        logger.info("No call_id in extraParams, trying SAN callId")
                        mongodb_credentials = await fetch_call_credentials(new_call_id)
                        metadata = mongodb_credentials.get("metadata", {})
                        
                        if metadata:
                            logger.info(f"Using MongoDB metadata with {len(metadata)} fields")
                            logger.info(f"MongoDB metadata keys: {list(metadata.keys())}")
                            metadata.update(extra_params)
                        else:
                            logger.info("No MongoDB metadata found, using extraParams")
                            metadata = extra_params
                    
                    logger.info("=== METADATA RESOLVED ===")
                    
                    logger.info("Agent configuration resolved:")
                    logger.info(f"  - agent_id: {agent_id}")
                    logger.info(f"  - public_key: {public_key[:10]}...")
                    logger.info(f"  - metadata: {metadata}")
                    
                    logger.info(f"=== CREATING PROCESSOR FOR CALL {new_call_id} ===")
                    processor = RealTimeAudioProcessor(agent_id, public_key, metadata)
                    processor.stream_id = stream_id
                    processor.call_id = new_call_id
                    
                    if "mediaFormat" in msg:
                        processor.media_format = msg["mediaFormat"]
                        logger.info(f"Captured media format from SAN: {processor.media_format}")
                    else:
                        logger.warning("No mediaFormat in 'start' event. Using default.")

                    logger.info(f"=== STARTING PROCESSOR ===")
                    tasks = await processor.start(ws)
                    if not tasks:
                        logger.error("Failed to start processor")
                        await ws.close(code=1011, reason="Could not connect to AI backend.")
                        return
                    active_call_id = new_call_id
                    logger.info(f"=== START EVENT COMPLETED - {len(tasks)} tasks started ===")
                continue

            elif event == "media" and processor:
                payload_b64 = msg.get("payload")
                if payload_b64:
                    pcm = base64.b64decode(payload_b64)
                    async with processor.in_lock: 
                        processor.inbound.extend(pcm)
                    # Reduce logging frequency for media processing
                    logger.debug(f"Added {len(pcm)} bytes to inbound buffer for call {active_call_id}")
                else:
                    logger.warning("Media event received but no payload found")
                continue

            elif event in ("hangup", "stop", "disconnect"):
                logger.info("=== END EVENT PROCESSING ===")
                logger.info(f"Call {active_call_id} ended via '{event}' event.")
                await stop_processor(processor, tasks)
                processor, tasks, active_call_id = None, [], None
                logger.info("=== END EVENT COMPLETED ===")
                continue

            elif event in ("connected", "answer", "ringing"):
                continue
            
            logger.warning(f"Received unhandled event: {event}")

    except WebSocketDisconnect:
        logger.info("=== WEBSOCKET DISCONNECT ===")
        logger.info("SAN system disconnected the WebSocket.")
    except Exception as e:
        logger.error("=== UNHANDLED ERROR ===")
        logger.error(f"Unhandled error in media_socket: {e}", exc_info=True)
    finally:
        logger.info("=== FINAL CLEANUP ===")
        await stop_processor(processor, tasks)
        try:
            if ws.client_state != WebSocketState.DISCONNECTED:
                await ws.close()
        except Exception as e:
            logger.warning(f"Error during WebSocket cleanup: {e}")
        logger.info("=== CLEANUP COMPLETE ===")

@app.get("/")
async def health():
    logger.info("Health check endpoint called")
    
    # Check MongoDB connection status
    mongodb_status = "connected" if mongodb_client is not None and mongodb_db is not None else "disconnected"
    
    return {
        "status": "ok", 
        "timestamp": datetime.now().isoformat(),
        "mongodb": {
            "status": mongodb_status,
            "database": MONGODB_DATABASE_NAME or "not_set",
            "collection": MONGODB_COLLECTION or "not_set"
        },
        "config": {
            "agent_id": AGENT_ID or "not_set",
            "public_key": PUBLIC_KEY[:10] + "..." if PUBLIC_KEY else "not_set",
            "millis_ws_uri": MILLIS_WS_URI
        }
    }

@app.get("/test")
async def test():
    logger.info("Test endpoint called")
    return {"message": "Server is working!", "timestamp": datetime.now().isoformat()}

@app.get("/test-mongodb")
async def test_mongodb():
    """Test MongoDB connectivity and add sample data."""
    logger.info("MongoDB test endpoint called")
    
    if mongodb_db is None:
        return {"error": "MongoDB not connected", "status": "failed"}
    
    try:
        collection = mongodb_db[MONGODB_COLLECTION]
        
        # Test insert
        test_doc = {
            "call_id": "test-call-123",
            "platform_agent_id": "test-agent-456",
            "public_key": "test-public-key-789",
            "created_at": datetime.now().isoformat()
        }
        
        result = await collection.insert_one(test_doc)
        logger.info(f"Test document inserted with ID: {result.inserted_id}")
        
        # Test query
        retrieved = await collection.find_one({"call_id": "test-call-123"})
        
        return {
            "status": "success",
            "mongodb_connected": True,
            "test_insert_id": str(result.inserted_id),
            "test_retrieved": retrieved is not None,
            "timestamp": datetime.now().isoformat()
        }
        
    except Exception as e:
        logger.error(f"MongoDB test failed: {e}")
        return {
            "error": str(e),
            "status": "failed",
            "mongodb_connected": mongodb_db is not None
        }



if __name__ == "__main__":
    print("Starting SAN to Millis AI Integration Server (v7 - MongoDB Integration)...")
    logger.info("=== SERVER STARTING ===")
    logger.info(f"Agent ID: {AGENT_ID or 'NOT_SET'}")
    logger.info(f"Public Key: {PUBLIC_KEY[:10] + '...' if PUBLIC_KEY else 'NOT_SET'}")
    logger.info(f"Millis URI: {MILLIS_WS_URI}")
    logger.info(f"MongoDB URI: {MONGODB_CONNECTION_STRING or 'NOT_SET'}")
    logger.info(f"MongoDB Database: {MONGODB_DATABASE_NAME or 'NOT_SET'}")
    logger.info(f"MongoDB Collection: {MONGODB_COLLECTION or 'NOT_SET'}")
    logger.info("=== SERVER READY ===")
    uvicorn.run(app, host="0.0.0.0", port=8000)