"""
Final comprehensive fix for ChromaDB 0.3.29 compatibility

This ensures ChromaDB works correctly with both direct usage and MCP server.
"""

import chromadb
from chromadb.config import Settings
from chromadb.api import API
import time


def create_working_chromadb_client(persist_directory=".mira/chromadb"):
    """Create a properly initialized ChromaDB 0.3.29 client."""
    
    # Ensure ChromaDB is properly initialized
    settings = Settings(
        chroma_db_impl="duckdb+parquet",
        persist_directory=persist_directory,
        anonymized_telemetry=False
    )
    
    # Create client
    client = chromadb.Client(settings)
    
    # IMPORTANT: For ChromaDB 0.3.29, we need to ensure the client is started
    if hasattr(client, '_db') and client._db is None:
        # Initialize the database backend
        client.start()
    
    return client


def test_comprehensive_fix():
    """Test the comprehensive ChromaDB fix."""
    print("Testing Comprehensive ChromaDB 0.3.29 Fix")
    print("="*60)
    
    # Test 1: Create client with proper initialization
    print("\n1. Creating properly initialized client...")
    client = create_working_chromadb_client(".mira/chromadb_final_test")
    print("✅ Client created")
    
    # Test 2: Test get_or_create_collection
    print("\n2. Testing get_or_create_collection...")
    collection = client.get_or_create_collection(
        "test_final_collection",
        metadata={"purpose": "final_test", "version": "1.0"}
    )
    print(f"✅ Collection created: {collection}")
    print(f"   Type: {type(collection)}")
    print(f"   Name: {collection.name if collection else 'None'}")
    
    # Test 3: Test collection operations
    if collection:
        print("\n3. Testing collection operations...")
        
        # Add document
        collection.add(
            documents=["Final test document"],
            metadatas=[{"type": "final_test", "timestamp": str(time.time())}],
            ids=["final_test_1"]
        )
        print("✅ Document added")
        
        # Get document
        results = collection.get(ids=["final_test_1"])
        print(f"✅ Document retrieved: {results['documents'][0]}")
        
        # Count
        count = collection.count()
        print(f"✅ Collection count: {count}")
    
    # Test 4: Test with MIRA integration
    print("\n4. Testing MIRA integration...")
    try:
        from core.storage.chroma_client import MIRAChromaClient
        
        # Patch MIRAChromaClient to use our initialization
        original_init = MIRAChromaClient.__init__
        
        def patched_init(self):
            """Patched initialization for ChromaDB 0.3.29."""
            # Ensure ChromaDB files are stored within MIRA's data directory
            self.chroma_path = Path(".mira/chromadb")
            self.chroma_path.mkdir(parents=True, exist_ok=True)
            
            logger.info(f"Initializing ChromaDB at {self.chroma_path}")
            
            # Use our working client creation
            self.client = create_working_chromadb_client(str(self.chroma_path))
            
            # Initialize core collections on startup
            self.initialize_collections()
        
        # Apply patch
        from pathlib import Path
        import logging
        logger = logging.getLogger(__name__)
        MIRAChromaClient.__init__ = patched_init
        
        # Test MIRA client
        mira_client = MIRAChromaClient()
        collections = mira_client.list_collections()
        print(f"✅ MIRA client working with {len(collections)} collections")
        
    except Exception as e:
        print(f"❌ MIRA integration error: {e}")
    
    # Cleanup
    try:
        client.delete_collection("test_final_collection")
        print("\n✅ Cleanup successful")
    except:
        pass
    
    print("\n" + "="*60)
    print("Comprehensive fix complete!")
    return True


if __name__ == "__main__":
    # Apply the comprehensive fix
    success = test_comprehensive_fix()
    
    if success:
        print("\n✅ ChromaDB 0.3.29 is now fully compatible!")
        print("\nTo use in your code:")
        print("1. Import: from fix_chromadb_final import create_working_chromadb_client")
        print("2. Create client: client = create_working_chromadb_client()")
        print("3. Use normally: collection = client.get_or_create_collection(...)")