import os
import struct
import time
import zlib
import hashlib
from typing import Tuple, Optional

def crc32_of(data: bytes) -> int:
    """
    Compute CRC-32 of 'data'.
    """
    return zlib.crc32(data) & 0xffffffff


# ---------------------------------------------------------------------------
# PING REQUEST (new format)
# Fields (in order):
#   - session_nonce: 129 bits (from the top 129 bits of 17 random bytes)
#   - version: 7 bits
#   - cipher: 4 bits   (0 = AES-256-GCM, 1 = ChaCha20-poly1305; for now only 0 is used)
#   - CRC: 32 bits
#
# Total bits: 129 + 7 + 4 + 32 = 172 bits. We pack into 22 bytes (176 bits) with 4 spare bits.
# ---------------------------------------------------------------------------
class PingRequest:
    """
    PING REQUEST format (172 bits / 22 bytes):
    - session_nonce: 129 bits (from top 129 bits of 17 random bytes)
    - version: 7 bits
    - cipher: 4 bits (0 = AES-256-GCM, 1 = ChaCha20-poly1305)
    - CRC: 32 bits
    """
    def __init__(self, version: int, cipher: int, session_nonce: bytes = None):
        if not (0 <= version < 128):
            raise ValueError("Version must fit in 7 bits (0..127)")
        if not (0 <= cipher < 16):
            raise ValueError("Cipher must fit in 4 bits (0..15)")
        
        self.version = version
        self.cipher = cipher
        
        # Generate session nonce if not provided
        if session_nonce is None:
            # Generate 17 random bytes
            nonce_full = os.urandom(17)
            # Use top 129 bits
            nonce_int_full = int.from_bytes(nonce_full, 'big')
            nonce_129_int = nonce_int_full >> 7  # drop lowest 7 bits
            self.session_nonce = nonce_129_int.to_bytes(17, 'big')
        else:
            if len(session_nonce) != 17:
                raise ValueError("Session nonce must be 17 bytes (136 bits)")
            self.session_nonce = session_nonce
    
    def serialize(self) -> bytes:
        """Serialize the ping request into a 22-byte packet."""
        # Convert session_nonce to integer (129 bits)
        nonce_int = int.from_bytes(self.session_nonce, 'big')
        
        # Pack fields: shift nonce left by 11 bits, add version and cipher
        partial_int = (nonce_int << 11) | (self.version << 4) | (self.cipher & 0x0F)
        # This creates 129+7+4 = 140 bits; pack into 18 bytes
        partial_bytes = partial_int.to_bytes(18, 'big')
        
        # Compute CRC over these 18 bytes
        cval = crc32_of(partial_bytes)
        
        # Combine partial data with 32-bit CRC
        final_int = (int.from_bytes(partial_bytes, 'big') << 32) | cval
        return final_int.to_bytes(22, 'big')
    
    @classmethod
    def deserialize(cls, data: bytes) -> Optional['PingRequest']:
        """Deserialize a 22-byte packet into a PingRequest object."""
        if len(data) != 22:
            return None
            
        # Extract 176-bit integer
        final_int = int.from_bytes(data, 'big')
        
        # Extract CRC and verify
        crc_in = final_int & 0xffffffff
        partial_int = final_int >> 32  # 140 bits
        partial_bytes = partial_int.to_bytes(18, 'big')
        crc_calc = crc32_of(partial_bytes)
        
        if crc_calc != crc_in:
            return None
            
        # Extract fields
        cipher = partial_int & 0x0F
        version = (partial_int >> 4) & 0x7F
        nonce_129_int = partial_int >> 11  # 129 bits
        session_nonce = nonce_129_int.to_bytes(17, 'big')
        
        return cls(version, cipher, session_nonce)



