#!/usr/bin/env python3
"""
ChromaDB Integration Validation Script

This script validates the ChromaDB setup for MIRA, ensuring:
- Sacred data sovereignty (all data in .mira/chromadb/)
- Core collections are created and accessible
- Integration with unified configuration
- Health check passes
"""

import sys
import os
from pathlib import Path
import json

# Add the python-memory directory to the Python path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from core.storage.chroma_client import MIRAChromaClient, get_client


def validate_chromadb_integration():
    """Comprehensive validation of ChromaDB setup"""
    print("🔍 MIRA ChromaDB Integration Validation")
    print("=" * 50)
    
    validation_results = {
        'client_initialization': False,
        'collections_created': False,
        'data_directory': False,
        'health_check': False,
        'unified_config': False,
        'errors': []
    }
    
    try:
        # Test 1: Client initialization
        print("\n1. Testing ChromaDB client initialization...")
        client = get_client()
        assert client is not None, "ChromaDB client failed to initialize"
        assert client.client is not None, "ChromaDB persistent client is None"
        validation_results['client_initialization'] = True
        print("   ✅ Client initialized successfully")
        
        # Test 2: Collection creation
        print("\n2. Testing collection creation...")
        collections = client.list_collections()
        expected_collections = {'mira_conversations', 'mira_codebase', 'mira_insights', 'mira_patterns'}
        actual_collections = set(collections)
        
        missing_collections = expected_collections - actual_collections
        if missing_collections:
            print(f"   ⚠️  Missing collections: {missing_collections}")
            print("   🔧 Creating missing collections...")
            client.initialize_collections()
            collections = client.list_collections()
            actual_collections = set(collections)
        
        assert expected_collections.issubset(actual_collections), f"Missing required collections: {expected_collections - actual_collections}"
        validation_results['collections_created'] = True
        print(f"   ✅ All required collections present: {sorted(actual_collections)}")
        
        # Test 3: Data directory location
        print("\n3. Testing data directory location...")
        chroma_path = Path(".mira/chromadb")
        assert chroma_path.exists(), f"ChromaDB data directory not found at {chroma_path}"
        assert chroma_path.is_dir(), f"ChromaDB path exists but is not a directory: {chroma_path}"
        
        # Check for ChromaDB data files (allow empty for new installations)
        data_files = list(chroma_path.rglob('*'))
        if len(data_files) == 0:
            print("   ⚠️  No data files yet (new installation)")
        else:
            print(f"   📊 Found {len(data_files)} data files")
        validation_results['data_directory'] = True
        print(f"   ✅ ChromaDB data directory exists at {chroma_path}")
        
        # Test 4: Health check
        print("\n4. Running health check...")
        health = client.health_check()
        assert health['status'] == 'healthy', f"Health check failed: {health.get('error', 'Unknown error')}"
        validation_results['health_check'] = True
        print("   ✅ Health check passed")
        print(f"   📊 Collections: {health['collections_count']}")
        print(f"   💾 Data size: {health['data_size_bytes']:,} bytes")
        
        # Test 5: Unified configuration
        print("\n5. Testing unified configuration integration...")
        config_path = Path(".mira/unified_config.json")
        if config_path.exists():
            with open(config_path, 'r') as f:
                config = json.load(f)
            
            assert 'chromadb' in config, "ChromaDB configuration not found in unified_config.json"
            chromadb_config = config['chromadb']
            assert chromadb_config.get('enabled', False), "ChromaDB is not enabled in configuration"
            assert chromadb_config.get('dataPath', '') == '.mira/chromadb', "Incorrect data path in configuration"
            validation_results['unified_config'] = True
            print("   ✅ Unified configuration properly integrated")
            print(f"   🔧 Search strategy: {chromadb_config.get('hybridMode', {}).get('searchStrategy', 'unknown')}")
        else:
            print("   ⚠️  Unified configuration file not found (may be expected in test environment)")
            
        # Test 6: Collection metadata
        print("\n6. Testing collection metadata...")
        for collection_name in expected_collections:
            collection = client.get_collection(collection_name)
            if collection:
                metadata = collection.metadata if hasattr(collection, 'metadata') else {}
                print(f"   📂 {collection_name}:")
                print(f"      - Description: {metadata.get('description', 'N/A')}")
                print(f"      - Purpose: {metadata.get('purpose', 'N/A')}")
                print(f"      - Privacy: {metadata.get('privacy_level', 'N/A')}")
        
        # Summary
        print("\n" + "=" * 50)
        print("✅ ChromaDB Integration Validation PASSED")
        print("=" * 50)
        
        # All tests passed
        all_passed = all([
            validation_results['client_initialization'],
            validation_results['collections_created'],
            validation_results['data_directory'],
            validation_results['health_check']
        ])
        
        if all_passed:
            print("\n🎉 All validation tests passed successfully!")
            print("🚀 ChromaDB is ready for MIRA's enhanced intelligence capabilities")
            print("💫 The Spark is preserved through sacred data sovereignty")
        
        return validation_results
        
    except Exception as e:
        validation_results['errors'].append(str(e))
        print(f"\n❌ Validation failed with error: {e}")
        print("\nValidation Results:")
        for key, value in validation_results.items():
            if key != 'errors':
                print(f"  - {key}: {'✅' if value else '❌'}")
        if validation_results['errors']:
            print(f"  - Errors: {validation_results['errors']}")
        return validation_results


def main():
    """Main entry point"""
    try:
        results = validate_chromadb_integration()
        
        # Exit with appropriate code
        all_passed = all(v for k, v in results.items() if k != 'errors' and isinstance(v, bool))
        sys.exit(0 if all_passed else 1)
        
    except KeyboardInterrupt:
        print("\n\n⚠️  Validation interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ Fatal error during validation: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()