import os
import chromadb
from dotenv import load_dotenv

load_dotenv()

def reset_chroma_db(persist_directory="chroma_db", collection_name="mangoit_docs"):
    """
    Reset the ChromaDB database by deleting and recreating the collection
    
    Args:
        persist_directory: Directory where ChromaDB is persisted
        collection_name: Name of the collection to reset
        
    Returns:
        bool: True if reset was successful, False otherwise
    """
    try:
        print(f"Connecting to ChromaDB at {persist_directory}...")
        client = chromadb.PersistentClient(path=persist_directory)
        
        # Get all collections
        collections = client.list_collections()
        collection_names = [collection.name for collection in collections]
        
        print(f"Found collections: {collection_names}")
        
        # Check if our collection exists
        if collection_name in collection_names:
            print(f"Deleting collection '{collection_name}'...")
            client.delete_collection(collection_name)
            print(f"Collection '{collection_name}' deleted successfully")
        else:
            print(f"Collection '{collection_name}' does not exist, nothing to delete")
        
        print("Database reset complete!")
        return True
        
    except Exception as e:
        print(f"Error resetting database: {str(e)}")
        return False

if __name__ == "__main__":
    # Reset the default collection
    reset_chroma_db()
    
    # Optionally reset other collections
    # If you have other collections, uncomment and modify the line below
    # reset_chroma_db(collection_name="other_collection")