# SPDX-License-Identifier: Apache-2.0
"""
vLLM Ascend Model Security — vLLM 0.17.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__)


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.
    """
    try:
        from vllm.model_executor.model_loader import register_model_loader
        from vllm.model_security_ascend.encrypted_loader import EncryptedModelLoader

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


# ---------------------------------------------------------------------------
# Public API - Re-export from lightr_vllm_core
# ---------------------------------------------------------------------------

from lightr_vllm_core import (
    # Crypto
    ModelEncryptor,
    ModelDecryptor,
    generate_dek,
    is_encrypted_file,
    # License
    LicenseManager,
    LicenseValidator,
    LicenseInfo,
    OfflineLicenseManager,
    # Base
    FingerprintPolicy,
    BaseFingerprint,
    BaseEncryptedLoader,
    # Constants
    MAGIC_BYTES,
    HEADER_SIZE,
    CHUNK_SIZE,
    ENC_SUFFIX,
    SAFETENSORS_ENC,
    GGUF_ENC,
)

# Platform-specific (NPU)
from vllm.model_security_ascend.fingerprint import (
    get_machine_fingerprint,
    get_npu_uuids,
    get_fingerprint_info,
)

__all__ = [
    # Loader
    "EncryptedModelLoader",
    # Crypto
    "ModelEncryptor",
    "ModelDecryptor",
    "generate_dek",
    "is_encrypted_file",
    "MAGIC_BYTES",
    "HEADER_SIZE",
    "CHUNK_SIZE",
    # License
    "LicenseManager",
    "LicenseValidator",
    "LicenseInfo",
    "OfflineLicenseManager",
    # Fingerprint
    "FingerprintPolicy",
    "get_machine_fingerprint",
    "get_npu_uuids",
    "get_fingerprint_info",
    # Base
    "BaseFingerprint",
    "BaseEncryptedLoader",
    # Constants
    "ENC_SUFFIX",
    "SAFETENSORS_ENC",
    "GGUF_ENC",
]