# ---------------------------------------------------------------------------
# PING RESPONSE (new format)
# Fields:
#   - timestamp: 32 bits (we take the lower 32 bits of the time in ms)
#   - version: 7 bits
#   - cipher: 4 bits
#   - answer: 1 bit
#   - CRC: 32 bits
#
# Total bits: 32 + 7 + 4 + 1 + 32 = 76 bits; pack into 10 bytes (80 bits) with 4 spare bits.
# ---------------------------------------------------------------------------
class PingResponse:
    """
    PING RESPONSE format (76 bits / 10 bytes):
    - timestamp: 32 bits (milliseconds since epoch, lower 32 bits)
    - version: 7 bits
    - cipher: 4 bits
    - answer: 1 bit (0 = no, 1 = yes)
    - CRC: 32 bits
    """
    def __init__(self, version: int, cipher: int, answer: int, timestamp: int = None):
        if not (0 <= version < 128):
            raise ValueError("Version must fit in 7 bits")
        if not (0 <= cipher < 16):
            raise ValueError("Cipher must fit in 4 bits")
        if answer not in (0, 1):
            raise ValueError("Answer must be 0 or 1")
            
        self.version = version
        self.cipher = cipher
        self.answer = answer
        self.timestamp = timestamp or (int(time.time() * 1000) & 0xffffffff)
    
    def serialize(self) -> bytes:
        """Serialize the ping response into a 10-byte packet."""
        # Pack timestamp, version, cipher, answer: 32+7+4+1 = 44 bits
        partial_val = (self.timestamp << (7+4+1)) | (self.version << (4+1)) | (self.cipher << 1) | self.answer
        partial_bytes = partial_val.to_bytes(6, 'big')  # 6 bytes = 48 bits, 4 spare bits
        
        # Compute CRC
        cval = crc32_of(partial_bytes)
        
        # Combine with CRC
        final_val = (int.from_bytes(partial_bytes, 'big') << 32) | cval
        return final_val.to_bytes(10, 'big')
    
    @classmethod
    def deserialize(cls, data: bytes) -> Optional['PingResponse']:
        """Deserialize a 10-byte packet into a PingResponse object."""
        if len(data) != 10:
            return None
            
        # Extract 80-bit integer
        final_int = int.from_bytes(data, 'big')
        
        # Extract CRC and verify
        crc_in = final_int & 0xffffffff
        partial_int = final_int >> 32  # 48 bits
        partial_bytes = partial_int.to_bytes(6, 'big')
        crc_calc = crc32_of(partial_bytes)
        
        if crc_calc != crc_in:
            return None
            
        # Extract fields (discard 4 spare bits)
        partial_int >>= 4  # now 44 bits
        answer = partial_int & 0x01
        cipher = (partial_int >> 1) & 0x0F
        version = (partial_int >> (1+4)) & 0x7F
        timestamp = partial_int >> (1+4+7)
        
        return cls(version, cipher, answer, timestamp)


# =============================================================================
# 3) Handshake
#    - 32-bit timestamp
#    - 64-byte ephemeral pubkey (raw x||y = 512 bits)
#    - 64-byte ephemeral signature (raw r||s = 512 bits)
#    - 32-byte PFS hash (256 bits)
#    - 32-bit CRC
#    => total 4 + 64 + 64 + 32 + 4 = 168 bytes = 1344 bits
# =============================================================================

class Handshake:
    """
    HANDSHAKE format (1344 bits / 168 bytes):
    - timestamp: 32 bits
    - ephemeral_pubkey: 512 bits (64 bytes, raw x||y format)
    - ephemeral_signature: 512 bits (64 bytes, raw r||s format)
    - pfs_hash: 256 bits (32 bytes)
    - CRC: 32 bits
    """
    def __init__(self, ephemeral_pubkey: bytes, ephemeral_signature: bytes, pfs_hash: bytes, timestamp: int = None):
        if len(ephemeral_pubkey) != 64:
            raise ValueError("ephemeral_pubkey must be 64 bytes (raw x||y)")
        if len(ephemeral_signature) != 64:
            raise ValueError("ephemeral_signature must be 64 bytes (raw r||s)")
        if len(pfs_hash) != 32:
            raise ValueError("pfs_hash must be 32 bytes")
            
        self.ephemeral_pubkey = ephemeral_pubkey
        self.ephemeral_signature = ephemeral_signature
        self.pfs_hash = pfs_hash
        self.timestamp = timestamp or (int(time.time() * 1000) & 0xffffffff)
    
    def serialize(self) -> bytes:
        """Serialize the handshake into a 168-byte packet."""
        # Pack timestamp and other fields
        partial = struct.pack("!I", self.timestamp) + self.ephemeral_pubkey + self.ephemeral_signature + self.pfs_hash
        
        # Compute CRC
        cval = crc32_of(partial)
        
        # Append CRC
        return partial + struct.pack("!I", cval)
    
    @classmethod
    def deserialize(cls, data: bytes) -> Optional['Handshake']:
        """Deserialize a 168-byte packet into a Handshake object."""
        if len(data) != 168:
            return None
            
        # Extract and verify CRC
        partial = data[:-4]
        crc_in = struct.unpack("!I", data[-4:])[0]
        crc_calc = crc32_of(partial)
        
        if crc_calc != crc_in:
            return None
            
        # Extract fields
        timestamp = struct.unpack("!I", partial[:4])[0]
        ephemeral_pubkey = partial[4:4+64]
        ephemeral_signature = partial[68:68+64]
        pfs_hash = partial[132:132+32]
        
        return cls(ephemeral_pubkey, ephemeral_signature, pfs_hash, timestamp)


# =============================================================================
# 4) PFS Hash Helper
#    If no previous session, return 32 zero bytes
#    Otherwise, compute sha256(session_number || last_shared_secret).
# =============================================================================

def compute_pfs_hash(session_number: int, shared_secret_hex: str) -> bytes:
    """
    Compute the PFS hash field for handshake messages:
    - If no previous session (session_number < 0), return 32 zero bytes
    - Otherwise, compute sha256(session_number || shared_secret)
    """
    if session_number < 0:
        return b"\x00" * 32

    # Convert shared_secret_hex to raw bytes
    secret_bytes = bytes.fromhex(shared_secret_hex)
    
    # Pack session_number as 4 bytes
    sn_bytes = struct.pack("!I", session_number)
    
    # Compute hash
    return hashlib.sha256(sn_bytes + secret_bytes).digest()
