#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""
Model encryption tool.

Encrypts safetensors and GGUF model files for secure distribution.

Usage:
    python -m vllm.model_security.cli.encrypt_model \\
        --input  ./my_model \\
        --output ./my_model_encrypted \\
        --key-out ./model.key

    # Or for a single GGUF file:
    python -m vllm.model_security.cli.encrypt_model \\
        --input  ./model.gguf \\
        --output ./model.gguf.enc \\
        --key-out ./model.key
"""

import argparse
import base64
import json
import os
import shutil
import sys
import uuid
from pathlib import Path

# Ensure vllm is importable
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))

from vllm.model_security.crypto import ModelEncryptor, generate_dek, compute_file_hash

# File patterns to encrypt
ENCRYPT_PATTERNS = [".safetensors", ".gguf", ".bin", ".pt"]
# File patterns to copy as-is (model config, tokenizer, etc.)
COPY_PATTERNS = [".json", ".txt", ".model", ".tiktoken", ".py", ".md"]


def encrypt_directory(
    input_dir: str,
    output_dir: str,
    encryptor: ModelEncryptor,
    key_id: bytes,
    verbose: bool = True,
) -> list:
    """
    Encrypt all model weight files in a directory.

    Returns:
        List of encrypted file paths (relative to output_dir)
    """
    os.makedirs(output_dir, exist_ok=True)
    encrypted_files = []

    for fname in sorted(os.listdir(input_dir)):
        src = os.path.join(input_dir, fname)
        if os.path.isdir(src):
            continue

        ext = os.path.splitext(fname)[1].lower()

        if ext in ENCRYPT_PATTERNS:
            dst = os.path.join(output_dir, fname + ".enc")
            if verbose:
                print(f"  Encrypting: {fname} -> {fname}.enc")
            encryptor.encrypt_file(src, dst, key_id)
            encrypted_files.append(fname + ".enc")

        elif any(fname.endswith(p) for p in COPY_PATTERNS) or ext in COPY_PATTERNS:
            dst = os.path.join(output_dir, fname)
            if verbose:
                print(f"  Copying:    {fname}")
            shutil.copy2(src, dst)

        else:
            # Copy unknown files as-is
            dst = os.path.join(output_dir, fname)
            if verbose:
                print(f"  Copying:    {fname} (unknown type)")
            shutil.copy2(src, dst)

    return encrypted_files


def write_model_key_file(
    output_dir: str,
    dek: bytes,
    key_id: bytes,
    encrypted_files: list,
) -> str:
    """
    Write the .model_key file containing the DEK and metadata.

    NOTE: In this implementation, DEK is stored in plaintext in .model_key.
    The security relies on:
    1. License file validation (signature + expiry + fingerprint)
    2. File system permissions on .model_key
    3. The encrypted model files being unreadable without the DEK

    For higher security, encrypt DEK with a vendor-held master key.
    See OfflineLicenseManager in license.py for symmetric encryption approach.
    """
    key_id_hex = key_id.hex()
    dek_b64 = base64.b64encode(dek).decode()

    key_data = {
        "key_id": key_id_hex,
        "algorithm": "AES-256-GCM",
        "dek_plaintext": dek_b64,  # Protected by license validation at runtime
        "encrypted_files": encrypted_files,
    }

    key_path = os.path.join(output_dir, ".model_key")
    with open(key_path, 'w') as f:
        json.dump(key_data, f, indent=2)

    # Restrict permissions on Unix
    try:
        os.chmod(key_path, 0o600)
    except Exception:
        pass

    return key_path


def encrypt_single_file(
    input_path: str,
    output_path: str,
    encryptor: ModelEncryptor,
    key_id: bytes,
    verbose: bool = True,
) -> None:
    """Encrypt a single model file."""
    if verbose:
        print(f"  Encrypting: {input_path} -> {output_path}")
    encryptor.encrypt_file(input_path, output_path, key_id)


def main():
    parser = argparse.ArgumentParser(
        description="Encrypt vLLM model files for secure distribution",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Encrypt a HuggingFace model directory
  python -m vllm.model_security.cli.encrypt_model \\
      --input ./llama-7b \\
      --output ./llama-7b-encrypted \\
      --key-out ./llama-7b.key

  # Encrypt a single GGUF file
  python -m vllm.model_security.cli.encrypt_model \\
      --input ./llama-7b.gguf \\
      --output ./llama-7b.gguf.enc \\
      --key-out ./llama-7b.key

  # Use existing DEK (for re-encryption or multi-file consistency)
  python -m vllm.model_security.cli.encrypt_model \\
      --input ./llama-7b \\
      --output ./llama-7b-encrypted \\
      --dek-file ./existing.key
        """,
    )
    parser.add_argument(
        "--input", "-i", required=True,
        help="Input model directory or single file path"
    )
    parser.add_argument(
        "--output", "-o", required=True,
        help="Output directory or file path for encrypted model"
    )
    parser.add_argument(
        "--key-out", default=None,
        help="Path to save the generated key info JSON (default: <output>/.model_key)"
    )
    parser.add_argument(
        "--dek-file", default=None,
        help="Path to existing .model_key file to reuse its DEK"
    )
    parser.add_argument(
        "--quiet", "-q", action="store_true",
        help="Suppress progress output"
    )

    args = parser.parse_args()
    verbose = not args.quiet

    input_path = os.path.abspath(args.input)
    output_path = os.path.abspath(args.output)

    if not os.path.exists(input_path):
        print(f"ERROR: Input path does not exist: {input_path}", file=sys.stderr)
        sys.exit(1)

    # Load or generate DEK
    if args.dek_file:
        with open(args.dek_file, 'r') as f:
            key_data = json.load(f)
        dek = base64.b64decode(key_data["dek_plaintext"])
        key_id = bytes.fromhex(key_data["key_id"])
        if verbose:
            print(f"Using existing DEK from: {args.dek_file}")
            print(f"Key ID: {key_id.hex()}")
    else:
        dek = generate_dek()
        key_id = uuid.uuid4().bytes
        if verbose:
            print(f"Generated new DEK")
            print(f"Key ID: {key_id.hex()}")

    encryptor = ModelEncryptor(dek)

    if os.path.isfile(input_path):
        # Single file mode
        if verbose:
            print(f"\nEncrypting single file: {input_path}")
        encrypt_single_file(input_path, output_path, encryptor, key_id, verbose)

        # Write key file
        key_out = args.key_out or (output_path + ".key")
        key_dir = os.path.dirname(key_out) or "."
        os.makedirs(key_dir, exist_ok=True)
        fname = os.path.basename(output_path)
        key_data = {
            "key_id": key_id.hex(),
            "algorithm": "AES-256-GCM",
            "dek_plaintext": base64.b64encode(dek).decode(),
            "encrypted_files": [fname],
        }
        with open(key_out, 'w') as f:
            json.dump(key_data, f, indent=2)
        try:
            os.chmod(key_out, 0o600)
        except Exception:
            pass

        if verbose:
            print(f"\nKey saved to: {key_out}")

    elif os.path.isdir(input_path):
        # Directory mode
        if verbose:
            print(f"\nEncrypting model directory: {input_path}")
            print(f"Output directory: {output_path}\n")

        encrypted_files = encrypt_directory(
            input_path, output_path, encryptor, key_id, verbose
        )

        # Write .model_key into output directory
        key_path = write_model_key_file(output_path, dek, key_id, encrypted_files)

        # Also save to --key-out if specified
        if args.key_out:
            shutil.copy2(key_path, args.key_out)
            try:
                os.chmod(args.key_out, 0o600)
            except Exception:
                pass
            if verbose:
                print(f"\nKey also saved to: {args.key_out}")

        if verbose:
            print(f"\nEncryption complete!")
            print(f"  Encrypted files: {len(encrypted_files)}")
            print(f"  Key ID:          {key_id.hex()}")
            print(f"  Key file:        {key_path}")
            print(f"\nNext step: generate a license file with gen_license.py")
            print(f"  python -m vllm.model_security.cli.gen_license \\")
            print(f"      --model-dir {output_path} \\")
            print(f"      --vendor-key vendor_private.pem \\")
            print(f"      --customer 'Your Customer' \\")
            print(f"      --out license.lic")
    else:
        print(f"ERROR: Input path is neither a file nor directory: {input_path}",
              file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
