"""
Type stubs for the session module.
"""

from dataclasses import dataclass
from enum import Enum
from typing import Optional

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

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

    INITIALIZED = ...
    JOINED = ...
    ACTIVE = ...
    LEFT = ...

@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."""
        ...

@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)
    ```
    """

    user_id: int
    key_id: int

    _client: MlsClient
    _group: Optional[MlsGroupWrapper]
    _sframe: Optional[SFrameContext]
    _multi_sframe: Optional[SFrameMultiContext]
    _ratchet_manager: Optional[KeyRatchetManager]
    _state: SessionState
    _epoch: int
    _participants: dict[int, DaveParticipant]

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

        Args:
            user_id: The local user's ID.
        """
        ...

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

    @property
    def epoch(self) -> int:
        """Current MLS epoch."""
        ...

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

    @property
    def participant_count(self) -> int:
        """Number of participants in the call."""
        ...

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

        Returns:
            Serialized key package bytes.
        """
        ...

    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.
        """
        ...

    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.
        """
        ...

    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.

        Raises:
            RuntimeError: If not in a call.
        """
        ...

    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.

        Raises:
            RuntimeError: If not in a call.
        """
        ...

    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.

        Raises:
            RuntimeError: If not in a call.
        """
        ...

    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.

        Raises:
            RuntimeError: If encryption not initialized.
        """
        ...

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

        Args:
            ciphertext: Encrypted media frame.

        Returns:
            DecryptedFrame with the decrypted data.

        Raises:
            RuntimeError: If decryption not initialized.
        """
        ...

    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.
        """
        ...

    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.
        """
        ...

    def leave_call(self) -> None:
        """Leave the current call."""
        ...

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

        Returns:
            The epoch authenticator bytes.

        Raises:
            RuntimeError: If not in a call.
        """
        ...

    def __repr__(self) -> str: ...
