# SPDX-License-Identifier: Apache-2.0
"""
Encrypted model loader for vLLM 0.18.0+

Integrates with vLLM's model loading pipeline to transparently decrypt
model weights at load time. Supports both safetensors and GGUF formats.

Uses the official register_model_loader plugin mechanism.
"""

import io
import os
import signal
import stat
import tempfile
import threading
from datetime import date
from typing import Generator, Optional, Tuple

import torch
from torch import nn

# vLLM 0.16.0 导入路径
from vllm.config import ModelConfig, VllmConfig
from vllm.config.load import LoadConfig
from vllm.logger import init_logger
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
from vllm.model_executor.model_loader.gguf_loader import GGUFModelLoader
from vllm.model_executor.model_loader.utils import (
    initialize_model,
    process_weights_after_loading,
)
from vllm.model_executor.model_loader.weight_utils import gguf_quant_weights_iterator
from vllm.utils.torch_utils import set_default_torch_dtype

from vllm.model_security.crypto import ModelDecryptor, MAGIC_BYTES
from vllm.model_security.license import LicenseValidator

logger = init_logger(__name__)

# 文件格式常量
ENC_SUFFIX = ".enc"
SAFETENSORS_ENC = ".safetensors.enc"
GGUF_ENC = ".gguf.enc"


# ---------------------------------------------------------------------------
# 工具函数
# ---------------------------------------------------------------------------

def is_encrypted_model(model_path: str) -> bool:
    """
    检查模型是否为加密模型

    Args:
        model_path: 模型目录或文件路径

    Returns:
        True 如果检测到加密文件
    """
    if os.path.isfile(model_path):
        return model_path.endswith(ENC_SUFFIX)

    if os.path.isdir(model_path):
        for fname in os.listdir(model_path):
            if fname.endswith(ENC_SUFFIX):
                return True

    return False


def _validate_enc_file(path: str) -> bool:
    """验证加密文件的 magic bytes"""
    try:
        with open(path, 'rb') as f:
            magic = f.read(len(MAGIC_BYTES))
            return magic == MAGIC_BYTES
    except Exception:
        return False


def _decrypt_to_buffer(enc_path: str, decryptor: ModelDecryptor) -> io.BytesIO:
    """
    将加密文件解密到内存缓冲区

    对于大模型，这是按 shard 调用的，内存使用受限于单个 shard 的大小。
    """
    buf = io.BytesIO()
    for chunk in decryptor.decrypt_file_stream(enc_path):
        buf.write(chunk)
    buf.seek(0)
    return buf


def _encrypted_safetensors_weights_iterator(
    enc_files: list,
    decryptor: ModelDecryptor,
) -> Generator[Tuple[str, torch.Tensor], None, None]:
    """
    加密 safetensors 权重迭代器

    逐 shard 解密到内存，然后读取张量。
    """
    from safetensors.torch import load as safetensors_load

    for enc_file in enc_files:
        logger.debug("Decrypting shard: %s", enc_file)
        buf = _decrypt_to_buffer(enc_file, decryptor)
        tensors = safetensors_load(buf.read())
        for name, tensor in tensors.items():
            yield name, tensor
        del buf, tensors


def _encrypted_gguf_weights_iterator(
    enc_file: str,
    decryptor: ModelDecryptor,
    gguf_to_hf_name_map: dict,
) -> Generator[Tuple[str, torch.Tensor], None, None]:
    """
    加密 GGUF 权重迭代器

    GGUF 需要文件路径，所以先解密到临时文件。
    """
    logger.debug("Decrypting GGUF file: %s", enc_file)

    tmp_fd = None
    tmp_path = None
    try:
        tmp_fd, tmp_path = tempfile.mkstemp(suffix='.gguf')
        # 设置限制性权限
        os.chmod(tmp_path, stat.S_IRUSR | stat.S_IWUSR)

        with os.fdopen(tmp_fd, 'wb') as tmp:
            tmp_fd = None  # fdopen 接管所有权
            for chunk in decryptor.decrypt_file_stream(enc_file):
                tmp.write(chunk)

        yield from gguf_quant_weights_iterator(tmp_path, gguf_to_hf_name_map)
    finally:
        if tmp_fd is not None:
            os.close(tmp_fd)
        if tmp_path is not None and os.path.exists(tmp_path):
            os.unlink(tmp_path)


