from cora_vision import CoraVision
from cora_memory import CoraMemory
from PIL import Image
import os

def create_dummy_image():
    img = Image.new('RGB', (256, 256), color = 'red')
    img.save('test_red.png')
    return 'test_red.png'

def test_stack():
    # 1. Init
    print("--- Initializing Stack ---")
    vision = CoraVision()
    memory = CoraMemory()

    # 2. Get Image
    img_path = create_dummy_image()
    prompt = "A red square test image"
    print(f"Created dummy image: {img_path}")

    # 3. Vision Analysis
    print("--- 1. The Eyes (Vision) ---")
    embedding = vision.embed_image(img_path)
    print(f"Embedding generated: Vector length {len(embedding)}")
    
    tags = vision.detect_tags(img_path)
    print(f"Tags detected: {tags}")

    # 4. Memory Storage
    print("--- 2. The Memory (Storage) ---")
    entry_id = memory.save(img_path, embedding, prompt, tags)
    print(f"Stored with ID: {entry_id}")

    # 5. Search
    print("--- 3. Retrieval (Search) ---")
    # Search by the same image embedding (should find itself)
    results = memory.search_by_vector(embedding, k=1)
    
    print("Search Results:")
    try:
        # ChromaDB API structure v0.4+
        ids = results['ids'][0]
        metadatas = results['metadatas'][0]
        distances = results['distances'][0]
        
        for i, uid in enumerate(ids):
            print(f"Found: {uid}")
            print(f"  Metadata: {metadatas[i]}")
            print(f"  Distance: {distances[i]}")
            
        if entry_id in ids:
            print("SUCCESS: Retrieved the saved entry!")
        else:
            print("FAILURE: Did not retrieve the entry.")
            
    except Exception as e:
        print(f"Error parsing results: {e}")
        print(f"Raw results: {results}")

    # Clean up
    if os.path.exists(img_path):
        os.remove(img_path)

if __name__ == "__main__":
    test_stack()