from __future__ import annotations

import io
import json
import os
import struct
import tempfile
from dataclasses import dataclass
from typing import Any, Protocol, runtime_checkable

import zstandard as zstd
from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

from ._exceptions import VaultError, VaultFormatError, VaultIntegrityError, VaultSerializerError, VaultConfigurationError


MAGIC_HEADER = b"NCDEV"
FORMAT_VERSION = 1

DEFAULT_EXTENSION = ".nc"
DEFAULT_CHUNK_SIZE = 64 * 1024
MAX_CHUNK_CEILING = 10 * 1024 * 1024
MAX_SERIALIZER_NAME = 64

FLAG_COMPRESSED = 0b00000001

KDF_KEY_HKDF = "hkdf-key"
KDF_PASSWORD_SCRYPT = "scrypt-password"


@runtime_checkable
class Serializer(Protocol):
    name: str

    def encode(self, obj: Any) -> bytes:
        ...

    def decode(self, raw: bytes) -> Any:
        ...


class BytesSerializer:
    name = "bytes"

    def encode(self, obj: Any) -> bytes:
        if isinstance(obj, bytes):
            return obj
        if isinstance(obj, bytearray):
            return bytes(obj)
        if isinstance(obj, memoryview):
            return obj.tobytes()
        raise TypeError("BytesSerializer expects bytes-like input")

    def decode(self, raw: bytes) -> bytes:
        return raw


class TextSerializer:
    name = "text"

    def encode(self, obj: Any) -> bytes:
        if not isinstance(obj, str):
            raise TypeError("TextSerializer expects str input")
        return obj.encode("utf-8")

    def decode(self, raw: bytes) -> str:
        try:
            return raw.decode("utf-8")
        except UnicodeDecodeError as e:
            raise ValueError(f"Failed to decode UTF-8 text payload: {e}") from e


class JsonSerializer:
    name = "json"

    def encode(self, obj: Any) -> bytes:
        try:
            return json.dumps(
                obj,
                ensure_ascii=False,
                separators=(",", ":"),
                sort_keys=True,
            ).encode("utf-8")
        except Exception as e:
            raise TypeError(f"Failed to JSON-encode object: {e}") from e

    def decode(self, raw: bytes) -> Any:
        try:
            return json.loads(raw.decode("utf-8"))
        except Exception as e:
            raise ValueError(f"Failed to JSON-decode payload: {e}") from e


class FlatBufferSerializer:
    def __init__(self, name: str, encode_func, decode_func):
        self.name = name
        self._encode_func = encode_func
        self._decode_func = decode_func

    def encode(self, obj: Any) -> bytes:
        raw = self._encode_func(obj)
        if not isinstance(raw, (bytes, bytearray, memoryview)):
            raise TypeError("FlatBuffer encode_func must return bytes-like data")
        return bytes(raw)

    def decode(self, raw: bytes) -> Any:
        return self._decode_func(raw)


@dataclass(slots=True)
class VaultOptions:
    compress: bool = True
    chunk_size: int = DEFAULT_CHUNK_SIZE
    extension: str = DEFAULT_EXTENSION
    compression_level: int = 3
    max_chunk_ceiling: int = MAX_CHUNK_CEILING

    def validate(self) -> None:
        if not isinstance(self.compress, bool):
            raise TypeError("options.compress must be bool")
        if not isinstance(self.chunk_size, int) or self.chunk_size <= 0:
            raise ValueError("options.chunk_size must be a positive int")
        if self.chunk_size > self.max_chunk_ceiling:
            raise ValueError(
                f"options.chunk_size cannot exceed max_chunk_ceiling ({self.max_chunk_ceiling})"
            )
        if not isinstance(self.extension, str) or not self.extension:
            raise ValueError("options.extension must be a non-empty string")
        if not isinstance(self.compression_level, int):
            raise TypeError("options.compression_level must be int")
        if not isinstance(self.max_chunk_ceiling, int) or self.max_chunk_ceiling <= 0:
            raise ValueError("options.max_chunk_ceiling must be a positive int")

