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

Creates signed license files for encrypted models.

Usage:
    # Generate vendor keypair (one-time)
    python -m vllm.model_security.cli.gen_license --gen-keypair \\
        --private-key vendor_private.pem \\
        --public-key vendor_public.pem

    # Generate license for a customer
    python -m vllm.model_security.cli.gen_license \\
        --model-dir ./my_model_encrypted \\
        --vendor-key vendor_private.pem \\
        --customer "Acme Corp" \\
        --fingerprint <hash_from_get_fingerprint.py> \\
        --expire 2027-12-31 \\
        --out acme_corp.lic
"""

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

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

from vllm.model_security.license import LicenseManager, FingerprintPolicy
from vllm.model_security.fingerprint import get_machine_fingerprint


def main():
    parser = argparse.ArgumentParser(
        description="Generate license files for encrypted vLLM models",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate vendor keypair (one-time setup)
  python -m vllm.model_security.cli.gen_license --gen-keypair \\
      --private-key vendor_private.pem \\
      --public-key vendor_public.pem

  # Generate license with no machine binding (for testing)
  python -m vllm.model_security.cli.gen_license \\
      --model-dir ./llama-7b-encrypted \\
      --vendor-key vendor_private.pem \\
      --customer "Test Customer" \\
      --policy none \\
      --out test.lic

  # Generate license bound to specific machine
  python -m vllm.model_security.cli.gen_license \\
      --model-dir ./llama-7b-encrypted \\
      --vendor-key vendor_private.pem \\
      --customer "Acme Corp" \\
      --fingerprint a1b2c3d4... \\
      --policy single \\
      --expire 2027-12-31 \\
      --out acme_corp.lic

  # Generate license for cluster (multiple machines)
  python -m vllm.model_security.cli.gen_license \\
      --model-dir ./llama-7b-encrypted \\
      --vendor-key vendor_private.pem \\
      --customer "BigCorp" \\
      --fingerprint fp1,fp2,fp3 \\
      --policy cluster \\
      --max-nodes 3 \\
      --out bigcorp.lic
        """,
    )

    # Keypair generation mode
    parser.add_argument(
        "--gen-keypair", action="store_true",
        help="Generate RSA keypair for license signing"
    )
    parser.add_argument(
        "--private-key", default="vendor_private.pem",
        help="Path to vendor private key (for signing)"
    )
    parser.add_argument(
        "--public-key", default="vendor_public.pem",
        help="Path to vendor public key (for verification)"
    )

    # License generation mode
    parser.add_argument(
        "--model-dir", default=None,
        help="Path to encrypted model directory (containing .model_key)"
    )
    parser.add_argument(
        "--vendor-key", default=None,
        help="Path to vendor private key for signing"
    )
    parser.add_argument(
        "--vendor-secret", default=None,
        help="Vendor secret for DEK encryption (hex string or file path starting with @)"
    )
    parser.add_argument(
        "--customer", "-c", default=None,
        help="Customer/organization name"
    )
    parser.add_argument(
        "--fingerprint", "-f", default=None,
        help="Machine fingerprint(s), comma-separated for cluster policy"
    )
    parser.add_argument(
        "--policy", "-p", default="none",
        choices=["none", "single", "cluster"],
        help="Fingerprint binding policy (default: none)"
    )
    parser.add_argument(
        "--expire", "-e", default=None,
        help="Expiry date in ISO format (YYYY-MM-DD), omit for no expiry"
    )
    parser.add_argument(
        "--max-nodes", type=int, default=1,
        help="Maximum nodes for cluster policy (default: 1)"
    )
    parser.add_argument(
        "--out", "-o", default="model.lic",
        help="Output license file path (default: model.lic)"
    )
    parser.add_argument(
        "--show-current-fingerprint", action="store_true",
        help="Show current machine fingerprint and exit"
    )

    args = parser.parse_args()

    # Show fingerprint mode
    if args.show_current_fingerprint:
        fp = get_machine_fingerprint()
        fp_gpu = get_machine_fingerprint(include_gpu=True)
        print(f"Current machine fingerprint (no GPU):   {fp}")
        print(f"Current machine fingerprint (with GPU): {fp_gpu}")
        print("\nUse this value with --fingerprint when generating a license.")
        return

    # Keypair generation mode
    if args.gen_keypair:
        print(f"Generating RSA-2048 keypair...")
        LicenseManager.generate_keypair(
            args.private_key,
            args.public_key,
            key_size=2048,
        )
        print(f"  Private key: {args.private_key}")
        print(f"  Public key:  {args.public_key}")
        print(f"\nKEEP THE PRIVATE KEY SECURE!")
        print(f"Distribute the public key with encrypted models.")
        return

    # License generation mode
    if not args.model_dir:
        print("ERROR: --model-dir is required for license generation", file=sys.stderr)
        parser.print_help()
        sys.exit(1)

    if not args.vendor_key:
        print("ERROR: --vendor-key is required for license generation", file=sys.stderr)
        parser.print_help()
        sys.exit(1)

    if not args.customer:
        print("ERROR: --customer is required for license generation", file=sys.stderr)
        parser.print_help()
        sys.exit(1)

    # Load or generate vendor_secret
    if args.vendor_secret:
        if args.vendor_secret.startswith('@'):
            # Load from file
            with open(args.vendor_secret[1:], 'rb') as f:
                vendor_secret = f.read()
        else:
            # Hex string
            vendor_secret = bytes.fromhex(args.vendor_secret)
    else:
        # Generate a random vendor_secret and save it
        import os
        vendor_secret = os.urandom(32)
        secret_path = "vendor_secret.bin"
        with open(secret_path, 'wb') as f:
            f.write(vendor_secret)
        print(f"Generated new vendor_secret and saved to: {secret_path}")
        print(f"KEEP THIS FILE SECURE! It's required to decrypt models at runtime.")

    # Load .model_key from encrypted model directory
    import os
    model_key_path = os.path.join(args.model_dir, ".model_key")
    if not os.path.isfile(model_key_path):
        print(f"ERROR: .model_key not found in {args.model_dir}", file=sys.stderr)
        print(f"Make sure the model directory is encrypted with encrypt_model.py", file=sys.stderr)
        sys.exit(1)

    with open(model_key_path, 'r') as f:
        key_data = json.load(f)

    key_id = key_data["key_id"]
    dek_b64 = key_data["dek_plaintext"]
    dek = base64.b64decode(dek_b64)

    # Parse fingerprints
    policy = FingerprintPolicy(args.policy)
    fingerprints = []
    if args.fingerprint:
        fingerprints = [fp.strip() for fp in args.fingerprint.split(",")]

    if policy != FingerprintPolicy.NONE and not fingerprints:
        print(f"ERROR: --fingerprint is required for policy '{policy}'", file=sys.stderr)
        print(f"Run with --show-current-fingerprint to get the machine fingerprint", file=sys.stderr)
        sys.exit(1)

    # Create license
    print(f"Generating license...")
    print(f"  Customer:    {args.customer}")
    print(f"  Model:       {args.model_dir}")
    print(f"  Key ID:      {key_id}")
    print(f"  Policy:      {policy}")
    if fingerprints:
        print(f"  Fingerprints: {len(fingerprints)}")
        for fp in fingerprints:
            print(f"    - {fp}")
    if args.expire:
        print(f"  Expires:     {args.expire}")
    else:
        print(f"  Expires:     Never")

    manager = LicenseManager(args.vendor_key)
    license_info = manager.create_license(
        customer=args.customer,
        model_key_id=key_id,
        dek=dek,
        vendor_secret=vendor_secret,
        fingerprint_policy=policy,
        allowed_fingerprints=fingerprints,
        expire_at=args.expire,
        max_nodes=args.max_nodes,
    )

    manager.save_license(license_info, args.out)
    print(f"\nLicense saved to: {args.out}")
    print(f"License ID: {license_info.license_id}")
    print(f"\nDistribute this license file along with:")
    print(f"  - Encrypted model directory: {args.model_dir}")
    print(f"  - Vendor public key: {args.public_key}")
    print(f"\nTo load the encrypted model, provide the vendor_secret via environment variable:")
    print(f"  export VLLM_VENDOR_SECRET={vendor_secret.hex()}")
    print(f"  vllm serve {args.model_dir}")
    print(f"\nOR via model-loader-extra-config:")
    print(f'  --model-loader-extra-config \'{{"vendor_secret": "{vendor_secret.hex()}"}}\'')
    print(f"\nWARNING: Keep vendor_secret as secure as vendor_private.pem!")


if __name__ == "__main__":
    main()
