"""Pydantic models for GA4GH Crypt4GH file encryption.

Crypt4GH is a standard for encrypting genomic data files using
an envelope encryption scheme with X25519 key exchange and
ChaCha20-Poly1305 symmetric encryption.
"""

from __future__ import annotations

from enum import Enum

from ai4gh.core.types import GA4GHBase

MAGIC_BYTES = b"crypt4gh"
VERSION = 1


class HeaderPacketType(int, Enum):
    """Crypt4GH header packet types."""

    DATA_ENCRYPTION_PARAMETERS = 0
    DATA_EDIT_LIST = 1


class DataEncryptionMethod(int, Enum):
    """Encryption methods for data segments."""

    CHACHA20_IETF_POLY1305 = 0


class KeyPair(GA4GHBase):
    """A Crypt4GH key pair (public + private)."""

    public_key: bytes
    private_key: bytes


class HeaderPacket(GA4GHBase):
    """A Crypt4GH header packet (after decryption)."""

    packet_type: HeaderPacketType
    data_encryption_method: DataEncryptionMethod | None = None
    data_key: bytes | None = None
    edit_list: list[tuple[int, int]] | None = None


class Crypt4GHHeader(GA4GHBase):
    """Parsed Crypt4GH file header."""

    magic: bytes = MAGIC_BYTES
    version: int = VERSION
    header_packet_count: int = 0
    packets: list[HeaderPacket] | None = None


class EncryptionResult(GA4GHBase):
    """Result of a Crypt4GH encryption operation."""

    encrypted_data: bytes
    header: Crypt4GHHeader
    data_key: bytes
