#!/usr/bin/env python3
"""
BC-250: Test CLIP with Gemma2 key fix — verify correct model detection.
"""
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")

def timeout_handler(sig, frame):
    print("\n[T] === TIMEOUT HIT ===", flush=True)
    os._exit(1)
signal.signal(signal.SIGALRM, timeout_handler)

print("[T] Importing...", flush=True)
import bc250_softmax_patch
import torch
import comfy.sd
import comfy.model_management

# Force CLIP to CPU to avoid GPU kernel compilation delays
comfy.model_management.text_encoder_device = lambda: torch.device("cpu")
comfy.model_management.text_encoder_offload_device = lambda: torch.device("cpu")

clip_path = "/home/fabian/ComfyUI/models/text_encoders/gemma2_2b_lumina2.safetensors"

# Load CLIP
print(f"[T] Loading CLIP...", flush=True)
t0 = time.time()
clip = comfy.sd.load_clip(
    ckpt_paths=[clip_path],
    embedding_directory=None,
    clip_type=comfy.sd.CLIPType.LUMINA2,
)
dt = time.time() - t0
print(f"[T] CLIP loaded in {dt:.1f}s", flush=True)
print(f"[T] CLIP type: {type(clip)}", flush=True)
print(f"[T] cond_stage_model type: {type(clip.cond_stage_model)}", flush=True)

# Check if it's Gemma2 now
csm = clip.cond_stage_model
print(f"[T] Has gemma2_2b attr: {hasattr(csm, 'gemma2_2b')}", flush=True)

# List attributes
attrs = [a for a in dir(csm) if not a.startswith('_') and not callable(getattr(csm, a, None))]
print(f"[T] CSM attrs (non-callable): {attrs[:15]}", flush=True)

# Tokenize
text = "a photo of a cat sitting on a windowsill"
print(f"\n[T] Tokenizing: '{text}'", flush=True)
t1 = time.time()
tokens = clip.tokenize(text)
dt = time.time() - t1
print(f"[T] Tokenized in {dt:.3f}s", flush=True)
print(f"[T] Token keys: {list(tokens.keys()) if isinstance(tokens, dict) else type(tokens)}", flush=True)
for k, v in tokens.items():
    if isinstance(v, list):
        for j, item in enumerate(v[:2]):
            if isinstance(item, list):
                print(f"[T]   {k}[{j}]: list len={len(item)}", flush=True)
            elif hasattr(item, 'shape'):
                print(f"[T]   {k}[{j}]: shape={item.shape}", flush=True)
            else:
                print(f"[T]   {k}[{j}]: {type(item)}", flush=True)
    elif hasattr(v, 'shape'):
        print(f"[T]   {k}: shape={v.shape}", flush=True)
    else:
        print(f"[T]   {k}: {type(v)}", flush=True)

# Encode with 180s timeout (Gemma-2 2B on CPU = slow!)
print(f"\n[T] Encoding (180s timeout)...", flush=True)
signal.alarm(180)
t2 = time.time()
try:
    output = clip.encode_from_tokens_scheduled(tokens)
    dt = time.time() - t2
    signal.alarm(0)
    print(f"[T] Encoded in {dt:.1f}s", flush=True)
    
    if isinstance(output, (list, tuple)):
        for i, item in enumerate(output):
            if isinstance(item, (list, tuple)):
                print(f"[T]   [{i}]: list/tuple len={len(item)}", flush=True)
                if len(item) > 0 and isinstance(item[0], dict):
                    for k, v in item[0].items():
                        if hasattr(v, 'shape'):
                            print(f"[T]     [{i}][0]['{k}']: shape={v.shape} dtype={v.dtype}", flush=True)
                        else:
                            print(f"[T]     [{i}][0]['{k}']: {type(v)} = {v}", flush=True)
                elif len(item) > 0 and hasattr(item[0], 'shape'):
                    print(f"[T]     [{i}][0]: shape={item[0].shape} dtype={item[0].dtype}", flush=True)
            elif hasattr(item, 'shape'):
                print(f"[T]   [{i}]: shape={item.shape} dtype={item.dtype}", flush=True)
            else:
                print(f"[T]   [{i}]: {type(item)}", flush=True)
    elif isinstance(output, dict):
        for k, v in output.items():
            if hasattr(v, 'shape'):
                print(f"[T]   {k}: shape={v.shape} dtype={v.dtype}", flush=True)
            else:
                print(f"[T]   {k}: {type(v)}", flush=True)
    
    print(f"\n[T] === CLIP ENCODE SUCCESS ===", flush=True)
except Exception as e:
    signal.alarm(0)
    print(f"[T] ERROR: {e}", flush=True)
    import traceback
    traceback.print_exc()

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