import os
import pickle
from .feature_extractor import FeatureExtractor
import time
from tqdm import tqdm

def precompute_embeddings():
    # Use absolute paths for Hugging Face Spaces
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    image_dir = os.path.join(base_dir, 'data', 'images')
    output_path = os.path.join(base_dir, 'data', 'embeddings.pkl')
    
    # Create directories if they don't exist
    os.makedirs(image_dir, exist_ok=True)
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    # Rest of your existing code...
    extractor = FeatureExtractor()
    embeddings = []
    image_paths = []

    valid_images = [f for f in os.listdir(image_dir) 
                   if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    total_images = len(valid_images)
    
    print(f"\nFound {total_images} images to process")
    
    start_time = time.time()
    for idx, filename in enumerate(tqdm(valid_images, desc="Processing images")):
        img_path = os.path.join(image_dir, filename)
        try:
            embedding = extractor.extract_features(img_path)
            embeddings.append(embedding)
            image_paths.append(img_path)
        except Exception as e:
            print(f"\nError processing {filename}: {e}")

    with open(output_path, 'wb') as f:
        pickle.dump({'embeddings': embeddings, 'image_paths': image_paths}, f)

    print(f"\nProcessing complete!")
    print(f"Successfully processed {len(embeddings)}/{total_images} images")
    
    return embeddings, image_paths