"""Seekable reader for account-bound .3dripper Scene IR containers."""

from collections import OrderedDict
import io
import json
import os
import re
import struct


MAGIC = b"3DRIPR01"
FORMAT = "3dripper/blender-scene-ir-protected"
VERSION = 1
ALGORITHM = "AES-256-GCM-CHUNKED"
KEY_VERSION = 1
AUTH_TAG_BYTES = 16
NONCE_PREFIX_BYTES = 8
MIN_CHUNK_SIZE = 64 * 1024
MAX_CHUNK_SIZE = 8 * 1024 * 1024
MAX_HEADER_BYTES = 64 * 1024
MAX_PLAINTEXT_BYTES = 8 * 1024 * 1024 * 1024
KEY_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]{16,200}$")


class ProtectedPackageError(ValueError):
    pass


def is_protected_package_path(path):
    try:
        with open(path, "rb") as source:
            return source.read(len(MAGIC)) == MAGIC
    except OSError:
        return False


class ProtectedPackageStream(io.RawIOBase):
    """Expose the encrypted inner ZIP as a seekable, lazily decrypted stream."""

    def __init__(self, path, key_resolver, cache_chunks=2):
        super().__init__()
        self._source = open(path, "rb")
        self._position = 0
        self._cache = OrderedDict()
        self._cache_chunks = max(1, int(cache_chunks))
        try:
            self.header, self._authenticated_header = self._read_header()
            self._payload_offset = len(self._authenticated_header)
            self._nonce_prefix = bytes.fromhex(self.header["noncePrefix"])
            self._chunk_count = (
                self.header["plaintextLength"] + self.header["chunkSize"] - 1
            ) // self.header["chunkSize"]
            expected_size = (
                self._payload_offset
                + self.header["plaintextLength"]
                + self._chunk_count * AUTH_TAG_BYTES
            )
            actual_size = os.fstat(self._source.fileno()).st_size
            if actual_size != expected_size:
                raise ProtectedPackageError(
                    "Protected package payload length does not match its header"
                )
            if not callable(key_resolver):
                raise ProtectedPackageError(
                    "This .3dripper package requires a signed-in 3DRipper account"
                )
            key = key_resolver(dict(self.header))
            if not isinstance(key, (bytes, bytearray, memoryview)) or len(key) != 32:
                raise ProtectedPackageError(
                    "3DRipper returned an invalid protected package key"
                )
            try:
                from cryptography.exceptions import InvalidTag
                from cryptography.hazmat.primitives.ciphers.aead import AESGCM
            except ImportError as error:
                raise ProtectedPackageError(
                    "The Add-on encryption runtime is unavailable; reinstall the official release"
                ) from error
            self._aesgcm = AESGCM(bytes(key))
            self._invalid_tag = InvalidTag
        except BaseException:
            self._source.close()
            raise

    def readable(self):
        return True

    def seekable(self):
        return True

    def tell(self):
        return self._position

    def seek(self, offset, whence=io.SEEK_SET):
        self._check_closed()
        if whence == io.SEEK_SET:
            position = offset
        elif whence == io.SEEK_CUR:
            position = self._position + offset
        elif whence == io.SEEK_END:
            position = self.header["plaintextLength"] + offset
        else:
            raise ValueError("Unsupported seek mode")
        if position < 0:
            raise ValueError("Cannot seek before the start of a protected package")
        self._position = min(position, self.header["plaintextLength"])
        return self._position

    def read(self, size=-1):
        self._check_closed()
        remaining = self.header["plaintextLength"] - self._position
        if size is None or size < 0:
            size = remaining
        else:
            size = min(size, remaining)
        if size <= 0:
            return b""

        output = bytearray()
        end = self._position + size
        while self._position < end:
            chunk_index = self._position // self.header["chunkSize"]
            chunk_offset = self._position % self.header["chunkSize"]
            chunk = self._read_chunk(chunk_index)
            take = min(end - self._position, len(chunk) - chunk_offset)
            output.extend(chunk[chunk_offset:chunk_offset + take])
            self._position += take
        return bytes(output)

    def readinto(self, target):
        data = self.read(len(target))
        target[:len(data)] = data
        return len(data)

    def close(self):
        if not self.closed:
            self._cache.clear()
            self._source.close()
        super().close()

    def _read_header(self):
        magic = self._source.read(len(MAGIC))
        if magic != MAGIC:
            raise ProtectedPackageError("Invalid protected package magic")
        raw_length = self._source.read(4)
        if len(raw_length) != 4:
            raise ProtectedPackageError("Protected package header is truncated")
        header_length = struct.unpack(">I", raw_length)[0]
        if header_length <= 0 or header_length > MAX_HEADER_BYTES:
            raise ProtectedPackageError("Protected package header length is invalid")
        header_bytes = self._source.read(header_length)
        if len(header_bytes) != header_length:
            raise ProtectedPackageError("Protected package header is truncated")
        try:
            header = json.loads(header_bytes.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as error:
            raise ProtectedPackageError(
                "Protected package header is not valid UTF-8 JSON"
            ) from error
        self._validate_header(header)
        return header, magic + raw_length + header_bytes

    def _validate_header(self, header):
        if not isinstance(header, dict):
            raise ProtectedPackageError("Protected package header must be an object")
        if header.get("format") != FORMAT:
            raise ProtectedPackageError("Unsupported protected package format")
        if header.get("version") != VERSION:
            raise ProtectedPackageError("Unsupported protected package version")
        if header.get("algorithm") != ALGORITHM:
            raise ProtectedPackageError(
                "Unsupported protected package encryption algorithm"
            )
        if header.get("keyVersion") != KEY_VERSION:
            raise ProtectedPackageError("Unsupported protected package key version")
        if not isinstance(header.get("keyId"), str) or not KEY_ID_PATTERN.fullmatch(
            header["keyId"]
        ):
            raise ProtectedPackageError("Protected package key ID is invalid")
        if (
            not isinstance(header.get("chunkSize"), int)
            or isinstance(header["chunkSize"], bool)
            or not MIN_CHUNK_SIZE <= header["chunkSize"] <= MAX_CHUNK_SIZE
        ):
            raise ProtectedPackageError("Protected package chunk size is invalid")
        if (
            not isinstance(header.get("plaintextLength"), int)
            or isinstance(header["plaintextLength"], bool)
            or not 0 < header["plaintextLength"] <= MAX_PLAINTEXT_BYTES
        ):
            raise ProtectedPackageError("Protected package plaintext length is invalid")
        nonce_prefix = header.get("noncePrefix")
        if (
            not isinstance(nonce_prefix, str)
            or not re.fullmatch(r"[0-9a-fA-F]{16}", nonce_prefix)
        ):
            raise ProtectedPackageError("Protected package nonce prefix is invalid")

    def _read_chunk(self, chunk_index):
        cached = self._cache.get(chunk_index)
        if cached is not None:
            self._cache.move_to_end(chunk_index)
            return cached

        chunk_size = self.header["chunkSize"]
        plaintext_offset = chunk_index * chunk_size
        plaintext_length = min(
            chunk_size, self.header["plaintextLength"] - plaintext_offset
        )
        encrypted_length = plaintext_length + AUTH_TAG_BYTES
        encrypted_offset = self._payload_offset + chunk_index * (
            chunk_size + AUTH_TAG_BYTES
        )
        self._source.seek(encrypted_offset)
        encrypted = self._source.read(encrypted_length)
        if len(encrypted) != encrypted_length:
            raise ProtectedPackageError("Protected package chunk is truncated")

        nonce = self._nonce_prefix + struct.pack(">I", chunk_index)
        additional_data = self._authenticated_header + struct.pack(">I", chunk_index)
        try:
            plaintext = self._aesgcm.decrypt(nonce, encrypted, additional_data)
        except self._invalid_tag as error:
            raise ProtectedPackageError(
                "Protected package authentication failed; sign in with its owner account "
                "and verify that the file was not modified"
            ) from error
        self._cache[chunk_index] = plaintext
        self._cache.move_to_end(chunk_index)
        while len(self._cache) > self._cache_chunks:
            self._cache.popitem(last=False)
        return plaintext

    def _check_closed(self):
        if self.closed:
            raise ValueError("I/O operation on closed protected package")
