"""NMDB message protocol helpers for SalTTS.

Uses NMDB's Protocol Buffers for message serialization.
"""

from typing import Optional

try:
    from nmdb.proto import message_pb2
    from nmdb import CICData
except ImportError:
    # Fallback if nmdb-py not installed
    message_pb2 = None

    class CICData:
        def __init__(self, emb=None):
            self.emb = emb or []


class MessageHelper:
    """Helper class for creating NMDB messages."""

    @staticmethod
    def create_text_message(
        text: str,
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_tokens: int = 100,
    ) -> bytes:
        """Create a text message with generation parameters.

        Args:
            text: Text content
            temperature: Temperature parameter
            top_p: Top-p parameter
            max_tokens: Maximum tokens

        Returns:
            Serialized message bytes
        """
        if message_pb2 is None:
            # Fallback: return text as bytes
            return text.encode('utf-8')

        msg = message_pb2.Message()
        msg.text = text
        msg.params.temperature = temperature
        msg.params.top_p = top_p
        msg.params.max_tokens = max_tokens

        return msg.SerializeToString()

    @staticmethod
    def parse_text_message(data: bytes) -> Optional[str]:
        """Parse a text message.

        Args:
            data: Serialized message bytes

        Returns:
            Text content or None
        """
        if message_pb2 is None:
            # Fallback: decode as UTF-8
            try:
                return data.decode('utf-8')
            except:
                return None

        try:
            msg = message_pb2.Message()
            msg.ParseFromString(data)
            return msg.text
        except:
            return None

    @staticmethod
    def create_cic_with_embedding(embedding: list) -> CICData:
        """Create CICData with embedding.

        Args:
            embedding: Embedding vector as list

        Returns:
            CICData object
        """
        return CICData(emb=embedding)

    @staticmethod
    def create_nested_cic(embeddings: list) -> CICData:
        """Create nested CICData.

        Args:
            embeddings: List of embedding vectors

        Returns:
            CICData with nested data
        """
        nested = [CICData(emb=emb) for emb in embeddings]
        return CICData(nested=nested)
