call_session.py

"""
src/ari/call_session.py

Manage resources and media loops for a single PSTN call.
A session owns RTP sockets, an ARI EM channel, and an ECS WebSocket connection.
It keeps a steady 20 ms RTP cadence regardless of websocket burstiness.
"""

import asyncio
import base64
import json
import random
import socket
import time
from uuid import uuid4

import aiohttp
import websockets

from env_config import EnvConfig
from utils.telephony_utils import rtp_build, rtp_strip_header
from utils.logger import get_logger

logger = get_logger(name=__name__)


class CallSession:
    """
    Manage resources and media loops for a single PSTN call.
    A session owns RTP sockets, an ARI EM channel, and an ECS WebSocket connection.
    It keeps a steady 20 ms RTP cadence regardless of websocket burstiness.
    """

    def __init__(self, pstn_channel_id: str, http: aiohttp.ClientSession) -> None:
        """
        Initialize session state.
        Args:
            pstn_channel_id: PSTN channel ID from Asterisk
            http: aiohttp.ClientSession for ARI HTTP calls
        """
        self.pstn_channel_id = pstn_channel_id
        self.http = http
        self.bridge_id: str | None = None
        self.em_channel_id: str | None = None

        # RTP sockets/addresses
        self.rtp_recv_sock: socket.socket | None = None  # Asterisk -> us
        self.rtp_send_addr: tuple[str, int] | None = None  # us -> Asterisk (EM local addr)
        self._rtp_send_sock: socket.socket | None = None

        # RTP state
        self.seq = random.randint(0, 65535)
        self.ts = random.randint(0, 2**32 - 1)
        self.ssrc = random.getrandbits(32)  # arbitrary stream id per call

        # ECS WebSocket
        self.ecs_ws: websockets.WebSocketClientProtocol | None = None
        self.stream_sid = str(uuid4())  # unique stream identifier

        # Tasks
        self.tasks: dict[str, asyncio.Task] = {}
        self.closed = asyncio.Event()

        # Outbound audio pacing/buffering
        self._out_buf: bytearray = bytearray()
        self._out_buf_lock: asyncio.Lock = asyncio.Lock()
        self._mark_next_real_frame: bool = False

        # Timestamp management for media events
        self._timestamp_counter: int = 0

        # Track pending marks for ACK back to ECS
        self._pending_marks: list[str] = []

        logger.info(
            "[call %s] init; seq=%d ts=%d ssrc=0x%08x stream_sid=%s",
            self.pstn_channel_id, self.seq, self.ts, self.ssrc, self.stream_sid
        )


    # ARI helpers _________________________________________________________________________________

    async def _create_bridge(self) -> None:
        """
        Create a mixing bridge for the PSTN and EM channels.
        """
        async with self.http.post(
            url=f"{EnvConfig.ARI_BASE}/bridges",
            params={"type": "mixing"}
        ) as r:
            r.raise_for_status()
            self.bridge_id = (await r.json())["id"]
        logger.info("[call %s] created bridge %s", self.pstn_channel_id, self.bridge_id)

    async def _add_to_bridge(self, channel_id: str) -> None:
        """
        Add a channel to the active bridge, retrying on transient errors.
        Args:
            channel_id: str representing the channel ID to add to the bridge.
        """
        assert self.bridge_id
        for _ in range(0, 9):
            try:
                async with self.http.post(
                    url=f"{EnvConfig.ARI_BASE}/bridges/{self.bridge_id}/addChannel",
                    params={"channel": channel_id},
                ) as r:
                    if r.status in (409, 422):  # HTTP 409=Conflict, 422=Unprocessable Entity
                        await asyncio.sleep(0.2)
                        continue
                    r.raise_for_status()
                    logger.info(
                        "[call %s] added channel %s to bridge %s",
                        self.pstn_channel_id, channel_id, self.bridge_id
                    )
                    return
            except Exception:
                await asyncio.sleep(0.2)
        raise RuntimeError(f"addChannel failed repeatedly for {channel_id}")

    async def _alloc_rtp_socket(self) -> int:
        """
        Create a UDP socket to receive RTP packets from Asterisk; return bound port.
        """
        # use the same port range that AWS Chime expects (10000-10299)
        start_port = EnvConfig.RTP_PORT_START
        end_port = EnvConfig.RTP_PORT_END

        for port in range(start_port, end_port + 1):
            sock = None
            try:
                # attempt to create a socket and bind it to the given port
                sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 1 << 20)
                sock.bind(address=("0.0.0.0", port))  # connect on local
                sock.setblocking(True)
                self.rtp_recv_sock = sock
                logger.info("[call %s] allocated RTP recv socket on port %d", self.pstn_channel_id, port)
                return port
            except OSError as e:
                # if current port is in use, try next one
                logger.error("[call %s] error allocating RTP recv socket on port %d: %s", self.pstn_channel_id, port, e)
                if sock:
                    try:
                        sock.close()
                    except Exception:
                        logger.error("[call %s] error closing RTP recv socket: %s", self.pstn_channel_id, e)
                        pass
                continue

        # if no available ports are found, raise an error
        raise RuntimeError(f"No available ports in range {start_port}-{end_port} for RTP socket")

    async def _create_external_media(self, host_port: str) -> None:
        """
        Start an ARI External Media channel pointing at host:port.
        Args:
            host_port: str representing the host and port to connect to.
        """
        params = {
            "app": EnvConfig.ARI_APP,
            "external_host": host_port,
            "format": "ulaw",
            "transport": "udp",
            "encapsulation": "rtp",
            "direction": "both",
            "connection_type": "client",
        }
        async with self.http.post(
            url=f"{EnvConfig.ARI_BASE}/channels/externalMedia",
            params=params
        ) as r:
            r.raise_for_status()
            self.em_channel_id = (await r.json())["id"]
        logger.info(
            "[call %s] created ExternalMedia channel %s -> %s",
            self.pstn_channel_id, self.em_channel_id, host_port
        )

    async def _wait_channel_up(self, channel_id: str, timeout: float = 3.0) -> bool:
        """
        Poll ARI until the channel state is Up or timeout expires.
        Args:
            channel_id: str representing the channel ID to wait for.
            timeout: float representing the timeout in seconds.
        Returns:
            bool: True if the channel is up, False otherwise.
        """
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            try:
                async with self.http.get(
                    url=f"{EnvConfig.ARI_BASE}/channels/{channel_id}"
                ) as r:
                    if r.status == 200 and (await r.json()).get("state") == "Up":
                        return True
            except Exception:
                pass
            await asyncio.sleep(0.1)
        return False

    async def _get_channel_var(self, channel_id: str, var: str) -> str | None:
        """
        Get a channel variable (e.g. EM return IP/port) from ARI.
        Args:
            channel_id: str representing the channel ID to get the variable from.
            var: str representing the variable to get.
        Returns:
            str | None: The value of the variable, or None if the variable is not found.
        """
        async with self.http.get(
            url=f"{EnvConfig.ARI_BASE}/channels/{channel_id}/variable",
            params={"variable": var},
        ) as r:
            r.raise_for_status()
            return (await r.json()).get("value")


    # ECS WebSocket _______________________________________________________________________________

    async def _open_ecs_websocket(self) -> None:
        """
        Connect to ECS WebSocket and send start event.
        """
        try:
            self.ecs_ws = await websockets.connect(uri=EnvConfig.ECS_MEDIA_WSS_URL)
            logger.info("[call %s] ECS WS connected: %s", self.pstn_channel_id, EnvConfig.ECS_MEDIA_WSS_URL)

            # send start event
            start_event = {
                "event": "start",
                "start": {
                    "streamSid": self.stream_sid,
                    "callSid": self.pstn_channel_id,
                    "customParameters": {
                        # add any custom parameters ECS might need
                        "source": "asterisk-shim",
                        "format": "ulaw",
                    }
                }
            }
            await self.ecs_ws.send(json.dumps(start_event))
            logger.info("[call %s] sent start event to ECS: %s", self.pstn_channel_id, self.stream_sid)

        except Exception as e:
            logger.error("[call %s] failed to connect to ECS: %s", self.pstn_channel_id, e)
            raise


    # RTP <-> ECS loops ___________________________________________________________________________

    def _rtp_send(self, payload: bytes, marker: int = 0) -> None:
        """
        Send one 20 ms μ-law frame over RTP and advance seq/ts.
        Args:
            payload: bytes containing the μ-law audio payload.
            marker: int representing the marker of the packet. Set to 1 to indicate a talkspurt start 
            (optional; used here to mark the first "real" frame after periods of silence/fillers).
        """
        if self._rtp_send_sock is None:
            self._rtp_send_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try:
                self._rtp_send_sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1 << 20)
            except OSError:
                pass

        pkt = rtp_build(payload, self.seq, self.ts, self.ssrc, marker=marker)
        self.seq = (self.seq + 1) & 0xFFFF  # 16-bit wrap
        self.ts = (self.ts + EnvConfig.FRAME_SAMPLES) & 0xFFFFFFFF  # 32-bit wrap
        assert self.rtp_send_addr is not None
        self._rtp_send_sock.sendto(pkt, self.rtp_send_addr)

    async def _rtp_to_ecs(self) -> None:
        """
        Receive μ-law RTP from Asterisk and stream to ECS as media events.
        """
        assert self.rtp_recv_sock is not None
        assert self.ecs_ws is not None
        loop = asyncio.get_running_loop()
        logger.info("[call %s] RTP→ECS loop start", self.pstn_channel_id)

        while not self.closed.is_set():
            try:
                datagram, _ = await loop.run_in_executor(None, self.rtp_recv_sock.recvfrom, 2048)
            except Exception as e:
                if self.closed.is_set():
                    break
                logger.debug("[call %s] recv exception: %s", self.pstn_channel_id, e)
                continue
            
            payload = rtp_strip_header(datagram)
            if not payload:
                continue

            # Increment timestamp counter based on actual payload length
            frames = max(1, len(payload) // EnvConfig.FRAME_BYTES)
            self._timestamp_counter += frames * 20  # ms

            # Send media event to ECS
            try:
                media_event = {
                    "event": "media",
                    "streamSid": self.stream_sid,
                    "media": {
                        "payload": base64.b64encode(payload).decode("ascii"),
                        "timestamp": self._timestamp_counter
                    }
                }
                await self.ecs_ws.send(json.dumps(media_event))
            except Exception as e:
                logger.info("[call %s] ECS send err: %s", self.pstn_channel_id, e)
                asyncio.create_task(  # don't await (avoids deadlock)
                    coro=self.stop()
                )
                return

    async def _ecs_to_rtp(self) -> None:
        """
        Receive media events from ECS and enqueue μ-law audio for RTP sending.
        """
        assert self.ecs_ws is not None
        logger.info("[call %s] ECS→RTP loop start", self.pstn_channel_id)

        # start RTP pacer
        pacer = asyncio.create_task(
            coro=self._rtp_pacer_loop(),
            name=f"rtpout-pacer-{self.pstn_channel_id}"
        )
        self.tasks["rtp_pacer"] = pacer

        try:
            async for raw in self.ecs_ws:
                if isinstance(raw, str):
                    msg = json.loads(raw)
                    event = msg.get("event")

                    if event == "media" and "media" in msg:
                        # decode and enqueue audio
                        try:
                            payload = msg["media"]["payload"]
                            decoded = base64.b64decode(payload)
                        except Exception:
                            continue
    
                        async with self._out_buf_lock:
                            self._out_buf.extend(decoded)

                        # mark next frame for talkspurt
                        if not self._mark_next_real_frame:
                            self._mark_next_real_frame = True

                    elif event == "clear":
                        # clear output buffer immediately
                        async with self._out_buf_lock:
                            buffer_size_before = len(self._out_buf)
                            self._out_buf.clear()
                        logger.info(
                            "[call %s] cleared %d bytes from output buffer",
                            self.pstn_channel_id, buffer_size_before
                        )

                    elif event == "mark":
                        # queue mark ACK; we'll send it when we actually transmit a real RTP frame
                        mark_name = (msg.get("mark") or {}).get("name", "responsePart")
                        self._pending_marks.append(mark_name)
                        logger.debug("[call %s] queued mark ack: %s", self.pstn_channel_id, mark_name)

                    else:
                        logger.debug("[call %s] unhandled ECS event: %s", self.pstn_channel_id, event)

        except websockets.ConnectionClosed:
            logger.info("[call %s] ECS WS closed", self.pstn_channel_id)
            asyncio.create_task(
                coro=self.stop()
            )
            return
        except Exception as e:
            logger.error("[call %s] _ecs_to_rtp error: %s", self.pstn_channel_id, e, exc_info=True)
            asyncio.create_task(
                coro=self.stop()
            )
            return

    async def _rtp_pacer_loop(self) -> None:
        """
        Tick every 20 ms and send exactly one RTP frame (silence or audio).
        """
        logger.info("[call %s] RTP pacer start", self.pstn_channel_id)
        tick_seconds = 0.02
        next_time = time.monotonic()
        try:
            while not self.closed.is_set():
                next_time += tick_seconds
                delay = max(0.0, next_time - time.monotonic())
                if delay > 0:
                    await asyncio.sleep(delay)

                have_real_frame = False
                async with self._out_buf_lock:
                    if len(self._out_buf) >= EnvConfig.FRAME_BYTES:
                        frame = bytes(self._out_buf[:EnvConfig.FRAME_BYTES])
                        del self._out_buf[:EnvConfig.FRAME_BYTES]
                        have_real_frame = True
                    else:
                        if self._out_buf:
                            take = min(len(self._out_buf), EnvConfig.FRAME_BYTES)
                            head = bytes(self._out_buf[:take])
                            del self._out_buf[:take]
                            pad = bytes([0xFF]) * (EnvConfig.FRAME_BYTES - take)
                            frame = head + pad
                            have_real_frame = True
                        else:
                            frame = bytes([0xFF]) * EnvConfig.FRAME_BYTES

                marker = 1 if (have_real_frame and self._mark_next_real_frame) else 0
                if marker:
                    self._mark_next_real_frame = False

                # Send mark ACKs when we actually transmit a real RTP frame
                if have_real_frame and self._pending_marks and self.ecs_ws:
                    try:
                        ack = {
                            "event": "mark",
                            "streamSid": self.stream_sid,
                            "mark": {"name": self._pending_marks.pop(0)},
                        }
                        await self.ecs_ws.send(json.dumps(ack))
                        logger.debug("[call %s] sent mark ack", self.pstn_channel_id)
                    except Exception:
                        # best-effort; if it fails, we'll drop the ack
                        pass

                self._rtp_send(frame, marker)
                
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error("[call %s] RTP pacer error: %s", self.pstn_channel_id, e, exc_info=True)


    # Lifecycle ___________________________________________________________________________________

    async def _setup(self) -> None:
        """
        Open ECS WS, create ARI/EM, and discover EM RTP.
        """
        # connect to ECS first
        await self._open_ecs_websocket()

        # create ARI bridge and add PSTN channel
        await self._create_bridge()
        await self._add_to_bridge(self.pstn_channel_id)
        
        # allocate RTP socket and create ExternalMedia
        port = await self._alloc_rtp_socket()
        await self._create_external_media(f"{EnvConfig.EXTERNAL_MEDIA_HOST}:{port}")

        # wait for EM channel to be up and add to bridge
        if self.em_channel_id and not await self._wait_channel_up(self.em_channel_id, timeout=5.0):
            logger.warning("[call %s] EM channel not Up before add; continuing", self.pstn_channel_id)
        await self._add_to_bridge(self.em_channel_id)

        # discover EM return address
        raddr = await self._get_channel_var(self.em_channel_id, "UNICASTRTP_LOCAL_ADDRESS")
        rport = await self._get_channel_var(self.em_channel_id, "UNICASTRTP_LOCAL_PORT")
        if not (raddr and rport):
            logger.error("[call %s] could not discover EM local RTP socket", self.pstn_channel_id)
            raise RuntimeError("Could not discover EM local RTP socket")

        self.rtp_send_addr = (raddr, int(rport))
        logger.info("[call %s] EM return socket %s:%s", self.pstn_channel_id, raddr, rport)

    async def start(self) -> None:
        """
        Public entry: initialize and launch background media tasks.
        """
        await self._setup()
        # launch media loops
        self.tasks["rtp_in"] = asyncio.create_task(
            coro=self._rtp_to_ecs(),
            name=f"rtpin-{self.pstn_channel_id}",
        )
        self.tasks["ecs_rx"] = asyncio.create_task(
            coro=self._ecs_to_rtp(),
            name=f"ecsrx-{self.pstn_channel_id}",
        )

    async def stop(self) -> None:
        """
        Idempotent shutdown: cancel tasks, close sockets and websocket.
        """
        if self.closed.is_set():
            logger.debug("[call %s] already closed, skipping stop", self.pstn_channel_id)
            return

        logger.info("[call %s] stopping session", self.pstn_channel_id)
        self.closed.set()

        # send stop event to ECS (best effort)
        try:
            if self.ecs_ws and self.ecs_ws.close_code is None:
                stop_event = {
                    "event": "stop",
                    "streamSid": self.stream_sid
                }
                await self.ecs_ws.send(json.dumps(stop_event))
                logger.debug("[call %s] sent stop event to ECS", self.pstn_channel_id)
        except Exception:
            pass

        # cancel tasks
        logger.info("[call %s] canceling %d background tasks", self.pstn_channel_id, len(self.tasks))
        tasks_to_cancel = list(self.tasks.values())
        self.tasks.clear()

        for task in tasks_to_cancel:
            if not task.done() and not task.cancelled():
                try:
                    task.cancel()
                except Exception:
                    pass

        # clean up ARI resources
        try:
            if self.em_channel_id and not self.http.closed:
                async with self.http.delete(
                    url=f"{EnvConfig.ARI_BASE}/channels/{self.em_channel_id}"
                ) as _:
                    logger.info("[call %s] deleted EM channel %s", self.pstn_channel_id, self.em_channel_id)
        except Exception as e:
            logger.debug("[call %s] error deleting EM channel: %s", self.pstn_channel_id, e)

        # remove channels from bridge and delete it
        try:
            if self.bridge_id:
                for channel_id in [self.pstn_channel_id, self.em_channel_id]:
                    if channel_id and not self.http.closed:
                        try:
                            async with self.http.delete(
                                url=f"{EnvConfig.ARI_BASE}/bridges/{self.bridge_id}/removeChannel", 
                                params={"channel": channel_id}
                            ) as _:
                                pass
                        except Exception:
                            pass

                if not self.http.closed:
                    try:
                        async with self.http.delete(
                            url=f"{EnvConfig.ARI_BASE}/bridges/{self.bridge_id}"
                        ) as _:
                            logger.info("[call %s] deleted bridge %s", self.pstn_channel_id, self.bridge_id)
                    except Exception:
                        pass
        except Exception:
            pass

        # close RTP sockets
        if self.rtp_recv_sock:
            try:
                self.rtp_recv_sock.close()
            except Exception:
                pass

        if self._rtp_send_sock:
            try:
                self._rtp_send_sock.close()
            except Exception:
                pass

        # close ECS WebSocket
        try:
            if self.ecs_ws and self.ecs_ws.close_code is None:
                await self.ecs_ws.close()
        except Exception:
            pass

        logger.info("[call %s] stopped", self.pstn_channel_id)