"""AES-256-CTR encryption for 0G Storage.

File format: [version (1 byte)] [nonce (16 bytes)] [encrypted data]

AES-256-CTR is symmetric: encrypt and decrypt are the same operation.
"""

from __future__ import annotations

import os
import struct

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

ENCRYPTION_HEADER_SIZE = 17
ENCRYPTION_VERSION = 1


class EncryptionHeader:
    """Encryption header containing version and nonce."""

    __slots__ = ("version", "nonce")

    def __init__(self, version: int = ENCRYPTION_VERSION, nonce: bytes | None = None) -> None:
        self.version = version
        self.nonce = nonce if nonce is not None else os.urandom(16)

    @classmethod
    def parse(cls, data: bytes) -> EncryptionHeader:
        """Parse a header from the first 17 bytes of encrypted data."""
        if len(data) < ENCRYPTION_HEADER_SIZE:
            raise ValueError(
                f"Data too short for encryption header: {len(data)} < {ENCRYPTION_HEADER_SIZE}"
            )
        version = data[0]
        if version != ENCRYPTION_VERSION:
            raise ValueError(f"Unsupported encryption version: {version}")
        nonce = bytes(data[1:17])
        return cls(version=version, nonce=nonce)

    def to_bytes(self) -> bytes:
        """Serialize to 17 bytes."""
        return struct.pack("B", self.version) + self.nonce


def crypt_at(key: bytes, nonce: bytes, offset: int, data: bytearray) -> None:
    """Encrypt or decrypt data in-place at a given byte offset.

    AES-256-CTR is symmetric: encrypt and decrypt are the same operation.
    ``offset`` is the byte offset within the plaintext data stream (not counting the header).
    """
    if len(key) != 32:
        raise ValueError(f"Encryption key must be 32 bytes, got {len(key)}")
    if len(nonce) != 16:
        raise ValueError(f"Nonce must be 16 bytes, got {len(nonce)}")
    if not data:
        return

    # Compute the CTR counter value at the given offset.
    # AES block size is 16 bytes. We need to advance the counter by offset // 16 blocks
    # and handle the partial-block prefix.
    block_size = 16
    block_offset = offset // block_size
    byte_offset_in_block = offset % block_size

    # The initial counter is the nonce interpreted as a 128-bit big-endian integer,
    # advanced by block_offset blocks.
    nonce_int = int.from_bytes(nonce, "big")
    counter_start = (nonce_int + block_offset) % (1 << 128)
    adjusted_nonce = counter_start.to_bytes(16, "big")

    cipher = Cipher(algorithms.AES(key), modes.CTR(adjusted_nonce))
    encryptor = cipher.encryptor()

    if byte_offset_in_block > 0:
        # Burn the partial block prefix to align the keystream
        encryptor.update(b"\x00" * byte_offset_in_block)

    result = encryptor.update(bytes(data)) + encryptor.finalize()
    data[:] = result


def encrypt_file(key: bytes, plaintext: bytes) -> bytes:
    """Encrypt plaintext and prepend the encryption header.

    Returns: header (17 bytes) + encrypted data.
    """
    header = EncryptionHeader()
    buf = bytearray(plaintext)
    crypt_at(key, header.nonce, 0, buf)
    return header.to_bytes() + bytes(buf)


def decrypt_file(key: bytes, encrypted: bytes) -> bytes:
    """Decrypt a full file: strip header, decrypt remaining bytes."""
    if len(encrypted) < ENCRYPTION_HEADER_SIZE:
        raise ValueError("Encrypted data too short")
    header = EncryptionHeader.parse(encrypted)
    buf = bytearray(encrypted[ENCRYPTION_HEADER_SIZE:])
    crypt_at(key, header.nonce, 0, buf)
    return bytes(buf)


def decrypt_segment(
    key: bytes,
    segment_index: int,
    segment_size: int,
    segment_data: bytes,
    header: EncryptionHeader,
) -> bytes:
    """Decrypt a single downloaded segment.

    For segment 0: skip the 17-byte header, decrypt at data offset 0.
    For other segments: decrypt using the header's nonce with the correct data offset.
    """
    if segment_index == 0:
        buf = bytearray(segment_data[ENCRYPTION_HEADER_SIZE:])
        crypt_at(key, header.nonce, 0, buf)
        return bytes(buf)
    else:
        data_offset = segment_index * segment_size - ENCRYPTION_HEADER_SIZE
        buf = bytearray(segment_data)
        crypt_at(key, header.nonce, data_offset, buf)
        return bytes(buf)
