#!/usr/bin/env python3
"""
BC-250: Test each ComfyUI component in isolation to find which one hangs.
Run from /home/fabian/ComfyUI with venv active.
"""
import os, sys, time, signal

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")

# Timeout handler
def timeout_handler(signum, frame):
    print(f"\n[TIMEOUT] Operation exceeded time limit!", flush=True)
    os._exit(1)

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

print("[T] Importing comfy...", flush=True)
t0 = time.time()
import torch
import comfy.sd
import comfy.model_management
import comfy.utils
import comfy.clip_model
import folder_paths
print(f"[T] Imports done in {time.time()-t0:.1f}s", flush=True)
print(f"[T] CUDA: {torch.cuda.is_available()}, Device: {torch.cuda.get_device_name(0)}", flush=True)

# === TEST 1: Load CLIP model ===
print(f"\n{'='*60}", flush=True)
print(f"[T] TEST 1: Load CLIP model (gemma2_2b_lumina2)", flush=True)
signal.alarm(60)  # 60s timeout
t1 = time.time()
try:
    clip_path = os.path.join(folder_paths.get_folder_paths("clip")[0], "gemma2_2b_lumina2.safetensors")
    if not os.path.exists(clip_path):
        # Try text_encoders folder
        for p in folder_paths.get_folder_paths("text_encoders"):
            cp = os.path.join(p, "gemma2_2b_lumina2.safetensors")
            if os.path.exists(cp):
                clip_path = cp
                break
    
    print(f"[T] CLIP path: {clip_path}", flush=True)
    print(f"[T] Loading CLIP...", flush=True)
    
    clip = comfy.sd.load_clip(
        ckpt_paths=[clip_path],
        embedding_directory=None,
        clip_type=comfy.sd.CLIPType.LUMINA2,
    )
    dt = time.time() - t1
    print(f"[T] CLIP loaded in {dt:.1f}s", flush=True)
    print(f"[T] CLIP type: {type(clip).__name__}", flush=True)
    
    # Check GPU memory after CLIP load
    print(f"[T] GPU VRAM after CLIP load:", flush=True)
    print(f"[T]   allocated: {torch.cuda.memory_allocated()/1e6:.1f} MB", flush=True)
    print(f"[T]   reserved: {torch.cuda.memory_reserved()/1e6:.1f} MB", flush=True)
    
except Exception as e:
    print(f"[T] TEST 1 ERROR: {e}", flush=True)
    import traceback
    traceback.print_exc()
    clip = None

signal.alarm(0)

# === TEST 2: Run CLIP text encoding ===
if clip is not None:
    print(f"\n{'='*60}", flush=True)
    print(f"[T] TEST 2: CLIP text encoding", flush=True)
    signal.alarm(120)  # 120s timeout
    t2 = time.time()
    try:
        print(f"[T] Encoding: 'a cat'...", flush=True)
        tokens = clip.tokenize({"g": "a cat"})
        print(f"[T] Tokenized in {time.time()-t2:.3f}s", flush=True)
        
        t2b = time.time()
        print(f"[T] Running CLIP encode (this is the suspected hang point)...", flush=True)
        output = clip.encode_from_tokens_scheduled(tokens)
        cond, pooled = output[:2]
        dt = time.time() - t2b
        print(f"[T] CLIP encoded in {dt:.1f}s", flush=True)
        print(f"[T] Cond shape: {cond.shape}, dtype: {cond.dtype}", flush=True)
        
    except Exception as e:
        print(f"[T] TEST 2 ERROR: {e}", flush=True)
        import traceback
        traceback.print_exc()
    signal.alarm(0)

# === TEST 3: Load GGUF UNet ===
print(f"\n{'='*60}", flush=True)
print(f"[T] TEST 3: Load GGUF UNet", flush=True)
signal.alarm(60)
t3 = time.time()
try:
    # Ensure GGUF patch is applied
    bc250_softmax_patch._try_patch_gguf()
    
    # Import GGUF nodes
    gguf_path = "/home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF"
    sys.path.insert(0, gguf_path)
    
    # Use the loader directly
    from loader import gguf_sd_loader
    from ops import GGMLOps
    
    unet_path = "/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf"
    print(f"[T] Loading GGUF state dict...", flush=True)
    sd, extra = gguf_sd_loader(unet_path)
    print(f"[T] State dict: {len(sd)} keys, arch={extra.get('arch_str')}", flush=True)
    
    # Now load through comfy
    print(f"[T] Creating diffusion model...", flush=True)
    ops = GGMLOps()
    model = comfy.sd.load_diffusion_model_state_dict(
        sd, model_options={"custom_operations": ops},
        metadata=extra.get("metadata", {}),
    )
    dt = time.time() - t3
    print(f"[T] UNet loaded in {dt:.1f}s", flush=True)
    if model is not None:
        print(f"[T] Model type: {type(model).__name__}", flush=True)
    else:
        print(f"[T] WARNING: model is None!", flush=True)
    
except Exception as e:
    print(f"[T] TEST 3 ERROR: {e}", flush=True)
    import traceback
    traceback.print_exc()
signal.alarm(0)

print(f"\n{'='*60}", flush=True)
print(f"[T] ALL TESTS COMPLETE in {time.time()-t0:.1f}s", flush=True)
os._exit(0)