from __future__ import annotations

import base64
import secrets
from abc import ABC, abstractmethod


class SecretStore(ABC):
    """Interface for secure storage of opaque secrets (API keys, tokens)."""

    @abstractmethod
    def put(self, plaintext: str) -> str:
        """Persist a secret and return a reference handle."""

    @abstractmethod
    def get(self, handle: str) -> str:
        """Retrieve a previously stored secret by handle."""

    @abstractmethod
    def delete(self, handle: str) -> None:
        """Best-effort deletion. For immutable backends this may be a no-op."""


class InMemorySecretStore(SecretStore):
    """Simple in-memory store useful for unit tests."""

    def __init__(self) -> None:
        self._secrets: dict[str, str] = {}

    def put(self, plaintext: str) -> str:
        handle = secrets.token_urlsafe(16)
        self._secrets[handle] = plaintext
        return f"mem://{handle}"

    def get(self, handle: str) -> str:
        key = handle.split("mem://", 1)[-1]
        if key not in self._secrets:
            raise KeyError(f"Secret handle not found: {handle}")
        return self._secrets[key]

    def delete(self, handle: str) -> None:
        key = handle.split("mem://", 1)[-1]
        self._secrets.pop(key, None)


class DevKmsStore(SecretStore):
    """Local AES-GCM envelope encryption driven by gateway.secrets.master_key."""

    PREFIX = "devkms:v1:"

    def __init__(
        self,
        master_key: str | bytes | None = None,
        *,
        associated_data: bytes | None = None,
    ) -> None:
        self._associated_data = associated_data
        key_material = master_key or _config_value("gateway.secrets.master_key")
        if not key_material:
            raise ValueError("gateway.secrets.master_key is required for DevKmsStore")
        key_bytes = self._normalize_key(key_material)
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        except ImportError as exc:  # pragma: no cover
            raise RuntimeError(
                "cryptography package is required for DevKmsStore; install via `uv sync`"
            ) from exc
        self._aesgcm = AESGCM(key_bytes)

    @staticmethod
    def _normalize_key(key_material: str | bytes) -> bytes:
        if isinstance(key_material, bytes):
            data = key_material
        else:
            candidate = key_material.strip()
            # Try URL-safe base64 first
            try:
                data = base64.urlsafe_b64decode(candidate + "==")
            except Exception:
                try:
                    data = bytes.fromhex(candidate)
                except Exception as err:
                    raise ValueError("Secret key must be base64 or hex encoded") from err
        if len(data) not in {16, 24, 32}:
            raise ValueError("Secret key must be 128, 192, or 256 bits")
        return data

    def put(self, plaintext: str) -> str:
        nonce = secrets.token_bytes(12)
        ciphertext = self._aesgcm.encrypt(
            nonce,
            plaintext.encode("utf-8"),
            self._associated_data,
        )
        payload = nonce + ciphertext
        handle = base64.urlsafe_b64encode(payload).decode("ascii")
        return f"{self.PREFIX}{handle}"

    def get(self, handle: str) -> str:
        if not handle.startswith(self.PREFIX):
            raise ValueError("Invalid secret handle for DevKmsStore")
        b64 = handle[len(self.PREFIX) :]
        payload = base64.urlsafe_b64decode(b64.encode("ascii"))
        nonce, ciphertext = payload[:12], payload[12:]
        plaintext = self._aesgcm.decrypt(nonce, ciphertext, self._associated_data)
        return plaintext.decode("utf-8")

    def delete(self, handle: str) -> None:
        # Stateless envelope store; nothing persisted locally.
        return None


__all__ = ["SecretStore", "InMemorySecretStore", "DevKmsStore"]


def _config_value(path: str) -> str | None:
    from ember._internal.context.runtime import EmberContext

    return EmberContext.current().get_config(path)
