"""Authenticated envelopes for retained private dashboard state."""

from __future__ import annotations

import base64
import binascii
import json
import os
import re
from collections.abc import Mapping

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


class PrivateStateEncryptionError(RuntimeError):
    """Private retained state could not be authenticated or decrypted."""


def decode_private_state_encryption_key(encoded: str) -> bytes:
    """Decode an independently injected 256-bit private-state key."""
    try:
        key = base64.b64decode(encoded, validate=True)
    except (binascii.Error, ValueError) as error:
        raise PrivateStateEncryptionError(
            "private state encryption key is invalid"
        ) from error
    if len(key) != 32:
        raise PrivateStateEncryptionError(
            "private state encryption key must decode to exactly 32 bytes"
        )
    return key


def private_state_encryption_key() -> bytes:
    """Load the legacy single key setting."""
    return decode_private_state_encryption_key(
        os.getenv("STACKCHAIN_PRIVATE_STATE_ENCRYPTION_KEY", "")
    )


_KEY_ID = re.compile(r"^[A-Za-z0-9_-]{1,32}$")
_MAX_KEYS = 4


def decode_private_state_encryption_keyring(
    encoded: str, active_key_id: str
) -> tuple[dict[str, bytes], str]:
    """Decode a bounded key-ID-to-key JSON object and validate its active key."""
    try:
        raw = json.loads(encoded)
    except (TypeError, json.JSONDecodeError) as error:
        raise PrivateStateEncryptionError(
            "private state encryption keyring is invalid"
        ) from error
    if not isinstance(raw, dict) or not raw:
        raise PrivateStateEncryptionError("private state encryption keyring is invalid")
    if len(raw) > _MAX_KEYS:
        raise PrivateStateEncryptionError(
            f"private state encryption keyring supports at most {_MAX_KEYS} keys"
        )
    keys: dict[str, bytes] = {}
    for key_id, value in raw.items():
        if not isinstance(key_id, str) or not _KEY_ID.fullmatch(key_id):
            raise PrivateStateEncryptionError("private state encryption key ID is invalid")
        if not isinstance(value, str):
            raise PrivateStateEncryptionError("private state encryption keyring is invalid")
        keys[key_id] = decode_private_state_encryption_key(value)
    if active_key_id not in keys:
        raise PrivateStateEncryptionError(
            "private state active key is missing from the keyring"
        )
    return keys, active_key_id


def private_state_encryption_config() -> bytes | tuple[dict[str, bytes], str]:
    """Load a rotation keyring when configured, otherwise the legacy single key."""
    encoded = os.getenv("STACKCHAIN_PRIVATE_STATE_ENCRYPTION_KEYS", "")
    active = os.getenv("STACKCHAIN_PRIVATE_STATE_ACTIVE_KEY_ID", "")
    if encoded or active:
        return decode_private_state_encryption_keyring(encoded, active)
    return private_state_encryption_key()


class PrivateStateCipher:
    """Seal JSON values with store-specific authenticated context."""

    def __init__(
        self,
        key: bytes | Mapping[str, bytes] | tuple[dict[str, bytes], str],
        *,
        store: str,
        active_key_id: str | None = None,
    ):
        if isinstance(key, tuple):
            key, active_key_id = key
        if isinstance(key, Mapping):
            if not key or len(key) > _MAX_KEYS or active_key_id not in key:
                raise PrivateStateEncryptionError(
                    "private state encryption keyring configuration is invalid"
                )
            ciphers: dict[str, AESGCM] = {}
            for key_id, raw_key in key.items():
                if not isinstance(key_id, str) or not _KEY_ID.fullmatch(key_id):
                    raise PrivateStateEncryptionError(
                        "private state encryption key ID is invalid"
                    )
                if not isinstance(raw_key, bytes) or len(raw_key) != 32:
                    raise PrivateStateEncryptionError(
                        "private state encryption requires exactly 32 key bytes"
                    )
                ciphers[key_id] = AESGCM(raw_key)
            self._ciphers = ciphers
            self._active_key_id = active_key_id
            self._legacy_cipher = None
        else:
            if not isinstance(key, bytes) or len(key) != 32:
                raise PrivateStateEncryptionError(
                    "private state encryption requires exactly 32 key bytes"
                )
            self._ciphers = {}
            self._active_key_id = None
            self._legacy_cipher = AESGCM(key)
        if not store or "\0" in store:
            raise ValueError("private state store identity is invalid")
        self._store = store

    def _aad(self, binding: str, *, version: str = "v1", key_id: str = "") -> bytes:
        prefix = f"stackchain:private-state:{version}\0{self._store}\0{binding}"
        if version == "v2":
            prefix += f"\0{key_id}"
        return prefix.encode()

    def seal(self, value: object, *, binding: str = "singleton") -> str:
        plaintext = json.dumps(value, separators=(",", ":")).encode()
        nonce = os.urandom(12)
        if self._active_key_id is not None:
            key_id = self._active_key_id
            sealed = nonce + self._ciphers[key_id].encrypt(
                nonce, plaintext, self._aad(binding, version="v2", key_id=key_id)
            )
            return f"v2:{key_id}:" + base64.urlsafe_b64encode(sealed).decode()
        sealed = nonce + self._legacy_cipher.encrypt(
            nonce, plaintext, self._aad(binding)
        )
        return "v1:" + base64.urlsafe_b64encode(sealed).decode()

    @staticmethod
    def _decode_sealed(payload: str) -> bytes:
        sealed = base64.b64decode(payload, altchars=b"-_", validate=True)
        if len(sealed) < 28:
            raise ValueError("encrypted payload is too short")
        return sealed

    def open(self, payload: str, *, binding: str = "singleton") -> tuple[object, bool]:
        """Return the decoded value and whether active-key rewrapping is required."""
        if not isinstance(payload, str):
            raise PrivateStateEncryptionError("private state could not be decrypted")
        if not payload.startswith(("v1:", "v2:")):
            try:
                return json.loads(payload), True
            except (TypeError, json.JSONDecodeError) as error:
                raise PrivateStateEncryptionError(
                    "private state could not be decrypted"
                ) from error
        try:
            if payload.startswith("v2:"):
                _, key_id, encoded = payload.split(":", 2)
                cipher = self._ciphers.get(key_id)
                if cipher is None:
                    raise ValueError("encrypted payload uses an unavailable key")
                sealed = self._decode_sealed(encoded)
                plaintext = cipher.decrypt(
                    sealed[:12],
                    sealed[12:],
                    self._aad(binding, version="v2", key_id=key_id),
                )
                return json.loads(plaintext), key_id != self._active_key_id

            sealed = self._decode_sealed(payload[3:])
            candidates = (
                [self._legacy_cipher]
                if self._legacy_cipher is not None
                else list(self._ciphers.values())
            )
            plaintext = None
            for cipher in candidates:
                try:
                    plaintext = cipher.decrypt(
                        sealed[:12], sealed[12:], self._aad(binding)
                    )
                    break
                except InvalidTag:
                    continue
            if plaintext is None:
                raise InvalidTag
            return json.loads(plaintext), self._active_key_id is not None
        except (
            binascii.Error,
            InvalidTag,
            UnicodeDecodeError,
            ValueError,
            json.JSONDecodeError,
        ) as error:
            raise PrivateStateEncryptionError(
                "private state could not be decrypted"
            ) from error
