#!/usr/bin/env python3
"""
BC-250: Test CLIP encoding on CPU — find out if Gemma-2 2B works.
"""
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 text encoder to CPU
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"
fsize = os.path.getsize(clip_path) / (1024*1024*1024)
print(f"[T] CLIP: {fsize:.2f} GB", flush=True)

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

# Check the clip object
print(f"[T] CLIP type: {type(clip)}", flush=True)
print(f"[T] CLIP cond_stage_model type: {type(clip.cond_stage_model)}", flush=True)

# Tokenize with just a string
text = "a photo of a cat"
print(f"[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)

# Encode with 120s timeout
print(f"[T] Encoding (120s timeout)...", flush=True)
signal.alarm(120)
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, 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)
    elif isinstance(output, (list, tuple)):
        for i, v in enumerate(output):
            if hasattr(v, 'shape'):
                print(f"[T]   [{i}]: shape={v.shape} dtype={v.dtype}", flush=True)
            else:
                print(f"[T]   [{i}]: {type(v)}", flush=True)
    
    print(f"\n[T] === CLIP ENCODE ON CPU: 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)