import os
import struct
from typing import Optional, Tuple
from cryptography.hazmat.primitives.ciphers.aead import AESGCM, ChaCha20Poly1305

class MessageHeader:
    """
    Header of an encrypted message (18 bytes total):
    
    Clear Text Section (4 bytes):
    - flag: 16 bits (0xBEEF by default)
    - data_len: 16 bits (length of encrypted payload excluding tag)
    
    Associated Data (14 bytes):
    - retry: 8 bits (retry counter)
    - connection_status: 4 bits (e.g., CRC required) + 4 bits padding
    - iv/messageID: 96 bits (12 bytes)
    """
    def __init__(self, flag: int, data_len: int, retry: int, connection_status: int, iv: bytes):
        if not (0 <= flag < 65536):
            raise ValueError("Flag must fit in 16 bits (0..65535)")
        if not (0 <= data_len < 65536):
            raise ValueError("Data length must fit in 16 bits (0..65535)")
        if not (0 <= retry < 256):
            raise ValueError("Retry must fit in 8 bits (0..255)")
        if not (0 <= connection_status < 16):
            raise ValueError("Connection status must fit in 4 bits (0..15)")
        if len(iv) != 12:
            raise ValueError("IV must be 12 bytes (96 bits)")
            
        self.flag = flag                  # 16 bits
        self.data_len = data_len          # 16 bits
        self.retry = retry                # 8 bits
        self.connection_status = connection_status  # 4 bits
        self.iv = iv                      # 96 bits (12 bytes)

    def pack(self) -> bytes:
        """Pack header into 18 bytes."""
        # Pack flag and data_len (4 bytes)
        header = struct.pack('>H H', self.flag, self.data_len)
        
        # Pack retry and connection_status (2 bytes)
        # connection_status in high 4 bits of second byte, 4 bits padding as zero
        ad_byte = (self.connection_status & 0x0F) << 4
        ad_packed = struct.pack('>B B', self.retry, ad_byte)
        
        # Append IV (12 bytes)
        return header + ad_packed + self.iv
        
    def get_associated_data(self) -> bytes:
        """Get the associated data for AEAD encryption (retry, conn_status, iv)."""
        # Pack retry and connection_status
        ad_byte = (self.connection_status & 0x0F) << 4
        ad_packed = struct.pack('>B B', self.retry, ad_byte)
        
        # Append IV
        return ad_packed + self.iv
        
    @classmethod
    def unpack(cls, data: bytes) -> 'MessageHeader':
        """Unpack 18 bytes into a MessageHeader object."""
        if len(data) < 18:
            raise ValueError(f"Header data too short: {len(data)} bytes, expected 18")
            
        flag, data_len = struct.unpack('>H H', data[:4])
        retry, ad_byte = struct.unpack('>B B', data[4:6])
        connection_status = (ad_byte >> 4) & 0x0F
        iv = data[6:18]
        
        return cls(flag, data_len, retry, connection_status, iv)

