import torch
import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
import open_clip
from PIL import Image
import threading

_model = None
_preprocess = None
_device = None
_lock = threading.Lock()
_gpu_lock = threading.Lock()

def get_model():
    global _model, _preprocess, _device
    if _model is None:
        with _lock:
            if _model is None:
                dev = "cuda" if torch.cuda.is_available() else "cpu"
                # ViT-H-14 is 1024-dim
                model, _, preprocess = open_clip.create_model_and_transforms('ViT-H-14', pretrained='laion2b_s32b_b79k')
                model = model.to(dev)
                model.eval()
                
                # Set globals only when fully ready to avoid race conditions
                _preprocess = preprocess
                _device = dev
                _model = model
    return _model, _preprocess, _device

def generate_embedding(image_path):
    model, preprocess, device = get_model()
    try:
        with _gpu_lock:
            with Image.open(image_path) as img:
                image = preprocess(img.convert("RGB")).unsqueeze(0).to(device)
                with torch.no_grad():
                    image_features = model.encode_image(image)
                    image_features /= image_features.norm(dim=-1, keepdim=True)
                    return image_features.cpu().numpy()[0].tolist()
    except Exception as e:
        print(f"Error generating embedding for {image_path}: {e}")
        return None