#!/usr/bin/env python3
"""
CUDA Compatibility Broker for helix64

Papers over CUDA/PyTorch version friction by:
1. Introspecting wheel CUDA vs host driver
2. Activating userspace compat profiles when needed
3. Generating receipts for all decisions
4. Failing closed when incompatible

Usage:
    from helix_cdc.cuda_compat import ensure_cuda_compat

    # Before importing torch or running workloads
    ensure_cuda_compat()
"""

import os
import sys
import json
import subprocess
import pathlib
import hashlib
from datetime import datetime, timezone
from typing import Optional, Dict, Any

def run_cmd(cmd: list) -> Optional[str]:
    """Run command and return output, or None if failed"""
    try:
        return subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL).strip()
    except:
        return None

def get_cpu_flags() -> str:
    """Get CPU flags for provenance"""
    try:
        with open("/proc/cpuinfo") as f:
            for line in f:
                if line.startswith("flags"):
                    return line.split(":", 1)[1].strip()
    except:
        pass
    return "unknown"

def get_git_commit() -> str:
    """Get current git commit hash"""
    try:
        result = subprocess.check_output(
            ["git", "rev-parse", "HEAD"],
            text=True,
            stderr=subprocess.DEVNULL
        )
        return result.strip()
    except:
        return "unknown"

def get_impl_sha256() -> str:
    """Get SHA256 hash of this implementation file"""
    try:
        impl_path = pathlib.Path(__file__)
        return hashlib.sha256(impl_path.read_bytes()).hexdigest()
    except:
        return "unknown"

def get_wheel_cuda_info() -> Dict[str, Any]:
    """
    Introspect PyTorch wheel for CUDA version

    Returns: {cuda_version, cudnn_version, cuda_available}
    """
    info = {
        "cuda_version": None,
        "cudnn_version": None,
        "cuda_available": False
    }

    try:
        import torch
        info["cuda_version"] = getattr(torch.version, "cuda", None)  # e.g., "12.1"
        info["cuda_available"] = torch.cuda.is_available()

        if torch.backends.cudnn.is_available():
            info["cudnn_version"] = torch.backends.cudnn.version()
    except ImportError:
        pass  # torch not installed
    except Exception:
        pass  # some other issue

    return info

def get_host_driver_info() -> Dict[str, Any]:
    """
    Query NVIDIA host driver for capabilities

    Returns: {driver_version, compute_caps, cuda_support}
    """
    info = {
        "driver_version": None,
        "compute_caps": [],
        "cuda_support": None
    }

    # Driver version
    driver = run_cmd(["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"])
    if driver:
        info["driver_version"] = driver.split('\n')[0].strip()

    # Compute capabilities (SM versions)
    caps = run_cmd(["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"])
    if caps:
        info["compute_caps"] = [c.strip() for c in caps.split('\n') if c.strip()]

    # Maximum CUDA version supported by driver
    # Approximation: driver >= 535 supports CUDA 12.x, >= 470 supports 11.x
    if info["driver_version"]:
        major = int(info["driver_version"].split('.')[0])
        if major >= 535:
            info["cuda_support"] = "12.x"
        elif major >= 470:
            info["cuda_support"] = "11.x"
        elif major >= 450:
            info["cuda_support"] = "10.x"
        else:
            info["cuda_support"] = "unknown"

    return info

def check_compatibility(wheel_info: Dict, driver_info: Dict) -> Dict[str, Any]:
    """
    Determine if wheel CUDA and host driver are compatible

    Returns: {compatible, action, reason}
    """
    result = {
        "compatible": True,
        "action": "ok",
        "reason": None
    }

    # No CUDA in wheel → CPU only
    if not wheel_info["cuda_version"]:
        result["action"] = "cpu_only"
        result["reason"] = "PyTorch wheel has no CUDA support"
        return result

    # No driver detected → CPU fallback
    if not driver_info["driver_version"]:
        result["compatible"] = False
        result["action"] = "cpu_fallback"
        result["reason"] = "No NVIDIA driver detected (nvidia-smi failed)"
        return result

    # Check if driver supports wheel's CUDA version
    wheel_cuda_major = int(wheel_info["cuda_version"].split('.')[0])
    driver_cuda = driver_info["cuda_support"]

    if driver_cuda == "unknown":
        result["compatible"] = False
        result["action"] = "driver_too_old"
        result["reason"] = f"Driver {driver_info['driver_version']} is too old for CUDA {wheel_info['cuda_version']}"
        return result

    driver_cuda_major = int(driver_cuda.split('.')[0])

    if wheel_cuda_major > driver_cuda_major:
        result["compatible"] = False
        result["action"] = "driver_too_old"
        result["reason"] = f"Wheel requires CUDA {wheel_info['cuda_version']}, driver supports max {driver_cuda}"
        return result

    # Compatible - may need userspace bundle
    if wheel_cuda_major == driver_cuda_major:
        result["action"] = "ok"
    else:
        result["action"] = "userspace_bundle_recommended"
        result["reason"] = f"Minor version mismatch (wheel: {wheel_info['cuda_version']}, driver: {driver_cuda})"

    return result

