# SPDX-License-Identifier: Apache-2.0
"""
vLLM Model Security — v0.18.0 实现

使用官方 register_model_loader 插件机制注册加密模型加载器。
仅影响 --load-format encrypted，不影响其他加载格式。

用法:
    vllm serve ./encrypted_model \
        --load-format encrypted \
        --model-loader-extra-config '{"license_path": "./model.lic"}'
"""

from __future__ import annotations

import logging
from typing import Any

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Plugin entry point for vllm.general_plugins
# ---------------------------------------------------------------------------
# This function is called by vLLM's plugin system in ALL processes
# (main process, EngineCore process, and worker processes).
# It MUST be re-entrant (safe to call multiple times).
# ---------------------------------------------------------------------------

def register_encrypted_loader() -> Any:
    """
    Plugin entry point for vLLM's general_plugins group.

    Registers the encrypted model loader via the official
    register_model_loader mechanism.

    This function is called in all vLLM processes:
    - Main process (process 0)
    - EngineCore process (V1 engine)
    - Worker processes

    Returns:
        None (general plugins don't need to return anything)
    """
    try:
        from vllm.model_executor.model_loader import register_model_loader
        from vllm.model_security.encrypted_loader import EncryptedModelLoader

        register_model_loader("encrypted")(EncryptedModelLoader)
        logger.info(
            "vllm-model-security: registered 'encrypted' model loader "
            "via register_model_loader"
        )
    except ImportError as exc:
        logger.warning(
            "vllm-model-security: failed to register encrypted loader. "
            "Ensure vllm>=0.16.0 is installed. Error: %s", exc
        )
    except Exception as exc:
        logger.warning(
            "vllm-model-security: unexpected error during registration: %s", exc
        )
    return None

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

from vllm.model_security.crypto import (
    ModelEncryptor,
    ModelDecryptor,
    generate_dek,
)
from vllm.model_security.license import (
    LicenseManager,
    LicenseValidator,
    LicenseInfo,
)
from vllm.model_security.fingerprint import (
    get_machine_fingerprint,
    FingerprintPolicy,
)

__all__ = [
    # Loader
    "EncryptedModelLoader",
    # Crypto
    "ModelEncryptor",
    "ModelDecryptor",
    "generate_dek",
    # License
    "LicenseManager",
    "LicenseValidator",
    "LicenseInfo",
    # Fingerprint
    "get_machine_fingerprint",
    "FingerprintPolicy",
]