# ---------------------------------------------------------------------------
# 加密模型加载器
# ---------------------------------------------------------------------------

class EncryptedModelLoader(BaseModelLoader):
    """
    加密模型加载器

    使用官方 register_model_loader 机制注册。
    支持:
    - 加密 safetensors (.safetensors.enc)
    - 加密 GGUF (.gguf.enc)
    - v0.16.0 load_weights() 独立接口
    - v0.16.0 load_config.device 配置

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

    def __init__(self, load_config: LoadConfig):
        super().__init__(load_config)
        extra = load_config.model_loader_extra_config or {}

        # 安全配置
        self.license_path: Optional[str] = extra.get("license_path")
        self.vendor_pubkey_path: Optional[str] = extra.get("vendor_pubkey_path")
        vendor_secret_hex = extra.get("vendor_secret") or os.environ.get("VLLM_VENDOR_SECRET")
        self.vendor_secret: Optional[bytes] = (
            bytes.fromhex(vendor_secret_hex) if vendor_secret_hex else None
        )

        # 状态
        self._decryptor: Optional[ModelDecryptor] = None
        self._license_info = None
        self._validator: Optional[LicenseValidator] = None
        self._monitor_thread: Optional[threading.Thread] = None

    # -----------------------------------------------------------------------
    # 许可证管理
    # -----------------------------------------------------------------------

    def _start_license_monitor(self) -> None:
        """启动许可证过期监控线程"""
        if self._monitor_thread is not None:
            return

        def _monitor_loop():
            import time
            while True:
                time.sleep(86400)  # 24 小时
                self._check_license_expiry()

        self._monitor_thread = threading.Thread(
            target=_monitor_loop,
            name="license-monitor",
            daemon=True,
        )
        self._monitor_thread.start()
        logger.info("License monitor started (checking every 24 hours)")

    def _check_license_expiry(self) -> None:
        """检查许可证过期状态"""
        if self._license_info is None or self._validator is None:
            return

        status = self._validator.check_expiry_with_grace(self._license_info, grace_days=3)
        expire_at = getattr(self._license_info, "expire_at", "N/A")

        if status == "valid":
            return

        if status == "warning":
            expire_date = date.fromisoformat(expire_at)
            days_over = (date.today() - expire_date).days
            days_left = 3 - days_over
            logger.warning("=" * 80)
            logger.warning("⚠️  LICENSE EXPIRED: License expired on %s", expire_at)
            logger.warning("   Grace period: %d day(s) remaining.", days_left)
            logger.warning("   Please contact your vendor for a license renewal.")
            logger.warning("=" * 80)
        elif status == "expired":
            logger.error("=" * 80)
            logger.error("🚫 LICENSE GRACE PERIOD ENDED: License expired on %s", expire_at)
            logger.error("   Service is shutting down. Please renew your license.")
            logger.error("=" * 80)
            os.kill(os.getpid(), signal.SIGTERM)

    def _get_decryptor(self, model_path: str) -> ModelDecryptor:
        """获取解密器（带许可证验证）"""
        if self._decryptor is not None:
            return self._decryptor

        # 解析路径
        license_path = self._resolve_path(self.license_path, model_path, "model.lic")
        pubkey_path = self._resolve_path(self.vendor_pubkey_path, model_path, "vendor_public.pem")

        if not license_path:
            raise ValueError(
                "License file not found. Provide via --model-loader-extra-config "
                '\'{"license_path": "/path/to/model.lic"}\' or place model.lic '
                "in the model directory."
            )
        if not pubkey_path:
            raise ValueError(
                "Vendor public key not found. Provide via --model-loader-extra-config "
                '\'{"vendor_pubkey_path": "/path/to/vendor_public.pem"}\' or place '
                "vendor_public.pem in the model directory."
            )

        logger.warning("=" * 80)
        logger.warning("🔐 MODEL SECURITY: Validating license: %s", license_path)
        logger.warning("=" * 80)

        validator = LicenseValidator(pubkey_path)
        license_info = validator.load_license(license_path)
        validator.validate(license_info)
        dek = validator.extract_dek_from_model_dir(
            license_info, model_path, self.vendor_secret
        )

        logger.warning("=" * 80)
        logger.warning("✅ LICENSE VALID - Model authorized for use")
        logger.warning("   License ID: %s", getattr(license_info, 'license_id', 'N/A'))
        logger.warning("   Expires: %s", getattr(license_info, 'expire_at', 'N/A'))
        logger.warning("=" * 80)

        self._license_info = license_info
        self._validator = validator
        self._decryptor = ModelDecryptor(dek)
        self._start_license_monitor()

        return self._decryptor

    def _resolve_path(self, provided: Optional[str], model_path: str, default_name: str) -> Optional[str]:
        """解析文件路径"""
        if provided:
            return provided
        candidate = os.path.join(model_path, default_name)
        if os.path.isfile(candidate):
            return candidate
        return None

    # -----------------------------------------------------------------------
    # 文件查找
    # -----------------------------------------------------------------------

    def _find_enc_files(self, model_path: str, suffix: str) -> list:
        """查找加密文件"""
        if os.path.isfile(model_path) and model_path.endswith(suffix):
            return [model_path]

        if os.path.isdir(model_path):
            return sorted([
                os.path.join(model_path, f)
                for f in os.listdir(model_path)
                if f.endswith(suffix)
            ])

        return []

    # -----------------------------------------------------------------------
    # BaseModelLoader 接口实现
    # -----------------------------------------------------------------------

    def download_model(self, model_config: ModelConfig) -> None:
        """加密模型不支持下载"""
        pass

    def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
        """
        加载解密后的权重

        v0.16.0 新增的独立接口，将权重加载与模型初始化分离。
        这是核心切入点：在这里注入加解密逻辑。
        """
        model_path = model_config.model
        decryptor = self._get_decryptor(model_path)

        # 查找加密文件
        gguf_enc_files = self._find_enc_files(model_path, GGUF_ENC)
        st_enc_files = self._find_enc_files(model_path, SAFETENSORS_ENC)

        if gguf_enc_files:
            if len(gguf_enc_files) != 1:
                raise ValueError(
                    f"Expected exactly one .gguf.enc file, found: {gguf_enc_files}"
                )
            logger.info("Loading encrypted GGUF model: %s", gguf_enc_files[0])

            # 复用 GGUFModelLoader 的映射逻辑
            gguf_loader = GGUFModelLoader(self.load_config)
            gguf_weights_map = gguf_loader._get_gguf_weights_map(model_config)

            weights_iter = _encrypted_gguf_weights_iterator(
                gguf_enc_files[0], decryptor, gguf_weights_map
            )
        elif st_enc_files:
            logger.info(
                "Loading encrypted safetensors model (%d shards)", len(st_enc_files)
            )
            weights_iter = _encrypted_safetensors_weights_iterator(
                st_enc_files, decryptor
            )
        else:
            raise ValueError(
                f"No encrypted model files found in: {model_path}. "
                f"Expected files ending with {SAFETENSORS_ENC} or {GGUF_ENC}."
            )

        model.load_weights(weights_iter)

    def load_model(
        self,
        vllm_config: VllmConfig,
        model_config: ModelConfig,
        prefix: str = ""
    ) -> nn.Module:
        """
        加载加密模型

        v0.16.0 签名变化：新增 model_config 和 prefix 参数。
        支持 load_config.device 配置指定加载设备。
        """
        device_config = vllm_config.device_config
        load_config = vllm_config.load_config

        # v0.16.0 新增：支持指定加载设备
        load_device = (
            device_config.device if load_config.device is None
            else load_config.device
        )
        target_device = torch.device(load_device)

        with set_default_torch_dtype(model_config.dtype):
            with target_device:
                model = initialize_model(
                    vllm_config=vllm_config,
                    model_config=model_config,
                    prefix=prefix
                )

            # 使用 load_weights 接口加载权重
            self.load_weights(model, model_config)

            process_weights_after_loading(model, model_config, target_device)

        return model.eval()