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

Generates a unique machine fingerprint for license binding.
Run this on the target deployment machine and send the output to the vendor.

Usage:
    vllm-get-fingerprint
    vllm-get-fingerprint --include-gpu
    vllm-get-fingerprint --json
"""

import argparse
import json
import sys
import hashlib
import platform
import uuid
from typing import List, Optional

try:
    import psutil
except ImportError:
    psutil = None


def get_cpu_id() -> str:
    """Get CPU/machine identifier (platform-dependent)."""
    system = platform.system()

    if system == "Linux":
        try:
            with open("/etc/machine-id", "r") as f:
                machine_id = f.read().strip()
                if machine_id:
                    return f"machine-id:{machine_id}"
        except Exception:
            pass

        try:
            with open("/sys/class/dmi/id/product_uuid", "r") as f:
                product_uuid = f.read().strip()
                if product_uuid:
                    return f"dmi-uuid:{product_uuid}"
        except Exception:
            pass

        try:
            with open("/proc/cpuinfo", "r") as f:
                for line in f:
                    if "Serial" in line:
                        serial = line.split(":")[-1].strip()
                        if serial and serial != "0000000000000000":
                            return f"cpu-serial:{serial}"
        except Exception:
            pass

    elif system == "Windows":
        try:
            import subprocess
            result = subprocess.run(
                ["wmic", "cpu", "get", "ProcessorId"],
                capture_output=True,
                text=True,
                check=True
            )
            lines = result.stdout.strip().split("\n")
            if len(lines) > 1:
                return lines[1].strip()
        except Exception:
            pass

    return f"{platform.processor()}_{platform.machine()}_{platform.node()}"


def get_mac_addresses() -> List[str]:
    """Get MAC addresses of network interfaces."""
    macs = []

    if psutil:
        try:
            for interface, addrs in psutil.net_if_addrs().items():
                for addr in addrs:
                    if addr.family == psutil.AF_LINK:
                        macs.append(addr.address)
        except Exception:
            pass

    if not macs:
        mac = uuid.getnode()
        macs.append(':'.join(f'{(mac >> i) & 0xff:02x}' for i in range(0, 48, 8)))

    return sorted(macs)


def get_gpu_uuids() -> List[str]:
    """Get GPU UUIDs (NVIDIA only for now)."""
    uuids = []

    try:
        import subprocess
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=uuid", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            check=True
        )
        uuids = [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
    except Exception:
        pass

    return uuids


def get_npu_uuids() -> List[str]:
    """Get NPU UUIDs for Ascend devices."""
    uuids = []

    try:
        import torch
        if hasattr(torch, 'npu') and torch.npu.is_available():
            device_count = torch.npu.device_count()
            for i in range(device_count):
                try:
                    props = torch.npu.get_device_properties(i)
                    if hasattr(props, 'uuid') and props.uuid:
                        uuids.append(props.uuid)
                    else:
                        uuids.append(f"npu-{i}:{torch.npu.get_device_name(i)}")
                except Exception:
                    uuids.append(f"npu-{i}")
    except ImportError:
        pass

    return uuids


def get_machine_fingerprint(
    include_gpu: bool = False,
    salt: Optional[str] = None
) -> str:
    """
    Generate a unique machine fingerprint.

    Args:
        include_gpu: Whether to include GPU UUIDs (may cause issues in containers)
        salt: Optional salt for hashing

    Returns:
        Hex-encoded SHA-256 hash of machine characteristics
    """
    components = []

    # CPU ID
    cpu_id = get_cpu_id()
    components.append(f"cpu:{cpu_id}")

    # MAC addresses
    macs = get_mac_addresses()
    components.append(f"mac:{','.join(macs)}")

    # Hostname
    hostname = platform.node()
    components.append(f"host:{hostname}")

    # NPU UUIDs (auto-detect Ascend devices)
    npu_uuids = get_npu_uuids()
    if npu_uuids:
        components.append(f"npu:{','.join(npu_uuids)}")

    # GPU UUIDs (optional)
    if include_gpu:
        gpu_uuids = get_gpu_uuids()
        if gpu_uuids:
            components.append(f"gpu:{','.join(gpu_uuids)}")

    # Add salt if provided
    if salt:
        components.append(f"salt:{salt}")

    # Hash all components
    fingerprint_str = "|".join(components)
    fingerprint_hash = hashlib.sha256(fingerprint_str.encode()).hexdigest()

    return fingerprint_hash


def get_fingerprint_info() -> dict:
    """
    Get detailed fingerprint information for debugging.

    Returns:
        Dictionary with fingerprint components
    """
    return {
        "cpu_id": get_cpu_id(),
        "mac_addresses": get_mac_addresses(),
        "hostname": platform.node(),
        "gpu_uuids": get_gpu_uuids(),
        "npu_uuids": get_npu_uuids(),
        "platform": platform.platform(),
        "fingerprint": get_machine_fingerprint(),
        "fingerprint_with_gpu": get_machine_fingerprint(include_gpu=True),
    }


def main():
    parser = argparse.ArgumentParser(
        description="Get machine fingerprint for license binding",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Get basic fingerprint (recommended for most cases)
  vllm-get-fingerprint

  # Include GPU UUIDs in fingerprint
  vllm-get-fingerprint --include-gpu

  # Output as JSON for scripting
  vllm-get-fingerprint --json

  # Show all components used in fingerprint
  vllm-get-fingerprint --verbose
        """,
    )
    parser.add_argument(
        "--include-gpu", action="store_true",
        help="Include GPU UUIDs in fingerprint (may change if GPUs are replaced)"
    )
    parser.add_argument(
        "--json", action="store_true",
        help="Output as JSON"
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true",
        help="Show all fingerprint components"
    )

    args = parser.parse_args()

    if args.verbose or args.json:
        info = get_fingerprint_info()
        if args.json:
            print(json.dumps(info, indent=2))
        else:
            print("Machine Fingerprint Components:")
            print(f"  CPU ID:          {info['cpu_id']}")
            print(f"  MAC Addresses:   {', '.join(info['mac_addresses'])}")
            print(f"  Hostname:        {info['hostname']}")
            print(f"  GPU UUIDs:       {', '.join(info['gpu_uuids']) or '(none)'}")
            print(f"  NPU UUIDs:       {', '.join(info['npu_uuids']) or '(none)'}")
            print(f"  Platform:        {info['platform']}")
            print()
            print(f"Fingerprint (no GPU):   {info['fingerprint']}")
            print(f"Fingerprint (with GPU): {info['fingerprint_with_gpu']}")
            print()
            if args.include_gpu:
                print(f"Use this fingerprint: {info['fingerprint_with_gpu']}")
            else:
                print(f"Use this fingerprint: {info['fingerprint']}")
    else:
        fp = get_machine_fingerprint(include_gpu=args.include_gpu)
        print(fp)


if __name__ == "__main__":
    main()