#!/usr/bin/env python3
"""
verify_xllm_moe_load.py — 在 BI-V100 真机验证 xllm_moe.so 能否加载

用法: python3 verify_xllm_moe_load.py
      或 Docker 内: python3 /workspace/qwen3_6_scripts/verify_xllm_moe_load.py

验证内容:
  1. 所有 prebuilt .so 的 dlopen 可行性
  2. xllm_moe.so 的 moe_fused_topk 是否可调用
  3. corex_moe_topk_softmax.so 的 moe_topk_softmax 是否可调用
  4. ix_startup_patch.apply() 完整执行结果
  5. _custom_ops.topk_softmax 实际走哪条路径
"""

import os
import sys
import importlib
import importlib.util
import traceback

def section(title):
    print(f"\n{'='*65}")
    print(f"  {title}")
    print(f"{'='*65}")

def try_load_so(name, path):
    """Try to load a .so file, return (module, error_string)."""
    if not os.path.isfile(path):
        return None, f"FILE NOT FOUND: {path}"
    try:
        spec = importlib.util.spec_from_file_location(name, path)
        if spec is None:
            return None, "spec_from_file_location returned None"
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        funcs = [x for x in dir(mod) if not x.startswith('_')]
        return mod, f"OK ({len(funcs)} functions: {', '.join(funcs[:10])})"
    except Exception as e:
        return None, f"{type(e).__name__}: {e}"

# =========================================================================
section("1. Locate all .so files")
# =========================================================================

# Search paths (same as xllm_ops.py)
search_dirs = []
try:
    import vllm
    vllm_root = os.path.dirname(vllm.__file__)
    search_dirs.append(vllm_root)
    print(f"  vllm root: {vllm_root}")
except ImportError:
    print("  vllm not installed")
    vllm_root = None

for d in [
    "/workspace/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10",
    "/workspace/qwen3_6_scripts/ex_engine/prebuilt",
    os.path.join(os.path.dirname(__file__), "prebuilt", "corex-3.2.3-ivcore10"),
    os.path.join(os.path.dirname(__file__), "qwen3_6_scripts", "prebuilt", "corex-3.2.3-ivcore10"),
]:
    if os.path.isdir(d):
        search_dirs.append(os.path.normpath(d))

search_dirs = list(dict.fromkeys(search_dirs))  # dedupe preserving order
print(f"  search dirs: {search_dirs}")

# =========================================================================
section("2. Load test: all xllm_*.so and key bridge .so")
# =========================================================================

test_modules = [
    "xllm_moe",
    "xllm_norm",
    "xllm_activation",
    "xllm_cache",
    "xllm_rope",
    "ix_moe_bridge",
    "ix_full_bridge",
    "corex_moe_topk_softmax",
    "corex_moe_index_combine",
    "corex_moe_direct_routed",
    "corex_moe_weight_gather",
    "corex_moe_exact_reduce",
    "gemm_grouped",
    "corex_batched_gemm",
]

loaded_modules = {}
for name in test_modules:
    found = False
    for d in search_dirs:
        path = os.path.join(d, f"{name}.so")
        if os.path.isfile(path):
            mod, status = try_load_so(name, path)
            tag = "✓" if mod else "✗"
            print(f"  {tag} {name:35s} {status}")
            if mod:
                loaded_modules[name] = mod
            found = True
            break
    if not found:
        print(f"  - {name:35s} NOT FOUND in any search dir")

# =========================================================================
section("3. Functional test: xllm_moe.moe_fused_topk")
# =========================================================================

if "xllm_moe" in loaded_modules:
    mod = loaded_modules["xllm_moe"]
    print(f"  Exported functions: {[x for x in dir(mod) if not x.startswith('_')]}")
    if hasattr(mod, "moe_fused_topk"):
        try:
            import torch
            if torch.cuda.is_available():
                # moe_fused_topk(gating_output, topk, renormalize, correction_bias, scoring_func)
                gating = torch.randn(4, 8, device="cuda", dtype=torch.float32)
                w, ids = mod.moe_fused_topk(gating, 2)
                print(f"  ✓ moe_fused_topk: weights shape={w.shape}, ids shape={ids.shape}")
                print(f"    weights={w[0].tolist()}, ids={ids[0].tolist()}")
            else:
                print(f"  - no CUDA device, skipping functional test")
        except Exception as e:
            print(f"  ✗ moe_fused_topk call FAILED: {e}")
            traceback.print_exc()
    else:
        print(f"  ✗ moe_fused_topk NOT in module attrs")
        print(f"    available: {[x for x in dir(mod) if not x.startswith('_')]}")
else:
    print(f"  ✗ xllm_moe.so not loaded, cannot test")

# =========================================================================
section("4. Functional test: corex_moe_topk_softmax.moe_topk_softmax")
# =========================================================================

