#!/usr/bin/env python3
"""
BC-250 ComfyUI GGUF Integration Test
Tests the actual ComfyUI loading pipeline step by step.
"""
import os, sys, time

os.environ["HSA_OVERRIDE_GFX_VERSION"] = "10.1.0"
os.environ["HSA_ENABLE_SDMA"] = "0"
os.environ["HIP_VISIBLE_DEVICES"] = "0"
os.environ["BC250_SOFTMAX_THRESHOLD"] = "512"

sys.path.insert(0, "/home/fabian/ComfyUI")

print("[TEST] Importing bc250_softmax_patch...", flush=True)
import bc250_softmax_patch

print("[TEST] Importing torch...", flush=True)
t0 = time.time()
import torch
print(f"[TEST] torch ready in {time.time()-t0:.1f}s", flush=True)

print(f"[TEST] CUDA available: {torch.cuda.is_available()}", flush=True)
print(f"[TEST] Device: {torch.cuda.get_device_name(0)}", flush=True)

# Step 1: Load GGUF using ComfyUI-GGUF loader
print(f"\n[TEST] === Step 1: gguf_sd_loader ===", flush=True)

# Import ComfyUI-GGUF properly as a package
import importlib
custom_nodes_path = "/home/fabian/ComfyUI/custom_nodes"
if custom_nodes_path not in sys.path:
    sys.path.insert(0, custom_nodes_path)

# Force import as package
gguf_pkg = importlib.import_module("ComfyUI-GGUF")
from importlib import import_module
gguf_loader = import_module("ComfyUI-GGUF.loader")
gguf_dequant = import_module("ComfyUI-GGUF.dequant")
gguf_ops = import_module("ComfyUI-GGUF.ops")

gguf_sd_loader = gguf_loader.gguf_sd_loader

t1 = time.time()
sd, extra = gguf_sd_loader("/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf")
dt = time.time() - t1
print(f"[TEST] State dict loaded in {dt:.1f}s", flush=True)
print(f"[TEST] Keys: {len(sd)}", flush=True)
print(f"[TEST] Architecture: {extra.get('arch_str')}", flush=True)

# Check some tensor info
is_quantized = gguf_dequant.is_quantized
q_count = sum(1 for v in sd.values() if is_quantized(v))
print(f"[TEST] Quantized tensors: {q_count}/{len(sd)}", flush=True)

# Step 2: Test a single dequantize on CPU
print(f"\n[TEST] === Step 2: Single tensor dequant ===", flush=True)
dequantize_tensor = gguf_dequant.dequantize_tensor
for k, v in sd.items():
    if is_quantized(v):
        print(f"[TEST] Dequantizing: {k} shape={v.tensor_shape} type={v.tensor_type}", flush=True)
        t2 = time.time()
        result = dequantize_tensor(v, dtype=torch.float16)
        dt = time.time() - t2
        print(f"[TEST] Done in {dt:.3f}s -> {result.shape} {result.dtype}", flush=True)
        
        # Move to GPU
        t3 = time.time()
        gpu = result.to("cuda:0")
        torch.cuda.synchronize()
        dt2 = time.time() - t3
        print(f"[TEST] GPU transfer in {dt2:.3f}s", flush=True)
        del gpu, result
        break

# Step 3: Test loading the model through ComfyUI model management
print(f"\n[TEST] === Step 3: ComfyUI model loading ===", flush=True)
try:
    import comfy.sd
    import comfy.model_management
    
    print(f"[TEST] Loading model config...", flush=True)
    t4 = time.time()
    
    # Use the GGMLOps
    GGMLOps = gguf_ops.GGMLOps
    
    # Try to load via comfy's model loading
    import comfy.supported_models
    import comfy.model_patcher
    
    # Detect model config from state dict
    print(f"[TEST] Detecting model type...", flush=True)
    model_config = comfy.model_detection.model_config_from_unet(sd, "")
    print(f"[TEST] Model config: {type(model_config).__name__}", flush=True)
    
    # Load into model skeleton
    print(f"[TEST] Loading into model skeleton...", flush=True)
    t5 = time.time()
    model = model_config.get_model(sd, "", device=comfy.model_management.unet_offload_device())
    model.model_config = model_config
    print(f"[TEST] Model skeleton in {time.time()-t5:.1f}s", flush=True)
    
    # Set operations
    print(f"[TEST] Setting model operations...", flush=True)
    ops = GGMLOps()
    model.model.diffusion_model = comfy.ops.load_model_gpu(model.model.diffusion_model, ops.__class__)
    
    # Load state dict
    print(f"[TEST] Loading state dict into model...", flush=True)
    t6 = time.time()
    model.model.diffusion_model.load_state_dict(sd, strict=False)
    dt = time.time() - t6
    print(f"[TEST] State dict loaded in {dt:.1f}s", flush=True)
    
    print(f"[TEST] Total model load: {time.time()-t4:.1f}s", flush=True)
    
except Exception as e:
    print(f"[TEST] Error in Step 3: {type(e).__name__}: {e}", flush=True)
    import traceback
    traceback.print_exc()

print(f"\n[TEST] COMPLETE in {time.time()-t0:.1f}s total", flush=True)
os._exit(0)