#!/usr/bin/env python3
"""
BC-250: Quick test - load GGUF into model via ComfyUI's actual pipeline.
Run from ComfyUI directory with venv active.
"""
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("[T] Importing patch...", flush=True)
import bc250_softmax_patch

print("[T] Importing comfy modules...", flush=True)
t0 = time.time()

import torch
import comfy.sd
import comfy.model_management
import comfy.utils

print(f"[T] Imports done in {time.time()-t0:.1f}s", flush=True)

# Load the GGUF model using ComfyUI's own loading
UNET_PATH = "/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf"

print(f"\n[T] Loading GGUF unet...", flush=True)
t1 = time.time()

try:
    # ComfyUI-GGUF has a hyphen, need to import differently
    import importlib
    gguf_nodes_path = "/home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF"
    if gguf_nodes_path not in sys.path:
        sys.path.insert(0, gguf_nodes_path)
    
    # Re-check if GGUF patch got applied after import  
    import nodes as gguf_nodes_mod
    import ops as gguf_ops_mod
    import dequant as gguf_dequant_mod
    
    # Register the dequant module so our patch can find it
    sys.modules['custom_nodes.ComfyUI_GGUF.dequant'] = gguf_dequant_mod
    sys.modules['custom_nodes.ComfyUI_GGUF.ops'] = gguf_ops_mod
    
    bc250_softmax_patch._try_patch_gguf()
    
    # Now load the model manually
    from loader import gguf_sd_loader
    from ops import GGMLOps, GGMLTensor
    
    print(f"[T] Loading GGUF state dict...", flush=True)
    t2 = time.time()
    sd, extra = gguf_sd_loader(UNET_PATH)
    print(f"[T] State dict loaded in {time.time()-t2:.1f}s ({len(sd)} keys, arch={extra.get('arch_str')})", flush=True)
    
    # Now detect model config
    import comfy.model_detection
    print(f"[T] Detecting model config...", flush=True)
    t3 = time.time()
    
    # model_config_from_unet needs the state dict
    parameters = comfy.utils.calculate_parameters(sd)
    print(f"[T] Parameters: {parameters/1e9:.2f}B", flush=True)
    
    unet_dtype = torch.float16
    load_device = comfy.model_management.get_torch_device()
    print(f"[T] Load device: {load_device}", flush=True)
    
    model_config = comfy.model_detection.model_config_from_unet(sd, "")
    print(f"[T] Model config: {type(model_config).__name__}", flush=True)
    
    if model_config is None:
        print(f"[T] ERROR: Could not detect model config!", flush=True)
    else:
        model_config.custom_operations = GGMLOps
        
        print(f"[T] Creating model...", flush=True)
        t4 = time.time()
        model = model_config.get_model(sd, "", device=comfy.model_management.unet_offload_device())
        print(f"[T] Model created in {time.time()-t4:.1f}s", flush=True)
        
        print(f"[T] Loading state dict into model...", flush=True)
        t5 = time.time()
        model.load_model_weights(sd, "")
        print(f"[T] Weights loaded in {time.time()-t5:.1f}s", flush=True)
        
        print(f"\n[T] === UNET LOADED SUCCESSFULLY ===", flush=True)
    
except Exception as e:
    print(f"[T] ERROR: {type(e).__name__}: {e}", flush=True)
    import traceback
    traceback.print_exc()

print(f"[T] Total: {time.time()-t0:.1f}s", flush=True)
os._exit(0)