import gradio as gr
import os
import faiss
import numpy as np
from pathlib import Path

# Constants
DATA_DIR = "data"
INDEX_PATH = os.path.join(DATA_DIR, "faiss_index.index")

def save_index(file_obj):
    """
    Save uploaded FAISS index to the data directory
    """
    # Create data directory if it doesn't exist
    Path(DATA_DIR).mkdir(exist_ok=True)
    
    # Check if index already exists
    if os.path.exists(INDEX_PATH):
        return "⚠️ A FAISS index already exists in the data directory. Please remove it first."
    
    try:
        # Save the uploaded file
        with open(INDEX_PATH, 'wb') as f:
            f.write(file_obj)
        
        # Verify the saved file is a valid FAISS index
        faiss.read_index(INDEX_PATH)
        return "✅ FAISS index successfully uploaded and saved!"
        
    except Exception as e:
        # If there was an error, remove the file if it was created
        if os.path.exists(INDEX_PATH):
            os.remove(INDEX_PATH)
        return f"❌ Error: Invalid FAISS index file - {str(e)}"

# Create Gradio interface
demo = gr.Interface(
    fn=save_index,
    inputs=gr.File(label="Upload FAISS Index", file_types=[".index"]),
    outputs=gr.Textbox(label="Status"),
    title="FAISS Index Uploader",
    description="Upload a FAISS index file to store in the HuggingFace Space data directory.",
)

if __name__ == "__main__":
    demo.launch()