"""
Twilio Media Streams connector for testing voice agents via phone calls.

Primary API:
    TwilioCall - Make phone calls via Twilio (recommended)

Internal classes:
    TwilioServer - WebSocket server for Twilio Media Streams
    TwilioConnector - Low-level Twilio connector

Usage:
    from rehearse import TwilioCall
    from rehearse.audio.tts import ElevenLabsTTS
    from rehearse.audio.stt import ElevenLabsSTT

    tts = ElevenLabsTTS(api_key="...")
    stt = ElevenLabsSTT(api_key="...")

    async with TwilioCall(
        to="+19251234567",
        account_sid="ACxxxxx",
        auth_token="xxxxx",
        from_number="+15551234567",
        ngrok_url="abc123.ngrok-free.app",
        tts=tts,
        stt=stt,
    ) as call:
        response = await call.listen()
"""

import asyncio
import base64
import json
import logging
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from ..audio.tts.base import BaseTTS
    from ..audio.stt.base import BaseSTT
    from ..conversation import Conversation

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from twilio.rest import Client
import uvicorn

from .base import BaseConnector, ReceivedAudio
from ..audio.formats import mulaw_8k_to_pcm16_24k, pcm16_24k_to_mulaw_8k

logger = logging.getLogger(__name__)


class TwilioMediaHandler:
    """Handles a single Twilio Media Streams WebSocket connection."""

    def __init__(
        self,
        incoming_queue: asyncio.Queue,
        outgoing_queue: asyncio.Queue,
        connected_event: asyncio.Event,
        shutdown_event: Optional[asyncio.Event] = None,
    ):
        self.incoming_queue = incoming_queue
        self.outgoing_queue = outgoing_queue
        self.connected_event = connected_event
        self.shutdown_event = shutdown_event or asyncio.Event()
        self.stream_sid: Optional[str] = None

    async def handle_connection(self, websocket: WebSocket):
        """Handle the WebSocket connection from Twilio."""
        logger.debug("Twilio WebSocket connection accepted")
        await websocket.accept()

        receive_task = asyncio.create_task(self._receive_from_twilio(websocket))
        send_task = asyncio.create_task(self._send_to_twilio(websocket))

        try:
            done, pending = await asyncio.wait(
                [receive_task, send_task],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
        except (WebSocketDisconnect, asyncio.CancelledError):
            pass
        finally:
            receive_task.cancel()
            send_task.cancel()
            await self.incoming_queue.put(None)
            try:
                await websocket.close()
            except Exception:
                pass

    async def _receive_from_twilio(self, websocket: WebSocket):
        """Receive audio from Twilio and put into incoming queue."""
        try:
            while not self.shutdown_event.is_set():
                try:
                    data = await asyncio.wait_for(
                        websocket.receive_text(), timeout=0.5
                    )
                except asyncio.TimeoutError:
                    continue

                event = json.loads(data)
                event_type = event.get("event")

                if event_type == "start":
                    self.stream_sid = event["start"]["streamSid"]
                    logger.info("Twilio stream started (streamSid=%s)", self.stream_sid)
                    self.connected_event.set()

                elif event_type == "media":
                    payload = event["media"]["payload"]
                    audio = base64.b64decode(payload)
                    await self.incoming_queue.put(audio)

                elif event_type == "stop":
                    logger.info("Twilio stream stopped")
                    break

        except WebSocketDisconnect:
            logger.debug("Twilio WebSocket disconnected")

    async def _send_to_twilio(self, websocket: WebSocket):
        """Send audio from outgoing queue to Twilio."""
        try:
            while not self.shutdown_event.is_set():
                try:
                    audio = await asyncio.wait_for(
                        self.outgoing_queue.get(), timeout=0.5
                    )
                except asyncio.TimeoutError:
                    continue

                if audio is None:
                    continue

                if self.stream_sid:
                    payload = base64.b64encode(audio).decode()
                    message = {
                        "event": "media",
                        "streamSid": self.stream_sid,
                        "media": {"payload": payload},
                    }
                    await websocket.send_text(json.dumps(message))

        except asyncio.CancelledError:
            pass


class TwilioServer:
    """
    WebSocket server for Twilio Media Streams.

    Usage:
        server = TwilioServer(port=8765)
        await server.start()

        # Get queues for connector
        queues = server.get_queues()

        # Stop when done
        await server.stop()
    """

    def __init__(self, host: str = "0.0.0.0", port: int = 8765):
        self.host = host
        self.port = port
        self.app = FastAPI()
        self.server = None
        self.server_task = None

        self.incoming_queue: asyncio.Queue = asyncio.Queue()
        self.outgoing_queue: asyncio.Queue = asyncio.Queue()
        self.connected_event = asyncio.Event()
        self.shutdown_event = asyncio.Event()

        self._setup_routes()

    def _setup_routes(self):
        @self.app.websocket("/media-stream")
        async def media_stream(websocket: WebSocket):
            handler = TwilioMediaHandler(
                self.incoming_queue,
                self.outgoing_queue,
                self.connected_event,
                self.shutdown_event,
            )
            await handler.handle_connection(websocket)

        @self.app.get("/health")
        async def health():
            return {"status": "ok"}

    def get_queues(self) -> dict:
        """Get queues for use with TwilioConnector."""
        return {
            "incoming": self.incoming_queue,
            "outgoing": self.outgoing_queue,
            "connected": self.connected_event,
        }

    async def start(self):
        """Start the server in the background."""
        logger.info("Starting WebSocket server on %s:%d", self.host, self.port)
        config = uvicorn.Config(
            self.app,
            host=self.host,
            port=self.port,
            log_level="critical",
        )
        self.server = uvicorn.Server(config)
        self.server_task = asyncio.create_task(self.server.serve())
        await asyncio.sleep(0.5)
        logger.debug("WebSocket server started")

    async def stop(self):
        """Stop the server."""
        logger.debug("Stopping WebSocket server")
        self.shutdown_event.set()
        await asyncio.sleep(0.1)

        if self.server:
            self.server.should_exit = True
            await asyncio.sleep(0.2)

        if self.server_task:
            self.server_task.cancel()
            try:
                await self.server_task
            except asyncio.CancelledError:
                pass
        logger.info("WebSocket server stopped")

    def reset(self):
        """Reset queues for a new test."""
        self.incoming_queue = asyncio.Queue()
        self.outgoing_queue = asyncio.Queue()
        self.connected_event = asyncio.Event()
        self.shutdown_event = asyncio.Event()


class TwilioConnector(BaseConnector):
    """
    Internal connector for Twilio phone calls.

    Note: For most use cases, use TwilioCall instead which provides
    a simpler, all-in-one interface.
    """

    def __init__(
        self,
        account_sid: str,
        auth_token: str,
        from_number: str,
        ngrok_url: str,
        send_digits: Optional[str] = None,
    ):
        """
        Create a TwilioConnector with explicit configuration.

        Args:
            account_sid: Twilio Account SID (starts with "AC")
            auth_token: Twilio Auth Token
            from_number: Your Twilio phone number to call FROM (e.g., "+15551234567")
            ngrok_url: Your ngrok URL domain (e.g., "abc123.ngrok-free.app")
            send_digits: DTMF digits to send after call connects (e.g., "www7" for pause then press 7)
        """
        self.account_sid = account_sid
        self.auth_token = auth_token
        self.from_number = from_number
        self.websocket_url = f"wss://{ngrok_url}/media-stream"
        self.send_digits = send_digits

        self.client = Client(account_sid, auth_token)

        # Set during call()
        self._phone_number: Optional[str] = None
        self._call_send_digits: Optional[str] = None  # Per-call override

        self.call_sid: Optional[str] = None
        self._connected = False

        # Queues from server
        self._server_queues: Optional[dict] = None
        self._incoming_queue: Optional[asyncio.Queue] = None
        self._outgoing_queue: Optional[asyncio.Queue] = None
        self._connected_event: Optional[asyncio.Event] = None

        self._tool_calls: list = []
        self._transcripts: list = []

    def set_call_params(self, phone_number: str, send_digits: Optional[str] = None) -> None:
        """
        Set the phone number and optional DTMF digits for the call.

        Args:
            phone_number: The agent's phone number to call
            send_digits: Per-call DTMF override (if None, uses connector's send_digits)
        """
        self._phone_number = phone_number
        self._call_send_digits = send_digits  # Per-call override

    def set_server_queues(self, queues: dict) -> None:
        """Set the server queues for audio exchange."""
        self._server_queues = queues

    async def connect(self) -> None:
        """Initiate call to the voice agent."""
        if self._server_queues is None:
            raise RuntimeError("Server queues not set. Use set_server_queues() first.")
        if self._phone_number is None:
            raise RuntimeError("Phone number not set. Use set_call_params() first.")

        self._incoming_queue = self._server_queues["incoming"]
        self._outgoing_queue = self._server_queues["outgoing"]
        self._connected_event = self._server_queues["connected"]

        logger.info("Initiating call from %s to %s", self.from_number, self._phone_number)
        logger.debug("WebSocket URL: %s", self.websocket_url)

        # TwiML for bidirectional audio streaming
        twiml = f"""
        <Response>
            <Connect>
                <Stream url="{self.websocket_url}" />
            </Connect>
            <Pause length="60"/>
        </Response>
        """

        # Build call parameters
        call_params = {
            "to": self._phone_number,
            "from_": self.from_number,
            "twiml": twiml,
        }
        # Use per-call override if set, otherwise use connector default
        digits = self._call_send_digits if self._call_send_digits is not None else self.send_digits
        if digits:
            call_params["send_digits"] = digits
            logger.debug("Will send DTMF digits: %s", digits)

        call = self.client.calls.create(**call_params)
        self.call_sid = call.sid
        logger.debug("Call initiated (callSid=%s)", self.call_sid)

        # Wait for Twilio to connect
        try:
            logger.debug("Waiting for Twilio stream to connect...")
            await asyncio.wait_for(self._connected_event.wait(), timeout=30.0)
            self._connected = True
            logger.info("Call connected (callSid=%s)", self.call_sid)
        except asyncio.TimeoutError:
            logger.error("Twilio did not connect within 30 seconds")
            raise RuntimeError("Twilio did not connect within 30 seconds")

    async def send_audio(self, audio: bytes) -> None:
        """Send PCM16 24kHz audio to the agent."""
        if not self._connected:
            raise RuntimeError("Not connected")

        # Convert to Twilio format (mulaw 8kHz)
        mulaw_audio = pcm16_24k_to_mulaw_8k(audio)

        # Send in 20ms chunks (160 bytes at 8kHz)
        chunk_size = 160
        for i in range(0, len(mulaw_audio), chunk_size):
            chunk = mulaw_audio[i : i + chunk_size]
            await self._outgoing_queue.put(chunk)

        await self._outgoing_queue.put(None)

    async def receive_audio(
        self,
        timeout: float = 15.0,
        silence_threshold: float = 1.5,
        max_duration: float = 8.0,
    ) -> ReceivedAudio:
        """Receive audio from the agent until silence is detected."""
        if not self._connected:
            raise RuntimeError("Not connected")

        audio_chunks = []
        start_time = asyncio.get_event_loop().time()
        last_speech_time = None
        speech_started = False
        packet_timeout = 0.5

        while True:
            try:
                chunk = await asyncio.wait_for(
                    self._incoming_queue.get(), timeout=packet_timeout
                )

                if chunk is None:
                    if not audio_chunks and not speech_started:
                        continue
                    break

                audio_chunks.append(chunk)
                current_time = asyncio.get_event_loop().time()

                is_speech = self._is_speech(chunk)
                if is_speech:
                    speech_started = True
                    last_speech_time = current_time

                if speech_started and last_speech_time:
                    silence_duration = current_time - last_speech_time
                    if silence_duration >= silence_threshold:
                        break

                if speech_started:
                    speech_duration = current_time - start_time
                    if speech_duration >= max_duration:
                        break

            except asyncio.TimeoutError:
                if audio_chunks:
                    break
                if asyncio.get_event_loop().time() - start_time > timeout:
                    break

        # Convert to PCM
        mulaw_audio = b"".join(audio_chunks)
        pcm_audio = mulaw_8k_to_pcm16_24k(mulaw_audio)

        return ReceivedAudio(audio=pcm_audio, transcript="", tool_calls=[])

    def _is_speech(self, mulaw_chunk: bytes, threshold: int = 5) -> bool:
        """Detect if a mulaw audio chunk contains speech."""
        if not mulaw_chunk:
            return False

        total_energy = 0
        for byte in mulaw_chunk:
            deviation = abs(byte - 255)
            if deviation > 128:
                deviation = 256 - deviation
            total_energy += deviation

        avg_energy = total_energy / len(mulaw_chunk)
        return avg_energy > threshold

    async def send_dtmf(self, digits: str) -> None:
        """Send DTMF tones via Twilio API."""
        if not self.call_sid:
            raise RuntimeError("Cannot send DTMF: not connected")

        self._connected_event.clear()

        # Drain queue
        while not self._incoming_queue.empty():
            try:
                self._incoming_queue.get_nowait()
            except asyncio.QueueEmpty:
                break

        twiml = f"""
        <Response>
            <Play digits="{digits}"/>
            <Connect>
                <Stream url="{self.config.websocket_url}" />
            </Connect>
            <Pause length="60"/>
        </Response>
        """

        self.client.calls(self.call_sid).update(twiml=twiml)

        try:
            await asyncio.wait_for(self._connected_event.wait(), timeout=10.0)
        except asyncio.TimeoutError:
            pass

    async def close(self) -> None:
        """Hang up the call."""
        logger.debug("Closing call (callSid=%s)", self.call_sid)
        if self.call_sid:
            try:
                self.client.calls(self.call_sid).update(status="completed")
                logger.info("Call ended (callSid=%s)", self.call_sid)
            except Exception as e:
                logger.warning("Failed to end call: %s", e)
        self._connected = False

    @property
    def is_connected(self) -> bool:
        return self._connected

    @property
    def tool_calls(self) -> list:
        return self._tool_calls

    @property
    def transcripts(self) -> list:
        return self._transcripts


class TwilioCall:
    """
    Make a phone call via Twilio for testing voice agents.

    This is the primary interface for testing voice agents over phone calls.
    It handles all the complexity of Twilio Media Streams, WebSocket servers,
    and audio format conversion.

    Usage:
        from rehearse import TwilioCall
        from rehearse.audio.tts import ElevenLabsTTS
        from rehearse.audio.stt import ElevenLabsSTT

        tts = ElevenLabsTTS(api_key="...")
        stt = ElevenLabsSTT(api_key="...")

        async with TwilioCall(
            to="+19251234567",
            account_sid="ACxxxxx",
            auth_token="xxxxx",
            from_number="+15551234567",
            ngrok_url="abc123.ngrok-free.app",
            tts=tts,
            stt=stt,
        ) as call:
            greeting = await call.listen()
            await call.say("Hello!")
    """

    def __init__(
        self,
        to: str,
        account_sid: str,
        auth_token: str,
        from_number: str,
        ngrok_url: str,
        tts: "BaseTTS",
        stt: "BaseSTT",
        send_digits: Optional[str] = None,
        audio_path: Optional[str] = None,
    ):
        """
        Create a TwilioCall.

        Args:
            to: Phone number to call (the voice agent's number)
            account_sid: Twilio Account SID (starts with "AC")
            auth_token: Twilio Auth Token
            from_number: Your Twilio phone number to call FROM
            ngrok_url: Your ngrok URL domain (e.g., "abc123.ngrok-free.app")
            tts: Text-to-speech provider (e.g., ElevenLabsTTS)
            stt: Speech-to-text provider (e.g., ElevenLabsSTT)
            send_digits: DTMF digits to send after connect (e.g., "www7")
            audio_path: Optional complete file path (e.g., "/tmp/debug.wav") to save WAV recording for debugging
        """
        self.to = to
        self.tts = tts
        self.stt = stt
        self.send_digits = send_digits
        self.audio_path = audio_path

        # Create internal connector
        self._connector = TwilioConnector(
            account_sid=account_sid,
            auth_token=auth_token,
            from_number=from_number,
            ngrok_url=ngrok_url,
            send_digits=send_digits,
        )

        self._server: Optional[TwilioServer] = None
        self._conversation: Optional["Conversation"] = None

    async def __aenter__(self) -> "Conversation":
        """Start the call and return the conversation interface."""
        # Import here to avoid circular imports
        from ..conversation import Conversation

        logger.info("Starting call to %s", self.to)

        # Start WebSocket server
        self._server = TwilioServer()
        await self._server.start()
        logger.debug("Started local WebSocket server")

        # Configure and connect
        self._connector.set_call_params(self.to, self.send_digits)
        self._connector.set_server_queues(self._server.get_queues())
        await self._connector.connect()
        logger.info("Call connected")

        self._conversation = Conversation(self._connector, self.tts, self.stt)
        return self._conversation

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """End the call and clean up."""
        import os

        logger.info("Ending call")

        # Save recorded audio if audio_path is specified
        if self.audio_path and self._conversation:
            from ..audio.formats import save_audio

            audio_data = self._conversation.get_all_audio()
            if audio_data:
                # Create parent directories if they don't exist
                audio_dir = os.path.dirname(self.audio_path)
                if audio_dir:
                    os.makedirs(audio_dir, exist_ok=True)
                # Save as stereo: agent on left channel, user on right channel
                save_audio(audio_data, self.audio_path, channels=2)
                logger.info("Saved stereo audio recording to %s", self.audio_path)

        # Close connector
        await self._connector.close()

        # Stop server
        if self._server:
            await self._server.stop()
            logger.debug("Stopped local WebSocket server")