class EncryptedMessage:
    """
    Encrypted message packet format:
    
    - Header (18 bytes):
      * flag: 16 bits
      * data_len: 16 bits
      * retry: 8 bits
      * connection_status: 4 bits (+ 4 bits padding)
      * iv/messageID: 96 bits (12 bytes)
    
    - Payload: variable length encrypted data
    
    - Footer:
      * Authentication tag: 128 bits (16 bytes)
      * CRC32: 32 bits (4 bytes) - optional, based on connection_status
    """
    def __init__(self, plaintext: bytes, key: bytes, flag: int = 0xBEEF, 
                 retry: int = 0, connection_status: int = 0, iv: bytes = None,
                 cipher_type: int = 0):
        self.plaintext = plaintext
        self.key = key
        self.flag = flag
        self.retry = retry
        self.connection_status = connection_status
        self.iv = iv or generate_iv(initial=True)
        self.cipher_type = cipher_type  # 0 = AES-256-GCM, 1 = ChaCha20-Poly1305
        
        # Will be set after encryption
        self.ciphertext = None
        self.tag = None
        self.header = None
        
    def encrypt(self) -> bytes:
        """Encrypt the plaintext and return the full encrypted message."""
        # Create header with correct data_len (which will be set after encryption)
        self.header = MessageHeader(
            flag=self.flag,
            data_len=0,  # Will be updated after encryption
            retry=self.retry,
            connection_status=self.connection_status,
            iv=self.iv
        )
        
        # Get associated data for AEAD
        aad = self.header.get_associated_data()
        
        # Encrypt using the appropriate cipher
        if self.cipher_type == 0:  # AES-256-GCM
            cipher = AESGCM(self.key)
            ciphertext_with_tag = cipher.encrypt(self.iv, self.plaintext, aad)
        elif self.cipher_type == 1:  # ChaCha20-Poly1305
            cipher = ChaCha20Poly1305(self.key)
            ciphertext_with_tag = cipher.encrypt(self.iv, self.plaintext, aad)
        else:
            raise ValueError(f"Unsupported cipher type: {self.cipher_type}")
        
        # Extract ciphertext and tag
        self.tag = ciphertext_with_tag[-16:]
        self.ciphertext = ciphertext_with_tag[:-16]
        
        # Update header with actual data length
        self.header.data_len = len(self.ciphertext)
        
        # Pack everything together
        packed_header = self.header.pack()
        
        # Check if CRC is required (based on connection_status)
        if self.connection_status & 0x01:  # Lowest bit indicates CRC required
            import zlib
            # Compute CRC32 of header + ciphertext + tag
            crc = zlib.crc32(packed_header + self.ciphertext + self.tag) & 0xffffffff
            crc_bytes = struct.pack('>I', crc)
            return packed_header + self.ciphertext + self.tag + crc_bytes
        else:
            return packed_header + self.ciphertext + self.tag
            
    @classmethod
    def decrypt(cls, data: bytes, key: bytes, cipher_type: int = 0) -> Tuple[bytes, MessageHeader]:
        """
        Decrypt an encrypted message and return the plaintext and header.
        
        Args:
            data: The full encrypted message
            key: The encryption key
            cipher_type: 0 for AES-256-GCM, 1 for ChaCha20-Poly1305
            
        Returns:
            Tuple of (plaintext, header)
        """
        if len(data) < 18 + 16:  # Header + minimum tag size
            raise ValueError("Message too short")
            
        # Extract header
        header_bytes = data[:18]
        header = MessageHeader.unpack(header_bytes)
        
        # Get ciphertext and tag
        data_len = header.data_len
        ciphertext_start = 18
        ciphertext_end = ciphertext_start + data_len
        
        if ciphertext_end + 16 > len(data):
            raise ValueError("Message length does not match header's data_len")
            
        ciphertext = data[ciphertext_start:ciphertext_end]
        tag = data[ciphertext_end:ciphertext_end + 16]
        
        # Get associated data for AEAD
        aad = header.get_associated_data()
        
        # Combine ciphertext and tag for decryption
        ciphertext_with_tag = ciphertext + tag
        
        # Decrypt using the appropriate cipher
        try:
            if cipher_type == 0:  # AES-256-GCM
                cipher = AESGCM(key)
                plaintext = cipher.decrypt(header.iv, ciphertext_with_tag, aad)
            elif cipher_type == 1:  # ChaCha20-Poly1305
                cipher = ChaCha20Poly1305(key)
                plaintext = cipher.decrypt(header.iv, ciphertext_with_tag, aad)
            else:
                raise ValueError(f"Unsupported cipher type: {cipher_type}")
                
            return plaintext, header
        except Exception as e:
            raise ValueError(f"Decryption failed: {e}")

def generate_iv(initial: bool = False, previous_iv: bytes = None) -> bytes:
    """
    Generate a 96-bit IV (12 bytes).
    
    Args:
        initial: If True, return a random IV
        previous_iv: The previous IV to increment
        
    Returns:
        A new IV
    """
    if initial or previous_iv is None:
        return os.urandom(12)  # 96 bits
    else:
        # Increment the previous IV by 1 modulo 2^96
        iv_int = int.from_bytes(previous_iv, 'big')
        iv_int = (iv_int + 1) % (1 << 96)
        return iv_int.to_bytes(12, 'big')

# Convenience functions to match original API
def encrypt_message(plaintext: bytes, key: bytes, flag: int = 0xBEEF, 
                   retry: int = 0, connection_status: int = 0, 
                   iv: bytes = None, cipher_type: int = 0) -> bytes:
    """
    Encrypt a message using the specified parameters.
    
    Args:
        plaintext: The data to encrypt
        key: The encryption key (32 bytes for AES-256-GCM, 32 bytes for ChaCha20-Poly1305)
        flag: 16-bit flag value (default: 0xBEEF)
        retry: 8-bit retry counter
        connection_status: 4-bit connection status
        iv: Optional 96-bit IV (if None, a random one will be generated)
        cipher_type: 0 for AES-256-GCM, 1 for ChaCha20-Poly1305
        
    Returns:
        The full encrypted message
    """
    message = EncryptedMessage(
        plaintext=plaintext,
        key=key,
        flag=flag,
        retry=retry,
        connection_status=connection_status,
        iv=iv,
        cipher_type=cipher_type
    )
    return message.encrypt()

def decrypt_message(message: bytes, key: bytes, cipher_type: int = 0) -> bytes:
    """
    Decrypt a message.
    
    Args:
        message: The full encrypted message
        key: The encryption key
        cipher_type: 0 for AES-256-GCM, 1 for ChaCha20-Poly1305
        
    Returns:
        The decrypted plaintext
    """
    plaintext, _ = EncryptedMessage.decrypt(message, key, cipher_type)
    return plaintext
