#!/usr/bin/env python3
"""
Diagnostic tool to check LightRAG initialization status.

This tool helps developers verify that their LightRAG instance is properly 
initialized before use, preventing common initialization errors.

Usage:
    python -m lightrag.tools.check_initialization
"""

import asyncio
import sys
from typing import Optional
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from lightrag import LightRAG
from lightrag.base import StoragesStatus


async def check_lightrag_setup(rag_instance: LightRAG, verbose: bool = False) -> bool:
    """
    Check if a LightRAG instance is properly initialized.
    
    Args:
        rag_instance: The LightRAG instance to check
        verbose: If True, print detailed diagnostic information
        
    Returns:
        True if properly initialized, False otherwise
    """
    issues = []
    warnings = []
    
    print("🔍 Checking LightRAG initialization status...\n")
    
    # Check storage initialization status
    if not hasattr(rag_instance, '_storages_status'):
        issues.append("LightRAG instance missing _storages_status attribute")
    elif rag_instance._storages_status != StoragesStatus.INITIALIZED:
        issues.append(f"Storages not initialized (status: {rag_instance._storages_status.name})")
    else:
        print("✅ Storage status: INITIALIZED")
    
    # Check individual storage components
    storage_components = [
        ('full_docs', 'Document storage'),
        ('text_chunks', 'Text chunks storage'),
        ('entities_vdb', 'Entity vector database'),
        ('relationships_vdb', 'Relationship vector database'),
        ('chunks_vdb', 'Chunks vector database'),
        ('doc_status', 'Document status tracker'),
        ('llm_response_cache', 'LLM response cache'),
        ('full_entities', 'Entity storage'),
        ('full_relations', 'Relation storage'),
        ('chunk_entity_relation_graph', 'Graph storage')
    ]
    
    if verbose:
        print("\n📦 Storage Components:")
    
    for component, description in storage_components:
        if not hasattr(rag_instance, component):
            issues.append(f"Missing storage component: {component} ({description})")
        else:
            storage = getattr(rag_instance, component)
            if storage is None:
                warnings.append(f"Storage {component} is None (might be optional)")
            elif hasattr(storage, '_storage_lock'):
                if storage._storage_lock is None:
                    issues.append(f"Storage {component} not initialized (lock is None)")
                elif verbose:
                    print(f"  ✅ {description}: Ready")
            elif verbose:
                print(f"  ✅ {description}: Ready")
    
    # Check pipeline status
    try:
        from lightrag.kg.shared_storage import get_namespace_data
        get_namespace_data("pipeline_status")
        print("✅ Pipeline status: INITIALIZED")
    except KeyError:
        issues.append("Pipeline status not initialized - call initialize_pipeline_status()")
    except Exception as e:
        issues.append(f"Error checking pipeline status: {str(e)}")
    
    # Print results
    print("\n" + "=" * 50)
    
    if issues:
        print("❌ Issues found:\n")
        for issue in issues:
            print(f"  • {issue}")
        
        print("\n📝 To fix, run this initialization sequence:\n")
        print("  await rag.initialize_storages()")
        print("  from lightrag.kg.shared_storage import initialize_pipeline_status")
        print("  await initialize_pipeline_status()")
        print("\n📚 Documentation: https://github.com/HKUDS/LightRAG#important-initialization-requirements")
        
        if warnings and verbose:
            print("\n⚠️  Warnings (might be normal):")
            for warning in warnings:
                print(f"  • {warning}")
        
        return False
    else:
        print("✅ LightRAG is properly initialized and ready to use!")
        
        if warnings and verbose:
            print("\n⚠️  Warnings (might be normal):")
            for warning in warnings:
                print(f"  • {warning}")
        
        return True


async def demo():
    """Demonstrate the diagnostic tool with a test instance."""
    from lightrag.llm.openai import openai_embed, gpt_4o_mini_complete
    from lightrag.kg.shared_storage import initialize_pipeline_status
    
    print("=" * 50)
    print("LightRAG Initialization Diagnostic Tool")
    print("=" * 50)
    
    # Create test instance
    rag = LightRAG(
        working_dir="./test_diagnostic",
        embedding_func=openai_embed,
        llm_model_func=gpt_4o_mini_complete,
    )
    
    print("\n🔴 BEFORE initialization:\n")
    await check_lightrag_setup(rag, verbose=True)
    
    print("\n" + "=" * 50)
    print("\n🔄 Initializing...\n")
    await rag.initialize_storages()
    await initialize_pipeline_status()
    
    print("\n🟢 AFTER initialization:\n")
    await check_lightrag_setup(rag, verbose=True)
    
    # Cleanup
    import shutil
    shutil.rmtree("./test_diagnostic", ignore_errors=True)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Check LightRAG initialization status"
    )
    parser.add_argument(
        "--demo", 
        action="store_true",
        help="Run a demonstration with a test instance"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show detailed diagnostic information"
    )
    
    args = parser.parse_args()
    
    if args.demo:
        asyncio.run(demo())
    else:
        print("Run with --demo to see the diagnostic tool in action")
        print("Or import this module and use check_lightrag_setup() with your instance")