"""
High-level DAVE session management, using the dataclasses type.
"""

# The dataclasses type might be changed to attrs later. This is more just prototyping.

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Optional

from openmls_dave._lib import (
    CommitResult,
    KeyRatchetManager,
    MlsClient,
    MlsGroupWrapper,
    ProcessedMessage,
    SFrameContext,
    SFrameMultiContext,
)


class SessionState(Enum):
    """DAVE session states."""

    INITIALIZED = auto()
    JOINED = auto()
    ACTIVE = auto()
    LEFT = auto()


@dataclass
class DaveParticipant:
    """Represents a participant in a DAVE session."""

    user_id: int
    key_id: int
    identity: bytes

    @classmethod
    def from_identity(cls, identity: bytes) -> "DaveParticipant":
        """Create participant from MLS identity."""
        # Parse user_id from identity if encoded
        try:
            user_id = int(identity.decode())
        except (ValueError, UnicodeDecodeError):
            user_id = hash(identity) & 0xFFFFFFFF

        return cls(
            user_id=user_id,
            key_id=user_id & 0xFFFF,
            identity=identity,
        )


@dataclass
class EncryptedFrame:
    """An encrypted media frame."""

    data: bytes
    key_id: int
    counter: int


@dataclass
class DecryptedFrame:
    """A decrypted media frame."""

    data: bytes
    sender_key_id: int
    counter: int


