src/utils/database_cleanup.py

# src/utils/database_cleanup.py
import chromadb
import shutil
from pathlib import Path
import asyncio
import gc
import random
from typing import List, Dict, Tuple
from src.utils.logger import logger
from config.config import settings


async def cleanup_chroma() -> Tuple[List[str], bool]:
    """Clean up ChromaDB data while maintaining connection"""
    details = []
    restart_needed = False

    try:
        # Get existing client
        client = chromadb.PersistentClient(
            path=settings.CHROMA_PATH,
            settings=chromadb.Settings(
                allow_reset=True,
                is_persistent=True,
                anonymized_telemetry=False
            )
        )

        # Get all collections
        collections = client.list_collections()

        if not collections:
            details.append("No collections found in ChromaDB")
            return details, restart_needed

        # Delete data from each collection
        for collection in collections:
            try:
                # Get all IDs in the collection
                all_ids = collection.get()['ids']

                if all_ids:
                    # Delete all documents in the collection
                    collection.delete(ids=all_ids)
                    details.append(
                        f"Deleted {len(all_ids)} documents from collection {collection.name}")
                else:
                    details.append(
                        f"Collection {collection.name} was already empty")

                # Delete the collection itself
                client.delete_collection(collection.name)
                details.append(f"Deleted collection {collection.name}")

            except Exception as e:
                logger.warning(
                    f"Error cleaning collection {collection.name}: {str(e)}")
                details.append(
                    f"Error cleaning collection {collection.name}: {str(e)}")
                restart_needed = True  # Set restart flag if any collection fails

        # Optional: Add a check to see if a full reset might be necessary
        if len(client.list_collections()) > 0:
            restart_needed = True
            details.append("Some collections might require manual reset")

        return details, restart_needed

    except Exception as e:
        raise Exception(f"ChromaDB cleanup failed: {str(e)}")


async def cleanup_mongodb(mongodb) -> List[str]:
    """Clean up MongoDB collections"""
    details = []

    try:
        # Get all collections in the database
        collections = await mongodb.db.list_collection_names()

        # Core collections from MongoDBStore initialization
        core_collections = {
            'chat_history': mongodb.chat_history,
            'conversations': mongodb.conversations,
            'knowledge_base': mongodb.documents,  # documents maps to knowledge_base
            # Direct access to vector_metadata collection
            'vector_metadata': mongodb.db.vector_metadata,
        }

        # Clean each core collection
        for name, collection in core_collections.items():
            try:
                result = await collection.delete_many({})
                details.append(
                    f"Cleared {name} ({result.deleted_count} documents)")
            except Exception as e:
                logger.error(f"Error clearing {name}: {str(e)}")
                details.append(f"Error clearing {name}: {str(e)}")

        # Clean any additional collections not in the core set
        for coll_name in collections:
            if coll_name not in core_collections:
                try:
                    result = await mongodb.db[coll_name].delete_many({})
                    details.append(
                        f"Cleared additional collection {coll_name} ({result.deleted_count} documents)")
                except Exception as e:
                    logger.error(
                        f"Error clearing additional collection {coll_name}: {str(e)}")

        return details
    except Exception as e:
        raise Exception(f"MongoDB cleanup failed: {str(e)}")


async def cleanup_files() -> List[str]:
    """Clean up uploaded files and temporary directories"""
    details = []

    # Directories to clean
    directories = {
        'uploads': Path("uploads"),
        'temp_downloads': Path(settings.TEMP_DOWNLOAD_DIR),
        # Additional temp directory used by some components
        'temp_dir': Path('./temp')
    }

    for dir_name, dir_path in directories.items():
        if dir_path.exists():
            try:
                # Delete all files in the directory
                for file in dir_path.glob('*'):
                    try:
                        if file.is_file():
                            file.unlink()
                            details.append(
                                f"Deleted file: {file.name} from {dir_name}")
                    except Exception as e:
                        details.append(
                            f"Error deleting file {file.name} from {dir_name}: {str(e)}")

                # Try to remove the empty directory
                if not any(dir_path.iterdir()):
                    dir_path.rmdir()
                    details.append(f"Removed empty {dir_name} directory")
            except Exception as e:
                details.append(
                    f"Error cleaning {dir_name} directory: {str(e)}")
        else:
            details.append(f"No {dir_name} directory found")

    return details


async def perform_cleanup(mongodb, include_files: bool = True) -> Dict:
    """
    Perform comprehensive cleanup of all databases and files

    Args:
        mongodb: MongoDB store instance
        include_files (bool): Whether to also delete uploaded files

    Returns:
        Dict: Cleanup operation summary with detailed status
    """
    cleanup_summary = {
        "chroma_db": {"status": "not_started", "details": []},
        "mongodb": {"status": "not_started", "details": []},
        "files": {"status": "not_started", "details": []}
    }

    try:
        # Clean ChromaDB
        try:
            details, restart_needed = await cleanup_chroma()
            cleanup_summary["chroma_db"] = {
                "status": "success" if not restart_needed else "partial",
                "details": details
            }
        except Exception as e:
            logger.error(f"Error cleaning ChromaDB: {str(e)}")
            cleanup_summary["chroma_db"] = {
                "status": "error",
                "details": [str(e)]
            }

        # Clean MongoDB
        try:
            details = await cleanup_mongodb(mongodb)
            cleanup_summary["mongodb"] = {
                "status": "success",
                "details": details
            }
        except Exception as e:
            logger.error(f"Error cleaning MongoDB: {str(e)}")
            cleanup_summary["mongodb"] = {
                "status": "error",
                "details": [str(e)]
            }

        # Clean files if requested
        if include_files:
            try:
                details = await cleanup_files()
                cleanup_summary["files"] = {
                    "status": "success",
                    "details": details
                }
            except Exception as e:
                logger.error(f"Error cleaning files: {str(e)}")
                cleanup_summary["files"] = {
                    "status": "error",
                    "details": [str(e)]
                }

        # Determine overall status
        overall_status = "success"
        if restart_needed:
            overall_status = "partial_success"
            cleanup_summary["message"] = "Cleanup partially completed. Server restart required to complete ChromaDB cleanup."
        elif any(item["status"] == "error" for item in cleanup_summary.values()):
            overall_status = "partial_success"
        elif all(item["status"] == "error" for item in cleanup_summary.values()):
            overall_status = "error"

        return {
            "status": overall_status,
            "message": cleanup_summary.get("message", "Cleanup operation completed"),
            "details": cleanup_summary,
            "restart_needed": restart_needed
        }

    except Exception as e:
        logger.error(f"Error in cleanup operation: {str(e)}")
        raise