def activate_compat_profile(cuda_version: str, profiles_root: pathlib.Path) -> Optional[str]:
    """
    Activate userspace CUDA compat profile

    Sets LD_LIBRARY_PATH to bundle with matching CUDA libraries.
    Returns: profile name if activated, None otherwise
    """
    # Normalize version to profile name (e.g., "12.1" → "cu121")
    profile_name = f"cu{cuda_version.replace('.', '')}"
    profile_path = profiles_root / profile_name

    if not profile_path.exists():
        return None

    # Prepend to LD_LIBRARY_PATH
    old_path = os.environ.get("LD_LIBRARY_PATH", "")
    new_path = f"{profile_path}:{old_path}" if old_path else str(profile_path)
    os.environ["LD_LIBRARY_PATH"] = new_path

    # Enable lazy loading (performance)
    os.environ.setdefault("CUDA_MODULE_LOADING", "LAZY")

    # Enable PTX JIT for cross-SM compatibility
    os.environ.setdefault("CUDA_FORCE_PTX_JIT", "1")

    return profile_name

def write_receipt(decision: Dict[str, Any], output_path: pathlib.Path):
    """Write CUDA compatibility receipt"""
    receipt = {
        "protocol_version": "helix_cdc:v0.1.0",
        "schema_version": "cuda_compat:v1",
        "claim": {
            "component": "helix64-cuda-compat",
            "operation": "cuda_compatibility_resolution",
            "wheel_cuda": decision["wheel_info"]["cuda_version"],
            "driver_version": decision["driver_info"]["driver_version"],
            "driver_cuda_support": decision["driver_info"]["cuda_support"],
            "compatible": decision["compat"]["compatible"],
            "action": decision["compat"]["action"],
            "reason": decision["compat"]["reason"],
            "profile_activated": decision.get("profile_activated"),
            "environment": {
                "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH"),
                "CUDA_MODULE_LOADING": os.environ.get("CUDA_MODULE_LOADING"),
                "CUDA_FORCE_PTX_JIT": os.environ.get("CUDA_FORCE_PTX_JIT")
            }
        },
        "provenance": {
            "timestamp_utc": datetime.now(timezone.utc).isoformat(),
            "hostname": os.uname().nodename if hasattr(os, 'uname') else "unknown",
            "git_commit": get_git_commit(),
            "impl_sha256": get_impl_sha256(),
            "cpu_flags": get_cpu_flags()
        },
        "acceptance_gates": {
            "passes": decision["compat"]["compatible"],
            "gates": {
                "driver_detected": decision["driver_info"]["driver_version"] is not None,
                "compatible_versions": decision["compat"]["compatible"],
                "impl_pinned": True,
                "determinism_ok": True  # CUDA compat is deterministic (same inputs → same decision)
            }
        }
    }

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(receipt, indent=2))

def ensure_cuda_compat(
    force_profile: Optional[str] = None,
    receipt_path: Optional[pathlib.Path] = None,
    fail_on_incompatible: bool = False
) -> Dict[str, Any]:
    """
    Main entry point: ensure CUDA compatibility before running workloads

    Args:
        force_profile: Override auto-detection (e.g., "cu121")
        receipt_path: Where to write receipt (default: artifacts/cuda_compat.receipt.json)
        fail_on_incompatible: Exit if incompatible (default: warn and continue)

    Returns:
        Decision dict with compatibility status
    """
    # Default receipt path
    if receipt_path is None:
        receipt_path = pathlib.Path("artifacts") / "cuda_compat.receipt.json"

    # Profiles root (alongside this file)
    profiles_root = pathlib.Path(__file__).parent.parent / "profiles" / "cuda"

    # 1. Introspect wheel
    wheel_info = get_wheel_cuda_info()

    # 2. Query host driver
    driver_info = get_host_driver_info()

    # 3. Check compatibility
    compat = check_compatibility(wheel_info, driver_info)

    # 4. Activate profile if needed/requested
    profile_activated = None

    if force_profile:
        # Manual override
        profile_activated = activate_compat_profile(force_profile.replace("cu", ""), profiles_root)
    elif compat["action"] == "userspace_bundle_recommended" and wheel_info["cuda_version"]:
        # Auto-activate matching profile
        profile_activated = activate_compat_profile(wheel_info["cuda_version"], profiles_root)

    # Allow env override
    env_profile = os.environ.get("HELIX64_CUDA_PROFILE")
    if env_profile:
        profile_activated = activate_compat_profile(env_profile.replace("cu", ""), profiles_root)

    # Build decision
    decision = {
        "wheel_info": wheel_info,
        "driver_info": driver_info,
        "compat": compat,
        "profile_activated": profile_activated
    }

    # 5. Write receipt
    write_receipt(decision, receipt_path)

    # 6. Handle incompatibility
    if not compat["compatible"]:
        msg = f"⚠️  CUDA compatibility issue: {compat['reason']}"
        if compat["action"] == "cpu_fallback":
            print(msg, file=sys.stderr)
            print("   Falling back to CPU backend", file=sys.stderr)
        elif compat["action"] == "driver_too_old":
            print(msg, file=sys.stderr)
            print(f"   Update NVIDIA driver or use --backend cpu", file=sys.stderr)
            if fail_on_incompatible:
                sys.exit(1)

    return decision
