"""TataTele Smartflo bidirectional audio streaming WebSocket protocol serializer for Pipecat."""

import base64
import json
from typing import Optional

from loguru import logger
from pydantic import BaseModel

from pipecat.audio.dtmf.types import KeypadEntry
from pipecat.audio.utils import create_stream_resampler, pcm_to_ulaw, ulaw_to_pcm
from pipecat.frames.frames import (
    AudioRawFrame,
    Frame,
    InputAudioRawFrame,
    InputDTMFFrame,
    InterruptionFrame,
    OutputTransportMessageFrame,
    OutputTransportMessageUrgentFrame,
    StartFrame,
)
from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType


class TataTeleFrameSerializer(FrameSerializer):
    """Serializer for TataTele Smartflo bidirectional audio streaming WebSocket protocol.

    This serializer handles converting between Pipecat frames and TataTele's WebSocket
    media streams protocol. It supports audio conversion and DTMF events.

    The protocol is similar to Twilio Media Streams:
    - Audio: G.711 μ-law (PCMU), 8000 Hz, 64 kbps, 8-bit mono
    - Media payloads sent TO TataTele must be multiples of 160 bytes
    - Events: connected → start → media → stop

    Call termination is handled by closing the WebSocket connection.
    """

    class InputParams(BaseModel):
        """Configuration parameters for TataTeleFrameSerializer.

        Parameters:
            tatatele_sample_rate: Sample rate used by TataTele, defaults to 8000 Hz.
            sample_rate: Optional override for pipeline input sample rate.
        """

        tatatele_sample_rate: int = 8000
        sample_rate: Optional[int] = None

    # TataTele requires media payloads to be multiples of 160 bytes
    PAYLOAD_ALIGNMENT = 160

    def __init__(
        self,
        stream_sid: str,
        params: Optional[InputParams] = None,
    ):
        """Initialize the TataTeleFrameSerializer.

        Args:
            stream_sid: The TataTele Media Stream SID.
            params: Configuration parameters.
        """
        self._params = params or TataTeleFrameSerializer.InputParams()

        self._stream_sid = stream_sid

        self._tatatele_sample_rate = self._params.tatatele_sample_rate
        self._sample_rate = 0  # Pipeline input rate

        self._input_resampler = create_stream_resampler()
        self._output_resampler = create_stream_resampler()
        self._output_buffer = bytearray()

    @property
    def type(self) -> FrameSerializerType:
        return FrameSerializerType.TEXT

    async def setup(self, frame: StartFrame):
        self._sample_rate = self._params.sample_rate or frame.audio_in_sample_rate

    async def serialize(self, frame: Frame) -> str | bytes | None:
        """Serializes a Pipecat frame to TataTele WebSocket format.

        Handles conversion of various frame types to TataTele WebSocket messages.
        Audio payloads are aligned to multiples of 160 bytes as required by TataTele.
        """
        if isinstance(frame, InterruptionFrame):
            answer = {"event": "clear", "streamSid": self._stream_sid}
            return json.dumps(answer)
        elif isinstance(frame, AudioRawFrame):
            data = frame.audio

            serialized_data = await pcm_to_ulaw(
                data, frame.sample_rate, self._tatatele_sample_rate, self._output_resampler
            )
            if serialized_data is None or len(serialized_data) == 0:
                return None

            # Buffer and align to multiples of 160 bytes
            self._output_buffer.extend(serialized_data)
            aligned_length = (
                len(self._output_buffer) // self.PAYLOAD_ALIGNMENT * self.PAYLOAD_ALIGNMENT
            )
            if aligned_length == 0:
                return None

            to_send = bytes(self._output_buffer[:aligned_length])
            self._output_buffer = self._output_buffer[aligned_length:]

            payload = base64.b64encode(to_send).decode("utf-8")
            answer = {
                "event": "media",
                "streamSid": self._stream_sid,
                "media": {"payload": payload},
            }

            return json.dumps(answer)
        elif isinstance(frame, (OutputTransportMessageFrame, OutputTransportMessageUrgentFrame)):
            return json.dumps(frame.message)

        return None

    async def deserialize(self, data: str | bytes) -> Frame | None:
        """Deserializes TataTele WebSocket data to Pipecat frames.

        Handles conversion of TataTele media events to appropriate Pipecat frames.
        """
        message = json.loads(data)

        if message.get("event") == "media":
            payload_base64 = message["media"]["payload"]
            payload = base64.b64decode(payload_base64)

            deserialized_data = await ulaw_to_pcm(
                payload, self._tatatele_sample_rate, self._sample_rate, self._input_resampler
            )
            if deserialized_data is None or len(deserialized_data) == 0:
                return None

            audio_frame = InputAudioRawFrame(
                audio=deserialized_data, num_channels=1, sample_rate=self._sample_rate
            )
            return audio_frame
        elif message.get("event") == "dtmf":
            digit = message.get("dtmf", {}).get("digit")
            if digit:
                try:
                    return InputDTMFFrame(KeypadEntry(digit))
                except ValueError:
                    logger.warning(f"Invalid DTMF digit received: {digit}")
                    return None
        else:
            return None
