"""
Basilisk OS — VetKeyService: on-chain encryption via ICP vetKeys.

Provides a high-level async API for deriving vetKeys from the IC
management canister's vetKD (Verifiably Encrypted Threshold Key
Derivation) protocol.

All inter-canister operations (public_key, derive_key) are async
generators that must be driven with ``yield``::

    from basilisk.os.vetkeys import VetKeyService

    vks = VetKeyService()

    # In an @update endpoint:
    pub_key = yield vks.public_key()
    encrypted_key = yield vks.derive_key(transport_public_key=tpk_bytes)

Key concepts:

  - **context**: Binds the derived key to the calling canister + a
    domain separator + an application-specific scope (e.g. caller
    principal).  Two different contexts always produce different keys.

  - **input**: An additional diversifier within the same context.
    Allows a single user to derive multiple independent keys (e.g.
    one per document or per field).

  - **transport_public_key**: An ephemeral public key generated by
    the client.  The management canister encrypts the derived key
    under this transport key so that no single subnet node ever sees
    the plaintext key.  The client decrypts it locally.

Available master keys on ICP:

  - ``key_1``          Production (34-node fiduciary subnet pzp6e)
  - ``test_key_1``     Test (13-node subnet fuqsr)
  - ``dfx_test_key``   Local dfx development only
"""

from basilisk import Async, Principal, ic
from basilisk.canisters.management import (
    management_canister,
    VetKDCurve,
    VetKDKeyId,
    VetKDPublicKeyArgs,
    VetKDDeriveKeyArgs,
)
try:
    from ic_python_logging import get_logger
except ImportError:
    import logging
    get_logger = logging.getLogger

logger = get_logger("basilisk.os.vetkeys")

# Well-known key names
VETKD_KEY_PRODUCTION = "key_1"
VETKD_KEY_TEST = "test_key_1"
VETKD_KEY_LOCAL = "dfx_test_key"


class VetKeyService:
    """
    High-level vetKeys service for Basilisk OS canisters.

    Wraps the IC management canister's ``vetkd_public_key`` and
    ``vetkd_derive_key`` endpoints with ergonomic helpers for
    context construction and key selection.

    Args:
        domain_separator: Application-specific domain separator bytes.
            Prevents key collisions between different applications
            using the same canister.  Defaults to ``b"basilisk"``.
        key_name: Which master key to use.  Defaults to
            ``"test_key_1"`` (safe for development).  Use ``"key_1"``
            for production deployments.
    """

    def __init__(self, domain_separator=b"basilisk", key_name=VETKD_KEY_TEST):
        self._domain_separator = (
            domain_separator
            if isinstance(domain_separator, bytes)
            else domain_separator.encode("utf-8")
        )
        self._key_name = key_name

    # ------------------------------------------------------------------
    # Key ID helper
    # ------------------------------------------------------------------

    def _key_id(self):
        """Return the VetKDKeyId for the configured master key."""
        return VetKDKeyId(
            curve=VetKDCurve(bls12_381_g2=None),
            name=self._key_name,
        )

    # ------------------------------------------------------------------
    # Context construction
    # ------------------------------------------------------------------

    def make_context(self, scope=None):
        """
        Build a vetKD context from the domain separator and an optional
        scope (typically the caller's principal bytes).

        The format follows the DFINITY convention::

            [len(domain_separator)] || domain_separator || scope

        Args:
            scope: Additional bytes to include in the context.
                Defaults to ``ic.caller()`` principal bytes.

        Returns:
            ``bytes`` suitable for the ``context`` field of vetKD args.
        """
        if scope is None:
            scope = ic.caller().bytes
        elif isinstance(scope, str):
            scope = scope.encode("utf-8")
        ds = self._domain_separator
        return bytes([len(ds)]) + ds + scope

    # ------------------------------------------------------------------
    # Async: vetkd_public_key (inter-canister call)
    # ------------------------------------------------------------------

    def public_key(self, scope=None):
        """
        Retrieve the derived vetKD public key for the given context.

        Must be called with ``yield``::

            pub_key_bytes = yield vks.public_key()
            pub_key_bytes = yield vks.public_key(scope=b"custom-scope")

        Args:
            scope: Optional scope bytes for context construction.
                Defaults to ``ic.caller()`` principal.

        Returns:
            Generator yielding an inter-canister call; final value
            is the public key as ``bytes``.
        """
        return self._public_key(scope)

    def _public_key(self, scope=None) -> Async[bytes]:
        context = self.make_context(scope)
        args = VetKDPublicKeyArgs(
            canister_id=None,
            context=context,
            key_id=self._key_id(),
        )
        logger.info(
            f"vetkd_public_key: key={self._key_name} "
            f"context_len={len(context)}"
        )
        result = yield management_canister.vetkd_public_key(args)
        pub_key = result.public_key
        logger.info(f"vetkd_public_key: received {len(pub_key)} bytes")
        return pub_key

    # ------------------------------------------------------------------
    # Async: vetkd_derive_key (inter-canister call)
    # ------------------------------------------------------------------

    def derive_key(self, transport_public_key, input=b"", scope=None):
        """
        Derive an encrypted vetKey for the given context.

        The returned key is encrypted under ``transport_public_key``
        and can only be decrypted by the holder of the corresponding
        transport secret key (typically the frontend/client).

        Must be called with ``yield``::

            encrypted_key = yield vks.derive_key(
                transport_public_key=tpk_bytes,
            )
            encrypted_key = yield vks.derive_key(
                transport_public_key=tpk_bytes,
                input=b"document-123",
                scope=b"custom-scope",
            )

        Args:
            transport_public_key: The client's ephemeral transport
                public key (bytes).  Generated client-side via
                ``TransportSecretKey.random().publicKeyBytes()``.
            input: Additional diversifier bytes for deriving multiple
                independent keys within the same context.  Defaults
                to empty bytes (single key per context).
            scope: Optional scope bytes for context construction.
                Defaults to ``ic.caller()`` principal.

        Returns:
            Generator yielding an inter-canister call; final value
            is the encrypted key as ``bytes``.
        """
        return self._derive_key(transport_public_key, input, scope)

    def _derive_key(self, transport_public_key, input=b"", scope=None) -> Async[bytes]:
        if isinstance(input, str):
            input = input.encode("utf-8")
        if isinstance(transport_public_key, str):
            transport_public_key = bytes.fromhex(transport_public_key)

        context = self.make_context(scope)
        args = VetKDDeriveKeyArgs(
            input=input,
            context=context,
            key_id=self._key_id(),
            transport_public_key=transport_public_key,
        )
        logger.info(
            f"vetkd_derive_key: key={self._key_name} "
            f"context_len={len(context)} input_len={len(input)} "
            f"transport_pk_len={len(transport_public_key)}"
        )
        result = yield management_canister.vetkd_derive_key(args)
        encrypted_key = result.encrypted_key
        logger.info(f"vetkd_derive_key: received {len(encrypted_key)} bytes")
        return encrypted_key