class DaveSession:
    """
    High-level DAVE session for managing encrypted voice/video calls.

    Example usage:

    ```python
    # Create a session
    session = DaveSession(user_id=12345)

    # Create a new call (as host)
    key_package = session.create_call(call_id=b"call_abc123")

    # Or join an existing call
    session.join_call(welcome_message=welcome_bytes)

    # Encrypt outgoing media
    encrypted = session.encrypt_frame(audio_data)

    # Decrypt incoming media
    decrypted = session.decrypt_frame(encrypted_data)
    ```
    """

    def __init__(self, user_id: int):
        """
        Initialize a DAVE session.

        Args:
            user_id: The local user's ID.
        """
        self.user_id = user_id
        self.key_id = user_id & 0xFFFF  # Use lower 16 bits as key ID

        # Create MLS client with user ID as identity
        self._client = MlsClient(str(user_id).encode())
        self._group: Optional[MlsGroupWrapper] = None
        self._sframe: Optional[SFrameContext] = None
        self._multi_sframe: Optional[SFrameMultiContext] = None
        self._ratchet_manager: Optional[KeyRatchetManager] = None

        self._state = SessionState.INITIALIZED
        self._epoch = 0
        self._participants: dict[int, DaveParticipant] = {}

    @property
    def state(self) -> SessionState:
        """Current session state."""
        return self._state

    @property
    def epoch(self) -> int:
        """Current MLS epoch."""
        return self._group.epoch if self._group else 0

    @property
    def group_id(self) -> Optional[bytes]:
        """Current group ID, if in a call."""
        return self._group.group_id if self._group else None

    @property
    def participant_count(self) -> int:
        """Number of participants in the call."""
        return self._group.member_count if self._group else 0

    def get_key_package(self) -> bytes:
        """
        Generate a key package for joining a call.

        Returns:
            Serialized key package bytes.
        """
        return self._client.generate_key_package()

    def create_call(self, call_id: bytes) -> bytes:
        """
        Create a new call as the initial member.

        Args:
            call_id: Unique identifier for the call.

        Returns:
            Key package for sharing with other participants.
        """
        self._group = MlsGroupWrapper(self._client, call_id)
        self._initialize_encryption()
        self._state = SessionState.ACTIVE

        return self.get_key_package()

    def join_call(self, welcome_message: bytes) -> None:
        """
        Join an existing call via Welcome message.

        Args:
            welcome_message: MLS Welcome message from the call creator.
        """
        self._group = MlsGroupWrapper.from_welcome(self._client, welcome_message)
        self._initialize_encryption()
        self._state = SessionState.ACTIVE

    def add_participant(self, key_package: bytes) -> CommitResult:
        """
        Add a participant to the call.

        Args:
            key_package: The new participant's key package.

        Returns:
            CommitResult containing commit and welcome messages.
        """
        if not self._group:
            raise RuntimeError("Not in a call")

        result = self._group.add_member(key_package)
        self._group.merge_pending_commit()
        self._refresh_encryption()

        return result

    def remove_participant(self, user_id: int) -> CommitResult:
        """
        Remove a participant from the call.

        Args:
            user_id: The user ID to remove.

        Returns:
            CommitResult containing the commit message.
        """
        if not self._group:
            raise RuntimeError("Not in a call")

        identity = str(user_id).encode()
        result = self._group.remove_member(identity)
        self._group.merge_pending_commit()
        self._refresh_encryption()

        return result

    def process_message(self, message: bytes) -> ProcessedMessage:
        """
        Process an incoming MLS message.

        Args:
            message: The MLS message bytes.

        Returns:
            ProcessedMessage with message type and sender info.
        """
        if not self._group:
            raise RuntimeError("Not in a call")

        result = self._group.process_message(message)

        # Refresh encryption on epoch change
        if result.message_type == "commit":
            self._refresh_encryption()

        return result

    def encrypt_frame(self, plaintext: bytes) -> EncryptedFrame:
        """
        Encrypt a media frame for transmission.

        Args:
            plaintext: Raw media frame data.

        Returns:
            EncryptedFrame with the encrypted data.
        """
        if not self._sframe:
            raise RuntimeError("Encryption not initialized")

        counter = self._sframe.current_counter
        encrypted = self._sframe.encrypt(plaintext)

        return EncryptedFrame(
            data=encrypted,
            key_id=self.key_id,
            counter=counter,
        )

    def decrypt_frame(self, ciphertext: bytes) -> DecryptedFrame:
        """
        Decrypt an incoming media frame.

        Args:
            ciphertext: Encrypted media frame.

        Returns:
            DecryptedFrame with the decrypted data.
        """
        if not self._multi_sframe:
            raise RuntimeError("Decryption not initialized")

        plaintext, header = self._multi_sframe.decrypt(ciphertext)

        return DecryptedFrame(
            data=plaintext,
            sender_key_id=header.key_id,
            counter=header.counter,
        )

    def register_sender(self, sender_key_id: int, base_key: bytes) -> None:
        """
        Register a sender's key for decryption.

        Args:
            sender_key_id: The sender's key ID.
            base_key: The sender's base key.
        """
        if self._multi_sframe:
            self._multi_sframe.add_sender(sender_key_id, base_key)

    def unregister_sender(self, sender_key_id: int) -> None:
        """
        Unregister a sender's key.

        Args:
            sender_key_id: The sender's key ID to remove.
        """
        if self._multi_sframe:
            self._multi_sframe.remove_sender(sender_key_id)

    def leave_call(self) -> None:
        """Leave the current call."""
        self._group = None
        self._sframe = None
        self._multi_sframe = None
        self._ratchet_manager = None
        self._participants.clear()
        self._state = SessionState.LEFT

    def get_epoch_authenticator(self) -> bytes:
        """
        Get the epoch authenticator for verification.

        Returns:
            The epoch authenticator bytes.
        """
        if not self._group:
            raise RuntimeError("Not in a call")
        return self._group.get_epoch_authenticator()

    def _initialize_encryption(self) -> None:
        """Initialize SFrame encryption contexts."""
        if not self._group:
            return

        # Derive media key from MLS group
        media_key = self._group.derive_media_key()

        # Create local SFrame context
        self._sframe = SFrameContext(media_key, self.key_id)

        # Create multi-sender context for decryption
        self._multi_sframe = SFrameMultiContext(self.key_id)
        self._multi_sframe.add_sender(self.key_id, media_key)

        # Initialize key ratchet manager
        self._ratchet_manager = KeyRatchetManager(media_key)

    def _refresh_encryption(self) -> None:
        """Refresh encryption contexts after epoch change."""
        self._initialize_encryption()

    def __repr__(self) -> str:
        return f"DaveSession(user_id={self.user_id}, " f"state={self._state.name}, " f"epoch={self.epoch})"
