"""Encrypted DataSource wrapper for transparent encryption during upload."""

from __future__ import annotations

from ..transfer.encryption import (
    ENCRYPTION_HEADER_SIZE,
    EncryptionHeader,
    crypt_at,
)
from .data import DataSource


class EncryptedDataSource:
    """Wraps a DataSource to transparently encrypt reads.

    The encrypted stream is: [header (17 bytes)] [AES-256-CTR encrypted data].
    Implements the DataSource protocol so it can be passed directly to the upload pipeline.
    """

    def __init__(self, inner: DataSource, key: bytes) -> None:
        if len(key) != 32:
            raise ValueError(f"Encryption key must be 32 bytes, got {len(key)}")
        self._inner = inner
        self._key = key
        self._header = EncryptionHeader()

    @property
    def header(self) -> EncryptionHeader:
        return self._header

    def size(self) -> int:
        """Return the encrypted size (original size + 17-byte header)."""
        return self._inner.size() + ENCRYPTION_HEADER_SIZE

    def read_at(self, offset: int, size: int) -> bytes:
        """Read encrypted bytes at the given offset."""
        encrypted_size = self.size()
        if offset < 0 or offset >= encrypted_size:
            return b""

        # Clamp to available size
        available = encrypted_size - offset
        size = min(size, available)

        header_size = ENCRYPTION_HEADER_SIZE
        header_bytes = self._header.to_bytes()
        result = bytearray()

        # If offset falls within the header region
        if offset < header_size:
            header_start = offset
            header_end = min(header_size, offset + size)
            n = header_end - header_start
            result.extend(header_bytes[header_start:header_end])
            size -= n
            offset = header_size  # advance past header

        # Read and encrypt data from inner source
        if size > 0:
            data_offset = offset - header_size
            inner_data = self._inner.read_at(data_offset, size)
            if inner_data:
                buf = bytearray(inner_data)
                crypt_at(self._key, self._header.nonce, data_offset, buf)
                result.extend(buf)

        return bytes(result)

    def close(self) -> None:
        """Close the underlying data source."""
        self._inner.close()

    def __enter__(self) -> EncryptedDataSource:
        if hasattr(self._inner, "__enter__"):
            self._inner.__enter__()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        if hasattr(self._inner, "__exit__"):
            self._inner.__exit__(exc_type, exc, tb)
        else:
            self.close()
