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

from __future__ import annotations

import base64
import binascii
import json
import os

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 the 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:
    return decode_private_state_encryption_key(
        os.getenv("STACKCHAIN_PRIVATE_STATE_ENCRYPTION_KEY", "")
    )


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

    def __init__(self, key: bytes, *, store: str):
        if not isinstance(key, bytes) or len(key) != 32:
            raise PrivateStateEncryptionError(
                "private state encryption requires exactly 32 key bytes"
            )
        if not store or "\0" in store:
            raise ValueError("private state store identity is invalid")
        self._cipher = AESGCM(key)
        self._store = store

    def _aad(self, binding: str) -> bytes:
        return f"stackchain:private-state:v1\0{self._store}\0{binding}".encode()

    def seal(self, value: object, *, binding: str = "singleton") -> str:
        plaintext = json.dumps(value, separators=(",", ":")).encode()
        nonce = os.urandom(12)
        sealed = nonce + self._cipher.encrypt(nonce, plaintext, self._aad(binding))
        return "v1:" + base64.urlsafe_b64encode(sealed).decode()

    def open(self, payload: str, *, binding: str = "singleton") -> tuple[object, bool]:
        """Return the decoded value and whether plaintext migration is required."""
        if not isinstance(payload, str):
            raise PrivateStateEncryptionError("private state could not be decrypted")
        if not payload.startswith("v1:"):
            try:
                return json.loads(payload), True
            except (TypeError, json.JSONDecodeError) as error:
                raise PrivateStateEncryptionError(
                    "private state could not be decrypted"
                ) from error
        try:
            sealed = base64.b64decode(payload[3:], altchars=b"-_", validate=True)
            if len(sealed) < 28:
                raise ValueError("encrypted payload is too short")
            plaintext = self._cipher.decrypt(
                sealed[:12], sealed[12:], self._aad(binding)
            )
            return json.loads(plaintext), False
        except (
            binascii.Error,
            InvalidTag,
            UnicodeDecodeError,
            ValueError,
            json.JSONDecodeError,
        ) as error:
            raise PrivateStateEncryptionError(
                "private state could not be decrypted"
            ) from error