if "corex_moe_topk_softmax" in loaded_modules:
    mod = loaded_modules["corex_moe_topk_softmax"]
    print(f"  Exported functions: {[x for x in dir(mod) if not x.startswith('_')]}")
    if hasattr(mod, "moe_topk_softmax"):
        try:
            import torch
            if torch.cuda.is_available():
                gating = torch.randn(4, 8, device="cuda", dtype=torch.float32)
                w, ids = mod.moe_topk_softmax(gating, 2, True)
                print(f"  ✓ moe_topk_softmax: weights shape={w.shape}, ids shape={ids.shape}")
                print(f"    weights={w[0].tolist()}, ids={ids[0].tolist()}")
            else:
                print(f"  - no CUDA device")
        except Exception as e:
            print(f"  ✗ moe_topk_softmax call FAILED: {e}")
            traceback.print_exc()
    else:
        print(f"  ✗ moe_topk_softmax NOT in module attrs")
else:
    print(f"  ✗ corex_moe_topk_softmax.so not loaded, cannot test")

# =========================================================================
section("5. ixformer.functions.vllm_moe_topk_softmax existence")
# =========================================================================

try:
    import ixformer.functions as ixf_F
    has_it = hasattr(ixf_F, "vllm_moe_topk_softmax")
    print(f"  ixformer.functions.vllm_moe_topk_softmax: {'✓ EXISTS' if has_it else '✗ MISSING (expected on BI-V100)'}")
    if not has_it:
        # Check what MoE-related functions DO exist
        moe_funcs = [n for n in dir(ixf_F) if 'moe' in n.lower() or 'topk' in n.lower()]
        print(f"  MoE-related functions that DO exist: {moe_funcs if moe_funcs else 'NONE'}")
except ImportError as e:
    print(f"  ixformer not available: {e}")

# =========================================================================
section("6. ix_startup_patch.apply() test")
# =========================================================================

try:
    from vllm import ix_startup_patch
    print(f"  ix_startup_patch imported from: {ix_startup_patch.__file__}")
    n = ix_startup_patch.apply()
    print(f"  ✓ apply() returned: {n} patches applied")
except ImportError as e:
    print(f"  ✗ import failed: {e}")
except Exception as e:
    print(f"  ✗ apply() failed: {type(e).__name__}: {e}")
    traceback.print_exc()

# =========================================================================
section("7. _custom_ops.topk_softmax dispatch test")
# =========================================================================

try:
    import vllm._custom_ops as ops
    print(f"  _custom_ops from: {ops.__file__}")

    import torch
    if torch.cuda.is_available():
        gating = torch.randn(4, 8, device="cuda", dtype=torch.float32)
        tw = torch.empty(4, 2, device="cuda", dtype=torch.float32)
        ti = torch.empty(4, 2, device="cuda", dtype=torch.int32)
        tei = torch.empty(4, 2, device="cuda", dtype=torch.int32)
        try:
            ops.topk_softmax(tw, ti, tei, gating)
            print(f"  ✓ topk_softmax succeeded")
            print(f"    weights={tw[0].tolist()}, ids={ti[0].tolist()}")
        except Exception as e:
            print(f"  ✗ topk_softmax FAILED: {e}")

        # Check which path was taken by inspecting the function
        import inspect
        src = inspect.getsource(ops.topk_softmax)
        if "hasattr(ixf_F" in src:
            print(f"  → using PATCHED _custom_ops.py (has fallback chain)")
        elif "ixf_F.vllm_moe_topk_softmax" in src and "hasattr" not in src:
            print(f"  → using ORIGINAL _custom_ops.py (NO fallback, will crash)")
        else:
            print(f"  → using UNKNOWN version of topk_softmax")
    else:
        print(f"  - no CUDA device")
except Exception as e:
    print(f"  ✗ {e}")
    traceback.print_exc()

# =========================================================================
section("8. Library compatibility check")
# =========================================================================

try:
    import torch
    print(f"  PyTorch: {torch.__version__}")
    print(f"  CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"  CUDA version: {torch.version.cuda}")
        print(f"  Device: {torch.cuda.get_device_name(0)}")

    # Check if libcudart.so.10.2 exists (required by all xllm .so)
    import ctypes
    for lib in ["libcudart.so.10.2", "libcuinfer.so.7", "libc10.so",
                "libtorch_python.so", "libtorch_cuda.so"]:
        try:
            ctypes.CDLL(lib)
            print(f"  ✓ {lib}")
        except OSError as e:
            print(f"  ✗ {lib}: {e}")
except Exception as e:
    print(f"  error: {e}")

print(f"\n{'='*65}")
print("  DONE — paste this entire output back")
print(f"{'='*65}")