#!/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`.
#  -------------------------------------------------------------------

import os
import json
import base64
import asyncio
import logging
from datetime import datetime

# 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

# 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")

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

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

# ---------------------------------------------------------------------------#
#                          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 = 20
    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)

    def __init__(self, agent_id: str, public_key: str):
        self.agent_id = agent_id
        self.public_key = public_key
        self.ws: websockets.WebSocketClientProtocol | None = None
        self.connected = False

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

        self.is_paused = False
        self.stream_id: str | None = None
        self.call_id: str | None = None
        self.media_format: dict = {
            "encoding": "PCM", "sampleRate": self.PHONE_RATE, "channels": 1
        }
        self._packet_counter = 0

    async def connect(self) -> bool:
        logger.info("🤖 Connecting to Millis AI...")
        try:
            self.ws = await websockets.connect(MILLIS_WS_URI, open_timeout=10)
            await self.ws.send(
                json.dumps({
                    "method": "initiate",
                    "data": {"agent": {"agent_id": self.agent_id}, "public_key": self.public_key},
                })
            )
            msg = await asyncio.wait_for(self.ws.recv(), timeout=10)
            if json.loads(msg).get("method") != "onready":
                raise RuntimeError("Millis AI did not send 'onready' confirmation.")
            self.connected = True
            logger.info("✅ Successfully connected to Millis AI.")
            return True
        except Exception as e:
            logger.error(f"❌ Millis AI connection failed: {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):
        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)
                await self.ws.send(chunk16)
                self._packet_counter += 1
                if self._packet_counter >= 1_000:
                    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):
        while self.connected and self.ws and self.ws.state == WsState.OPEN:
            try:
                msg = await self.ws.recv()
                if isinstance(msg, bytes):
                    async with self.out_lock: self.outbound.extend(msg)
                    continue
                evt = json.loads(msg)
                method = evt.get("method")
                logger.info(f"🤖 JSON from Millis: {evt}")
                if method == "pause": self.is_paused = True
                elif method == "unpause": self.is_paused = False
                elif method in ("clear", "start_answering"):
                    async with self.out_lock: self.outbound.clear()
                    self.is_paused = False
            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: {e}")
                self.connected = False

    async def _pump_outbound_to_carrier(self, client_ws: WebSocket):
        sent_packets = 0
        while self.connected:
            if self.is_paused:
                await asyncio.sleep(0.01)
                continue
            chunk16 = None
            async with self.out_lock:
                if len(self.outbound) >= self.MILLIS_CHUNK_SIZE:
                    chunk16 = self.outbound[:self.MILLIS_CHUNK_SIZE]
                    del self.outbound[:self.MILLIS_CHUNK_SIZE]
            if not chunk16:
                await asyncio.sleep(0.005)
                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
                if sent_packets % 100 == 1:
                    logger.info(f"⬆️  Sending upstream audio packet #{sent_packets} to SAN...")

                ### --- FIX: Modified the JSON payload to match the provider's simple format --- ###
                await client_ws.send_json({
                    "event": "reverse-media",
                    "callid": self.call_id,  # Changed from "callId" to "callid"
                    "payload": payload,
                    # Removed "streamId" and "mediaFormat" fields
                })
                ### --- END FIX --- ###

            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]:
        if not await self.connect(): return []
        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)),
        ]
        return tasks

async def stop_processor(proc: RealTimeAudioProcessor | None, tasks: list[asyncio.Task]):
    if not proc: return
    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: RealTimeAudioProcessor | None = None
    tasks: list[asyncio.Task] = []
    active_call_id: str | None = None

    try:
        while True:
            raw = await ws.receive_text()
            msg = json.loads(raw)
            event = msg.get("event")

            if event == "start":
                new_call_id = msg.get("callId")
                stream_id = msg.get("streamId")

                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}")
                    processor = RealTimeAudioProcessor(AGENT_ID, PUBLIC_KEY)
                    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.")

                    tasks = await processor.start(ws)
                    if not tasks:
                        await ws.close(code=1011, reason="Could not connect to AI backend.")
                        return
                    active_call_id = new_call_id
                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)
                continue

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

            elif event in ("connected", "answer", "ringing"):
                logger.debug(f"ℹ️  Informational event received: {event}")
                continue
            
            logger.warning(f"⚠️ Received unhandled event: {event}")

    except WebSocketDisconnect:
        logger.info("🚪 SAN system disconnected the WebSocket.")
    except Exception as e:
        logger.error(f"❌ Unhandled error in media_socket: {e}", exc_info=True)
    finally:
        await stop_processor(processor, tasks)
        if ws.client_state != WebSocketState.DISCONNECTED:
            await ws.close()
        logger.info("✅ Cleanup complete for this WebSocket connection.")

@app.get("/")
async def health():
    return {"status": "ok", "timestamp": datetime.now().isoformat()}

if __name__ == "__main__":
    print("🚀 Starting SAN to Millis AI Integration Server (v5 - Provider Format Fix)...")
    uvicorn.run(app, host="0.0.0.0", port=7860)