class Enclave:
    def __init__(
        self,
        master_secret: bytes,
        *,
        mode: str,
        options: VaultOptions | None = None,
    ):
        if not isinstance(master_secret, (bytes, bytearray)) or not master_secret:
            raise TypeError("master_secret must be non-empty bytes-like data")
        if mode not in {KDF_KEY_HKDF, KDF_PASSWORD_SCRYPT}:
            raise ValueError(f"Unsupported mode: {mode}")

        self._master_secret = bytes(master_secret)
        self._mode = mode
        self.options = options or VaultOptions()
        self.options.validate()

        self._compressor = zstd.ZstdCompressor(level=self.options.compression_level)
        self._decompressor = zstd.ZstdDecompressor()
        self._serializers: dict[str, Serializer] = {}

        self.register_serializer(BytesSerializer())
        self.register_serializer(TextSerializer())
        self.register_serializer(JsonSerializer())

    @classmethod
    def from_key(
        cls,
        key: bytes | bytearray,
        *,
        options: VaultOptions | None = None,
    ) -> "Enclave":
        if not isinstance(key, (bytes, bytearray)) or not key:
            raise TypeError("key must be non-empty bytes-like data")
        return cls(bytes(key), mode=KDF_KEY_HKDF, options=options)

    @classmethod
    def from_password(
        cls,
        password: str,
        *,
        options: VaultOptions | None = None,
    ) -> "Enclave":
        if not isinstance(password, str) or not password:
            raise TypeError("password must be a non-empty str")
        return cls(password.encode("utf-8"), mode=KDF_PASSWORD_SCRYPT, options=options)

    @property
    def kdf_mode(self) -> str:
        return self._mode

    def register_serializer(self, serializer: Serializer) -> None:
        if not isinstance(serializer.name, str) or not serializer.name:
            raise ValueError("serializer.name must be a non-empty string")

        serializer_name_bytes = serializer.name.encode("utf-8")
        if len(serializer_name_bytes) > MAX_SERIALIZER_NAME:
            raise ValueError(f"serializer name too long (max {MAX_SERIALIZER_NAME} bytes)")

        self._serializers[serializer.name] = serializer

    def unregister_serializer(self, name: str) -> None:
        if name in {"bytes", "text", "json"}:
            raise VaultSerializerError(f"Cannot unregister built-in serializer: {name}")
        self._serializers.pop(name, None)

    def list_serializers(self) -> list[str]:
        return sorted(self._serializers.keys())

    def get_serializer(self, name: str) -> Serializer:
        try:
            return self._serializers[name]
        except KeyError as e:
            raise VaultSerializerError(f"Unknown serializer: {name}") from e

    def _derive_key(self, salt: bytes) -> bytes:
        if not isinstance(salt, bytes) or len(salt) != 16:
            raise VaultConfigurationError("salt must be exactly 16 bytes")

        try:
            if self._mode == KDF_KEY_HKDF:
                hkdf = HKDF(
                    algorithm=hashes.SHA256(),
                    length=32,
                    salt=salt,
                    info=b"Enclave/ChaCha20Poly1305/v4",
                )
                return hkdf.derive(self._master_secret)

            if self._mode == KDF_PASSWORD_SCRYPT:
                scrypt = Scrypt(
                    salt=salt,
                    length=32,
                    n=2**14,
                    r=8,
                    p=1,
                )
                return scrypt.derive(self._master_secret)

            raise VaultConfigurationError(f"Unsupported KDF mode: {self._mode}")
        except Exception as e:
            raise VaultError(f"Failed to derive encryption key: {e}") from e

    def _normalize_path(self, filename: str) -> str:
        if not isinstance(filename, str) or not filename.strip():
            raise ValueError("filename must be a non-empty string")

        root, ext = os.path.splitext(filename)
        if not ext:
            return filename + self.options.extension
        return filename

    def _build_header(self, serializer_name: str, compress: bool) -> tuple[bytes, bytes]:
        serializer_name_bytes = serializer_name.encode("utf-8")
        if len(serializer_name_bytes) > MAX_SERIALIZER_NAME:
            raise VaultFormatError("Serializer name too long")

        salt = os.urandom(16)
        flags = FLAG_COMPRESSED if compress else 0

        header = (
            MAGIC_HEADER
            + struct.pack("B", FORMAT_VERSION)
            + salt
            + struct.pack("B", flags)
            + struct.pack("B", len(serializer_name_bytes))
            + serializer_name_bytes
        )
        return header, salt

    def _parse_header(self, blob: bytes | bytearray | memoryview) -> dict[str, Any]:
        stream = io.BytesIO(bytes(blob))

        magic = stream.read(len(MAGIC_HEADER))
        if magic != MAGIC_HEADER:
            raise VaultFormatError("Invalid magic header")

        version_data = stream.read(1)
        if len(version_data) != 1:
            raise VaultFormatError("Missing format version")
        version = struct.unpack("B", version_data)[0]
        if version != FORMAT_VERSION:
            raise VaultFormatError(f"Unsupported vault version: {version}")

        salt = stream.read(16)
        if len(salt) != 16:
            raise VaultFormatError("Missing or truncated salt")

        flags_data = stream.read(1)
        if len(flags_data) != 1:
            raise VaultFormatError("Missing flags")
        flags = struct.unpack("B", flags_data)[0]
        compress = bool(flags & FLAG_COMPRESSED)

        serializer_len_data = stream.read(1)
        if len(serializer_len_data) != 1:
            raise VaultFormatError("Missing serializer length")
        serializer_len = struct.unpack("B", serializer_len_data)[0]

        if serializer_len > MAX_SERIALIZER_NAME:
            raise VaultFormatError("Serializer name exceeds maximum allowed length")

        serializer_name_bytes = stream.read(serializer_len)
        if len(serializer_name_bytes) != serializer_len:
            raise VaultFormatError("Truncated serializer name")

        try:
            serializer_name = serializer_name_bytes.decode("utf-8")
        except UnicodeDecodeError as e:
            raise VaultFormatError("Serializer name is not valid UTF-8") from e

        header_len = stream.tell()
        header_bytes = bytes(blob[:header_len])

        return {
            "magic": magic,
            "version": version,
            "salt": salt,
            "flags": flags,
            "compressed": compress,
            "serializer": serializer_name,
            "header_len": header_len,
            "header_bytes": header_bytes,
        }

    @staticmethod
    def _normalize_aad_context(aad_context: bytes | bytearray | memoryview | None) -> bytes:
        if aad_context is None:
            return b""
        if isinstance(aad_context, bytes):
            return aad_context
        if isinstance(aad_context, bytearray):
            return bytes(aad_context)
        if isinstance(aad_context, memoryview):
            return aad_context.tobytes()
        raise TypeError("aad_context must be bytes-like or None")

    @staticmethod
    def _chunk_nonce(chunk_idx: int) -> bytes:
        if chunk_idx < 0:
            raise ValueError("chunk_idx cannot be negative")
        return struct.pack(">Q", chunk_idx) + struct.pack(">I", 0)

    def _chunk_aad(self, header_bytes: bytes, chunk_idx: int, aad_context: bytes) -> bytes:
        return header_bytes + struct.pack(">I", chunk_idx) + aad_context

    def inspect(self, blob: bytes | bytearray | memoryview) -> dict[str, Any]:
        meta = self._parse_header(blob)
        return {
            "magic": meta["magic"].decode("ascii", errors="replace"),
            "version": meta["version"],
            "serializer": meta["serializer"],
            "compressed": meta["compressed"],
            "kdf_mode": self._mode,
            "header_len": meta["header_len"],
        }

    def inspect_file(self, filename: str) -> dict[str, Any]:
        filename = self._normalize_path(filename)
        if not os.path.exists(filename):
            raise FileNotFoundError(filename)

        with open(filename, "rb") as f:
            data = f.read(256)
        return self.inspect(data)

    def encrypt_bytes(
        self,
        raw: bytes | bytearray | memoryview,
        *,
        serializer_name: str = "bytes",
        compress: bool | None = None,
        aad_context: bytes | bytearray | memoryview | None = None,
    ) -> bytes:
        if not isinstance(raw, (bytes, bytearray, memoryview)):
            raise TypeError("raw must be bytes-like")

        if serializer_name not in self._serializers:
            raise VaultSerializerError(f"Unknown serializer: {serializer_name}")

        raw = bytes(raw)
        compress = self.options.compress if compress is None else compress
        aad_context_bytes = self._normalize_aad_context(aad_context)

        header, salt = self._build_header(serializer_name=serializer_name, compress=compress)
        key = self._derive_key(salt)
        cipher = ChaCha20Poly1305(key)

        out = io.BytesIO()
        out.write(header)

        chunk_size = self.options.chunk_size
        max_chunk_ceiling = self.options.max_chunk_ceiling

        chunk_idx = 0
        for i in range(0, len(raw), chunk_size):
            chunk = raw[i : i + chunk_size]
            payload_chunk = self._compressor.compress(chunk) if compress else chunk
            if len(payload_chunk) > max_chunk_ceiling:
                raise VaultFormatError(
                    f"Chunk {chunk_idx} exceeds max_chunk_ceiling ({max_chunk_ceiling}) before encryption"
                )

            nonce = self._chunk_nonce(chunk_idx)
            aad = self._chunk_aad(header, chunk_idx, aad_context_bytes)
            encrypted_chunk = cipher.encrypt(nonce, payload_chunk, aad)

            if len(encrypted_chunk) > max_chunk_ceiling:
                raise VaultFormatError(
                    f"Encrypted chunk {chunk_idx} exceeds max_chunk_ceiling ({max_chunk_ceiling})"
                )

            out.write(struct.pack(">I", len(encrypted_chunk)))
            out.write(nonce)
            out.write(encrypted_chunk)

            chunk_idx += 1

        return out.getvalue()

    def decrypt_bytes(
        self,
        blob: bytes | bytearray | memoryview,
        *,
        aad_context: bytes | bytearray | memoryview | None = None,
    ) -> tuple[bytes, str]:
        if not isinstance(blob, (bytes, bytearray, memoryview)):
            raise TypeError("blob must be bytes-like")

        blob = bytes(blob)
        if len(blob) < len(MAGIC_HEADER) + 1 + 16 + 1 + 1:
            raise VaultFormatError("Blob too short to contain a valid vault header")

        aad_context_bytes = self._normalize_aad_context(aad_context)
        meta = self._parse_header(blob)

        key = self._derive_key(meta["salt"])
        cipher = ChaCha20Poly1305(key)

        stream = io.BytesIO(blob)
        stream.seek(meta["header_len"])
        out = bytearray()

        chunk_idx = 0
        while True:
            size_data = stream.read(4)
            if not size_data:
                break
            if len(size_data) != 4:
                raise VaultFormatError("Truncated chunk length")

            enc_size = struct.unpack(">I", size_data)[0]
            if enc_size <= 0 or enc_size > self.options.max_chunk_ceiling:
                raise VaultFormatError(f"Invalid encrypted chunk length: {enc_size}")

            nonce = stream.read(12)
            if len(nonce) != 12:
                raise VaultFormatError("Truncated nonce")

            expected_nonce = self._chunk_nonce(chunk_idx)
            if nonce != expected_nonce:
                raise VaultIntegrityError(
                    f"Unexpected nonce for chunk {chunk_idx}; file may be tampered with"
                )

            encrypted_chunk = stream.read(enc_size)
            if len(encrypted_chunk) != enc_size:
                raise VaultFormatError("Truncated encrypted chunk")

            aad = self._chunk_aad(meta["header_bytes"], chunk_idx, aad_context_bytes)

            try:
                chunk = cipher.decrypt(nonce, encrypted_chunk, aad)
            except InvalidTag as e:
                raise VaultIntegrityError(
                    f"Authentication failed for chunk {chunk_idx}"
                ) from e
            except Exception as e:
                raise VaultError(f"Decryption failed for chunk {chunk_idx}: {e}") from e

            if meta["compressed"]:
                try:
                    chunk = self._decompressor.decompress(chunk)
                except zstd.ZstdError as e:
                    raise VaultFormatError(
                        f"Decompression failed for chunk {chunk_idx}: {e}"
                    ) from e

            out.extend(chunk)
            chunk_idx += 1

        return bytes(out), meta["serializer"]

    def pack(
        self,
        obj: Any,
        *,
        serializer: str | Serializer = "json",
        compress: bool | None = None,
        aad_context: bytes | bytearray | memoryview | None = None,
    ) -> bytes:
        ser = serializer if not isinstance(serializer, str) else self.get_serializer(serializer)

        try:
            raw = ser.encode(obj)
        except Exception as e:
            raise VaultSerializerError(f"Failed to encode object with {ser.name}: {e}") from e

        return self.encrypt_bytes(
            raw,
            serializer_name=ser.name,
            compress=compress,
            aad_context=aad_context,
        )

    def unpack(
        self,
        blob: bytes | bytearray | memoryview,
        *,
        serializer: str | Serializer | None = None,
        aad_context: bytes | bytearray | memoryview | None = None,
    ) -> Any:
        raw, embedded_serializer_name = self.decrypt_bytes(blob, aad_context=aad_context)

        if serializer is None:
            ser = self.get_serializer(embedded_serializer_name)
        elif isinstance(serializer, str):
            ser = self.get_serializer(serializer)
        else:
            ser = serializer

        try:
            return ser.decode(raw)
        except Exception as e:
            raise VaultSerializerError(f"Failed to decode payload with {ser.name}: {e}") from e

    def _atomic_write(self, filename: str, data: bytes) -> None:
        directory = os.path.dirname(os.path.abspath(filename)) or "."
        os.makedirs(directory, exist_ok=True)

        temp_name: str | None = None
        try:
            with tempfile.NamedTemporaryFile(dir=directory, delete=False) as tmp:
                temp_name = tmp.name
                tmp.write(data)
                tmp.flush()
                os.fsync(tmp.fileno())
            os.replace(temp_name, filename)
        finally:
            if temp_name and os.path.exists(temp_name):
                try:
                    os.remove(temp_name)
                except OSError:
                    pass

    def pack_to_file(
        self,
        filename: str,
        obj: Any,
        *,
        serializer: str | Serializer = "json",
        compress: bool | None = None,
        aad_context: bytes | bytearray | memoryview | None = None,
        atomic: bool = True,
    ) -> str:
        filename = self._normalize_path(filename)
        blob = self.pack(
            obj,
            serializer=serializer,
            compress=compress,
            aad_context=aad_context,
        )

        if atomic:
            self._atomic_write(filename, blob)
        else:
            with open(filename, "wb") as f:
                f.write(blob)

        return filename

    def unpack_from_file(
        self,
        filename: str,
        *,
        serializer: str | Serializer | None = None,
        aad_context: bytes | bytearray | memoryview | None = None,
    ) -> Any:
        filename = self._normalize_path(filename)
        if not os.path.exists(filename):
            raise FileNotFoundError(filename)

        with open(filename, "rb") as f:
            blob = f.read()

        return self.unpack(blob, serializer=serializer, aad_context=aad_context)

    def read_raw_from_file(
        self,
        filename: str,
        *,
        aad_context: bytes | bytearray | memoryview | None = None,
    ) -> tuple[bytes, str]:
        filename = self._normalize_path(filename)
        if not os.path.exists(filename):
            raise FileNotFoundError(filename)

        with open(filename, "rb") as f:
            blob = f.read()

        return self.decrypt_bytes(blob, aad_context=aad_context)