"""Data-embedding pair handling using NMDB CICData."""

from dataclasses import dataclass
from typing import Optional, List
import numpy as np

try:
    from nmdb import CICData
except ImportError:
    # Fallback if nmdb-py not installed
    class CICData:
        def __init__(self, emb=None, nested=None):
            self.emb = emb or []
            self.nested = nested or []

        def to_bytes(self):
            return b""

        @classmethod
        def from_bytes(cls, data):
            return cls()


@dataclass
class DataEmbeddingPair:
    """Represents a data-embedding pair from NMDB channel.

    Uses NMDB's CICData format for embedding data.

    Attributes:
        cic_data: CICData object containing embedding
        channel_name: Source channel name
        timestamp: Timestamp of message
    """
    cic_data: CICData
    channel_name: str
    timestamp: float

    @property
    def embedding(self) -> np.ndarray:
        """Get embedding as numpy array."""
        return np.array(self.cic_data.emb, dtype=np.float32)

    @property
    def nested_embeddings(self) -> List[np.ndarray]:
        """Get nested embeddings as list of numpy arrays."""
        return [np.array(nested.emb, dtype=np.float32) for nested in self.cic_data.nested]

    def has_nested(self) -> bool:
        """Check if has nested CIC data."""
        return len(self.cic_data.nested) > 0


class DataEmbeddingPairProcessor:
    """Process data-embedding pairs from NMDB channels."""

    def __init__(self):
        """Initialize processor."""
        self.channel_configs = {
            "text1": {"type": "text", "embedding_dim": 384},
            "text2": {"type": "text", "embedding_dim": 512},
            "data3": {"type": "binary", "embedding_dim": 512},
        }

    def validate_pair(self, pair: DataEmbeddingPair) -> bool:
        """Validate a data-embedding pair.

        Args:
            pair: Data-embedding pair to validate

        Returns:
            True if valid
        """
        if pair.channel_name not in self.channel_configs:
            return False

        expected_dim = self.channel_configs[pair.channel_name]["embedding_dim"]
        if len(pair.embedding) != expected_dim:
            return False

        return True

    def merge_embeddings(self, pairs: List[DataEmbeddingPair]) -> np.ndarray:
        """Merge embeddings from multiple pairs.

        Args:
            pairs: List of data-embedding pairs

        Returns:
            Merged embedding vector
        """
        embeddings = [pair.embedding for pair in pairs]
        return np.concatenate(embeddings)
