"""
Debug script to test LLM function configuration in LightRAG
"""

import asyncio
import sys
import os
import logging

# Add the current directory to the path so we can import LightRAG
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def create_mock_llm_and_embedding():
    """Create mock LLM and embedding functions for testing"""
    
    # Set up mock LLM function
    async def mock_llm_func(prompt, system_prompt=None, **kwargs):
        logger.info(f"Mock LLM called with prompt: {prompt[:100]}...")
        return "Mock LLM response for testing"
    
    # Set up mock embedding function wrapped with EmbeddingFunc
    async def mock_embedding_func(texts):
        logger.info(f"Mock embedding called with {len(texts)} texts")
        import numpy as np
        # Return random embeddings of dimension 384
        return [np.random.rand(384).astype(np.float32) for _ in texts]
    
    llm_func = mock_llm_func
    
    # Wrap the embedding function with EmbeddingFunc class
    from lightrag.utils import EmbeddingFunc
    embedding_func = EmbeddingFunc(
        embedding_dim=384,
        func=mock_embedding_func
    )
    
    return llm_func, embedding_func

async def test_llm_function_setup():
    """Test LLM function setup and LightRAG initialization"""
    
    llm_func, embedding_func = create_mock_llm_and_embedding()
    
    logger.info(f"LLM function type: {type(llm_func)}")
    logger.info(f"LLM function callable: {callable(llm_func)}")
    logger.info(f"Embedding function type: {type(embedding_func)}")
    logger.info(f"Embedding function callable: {callable(embedding_func)}")
    logger.info(f"Embedding function embedding_dim: {embedding_func.embedding_dim}")
    
    try:
        from lightrag import LightRAG
        
        # Test creating LightRAG instance
        logger.info("Creating LightRAG instance...")
        rag = LightRAG(
            working_dir="./debug_test_storage",
            kv_storage="JsonKVStorage",
            vector_storage="NanoVectorDBStorage",
            graph_storage="NetworkXStorage",
            doc_status_storage="JsonDocStatusStorage",
            llm_model_func=llm_func,
            embedding_func=embedding_func
        )
        
        logger.info("LightRAG instance created successfully!")
        
        # Test storage initialization
        logger.info("Initializing storages...")
        await rag.initialize_storages()
        logger.info("Storages initialized successfully!")
        
        # Initialize pipeline status
        logger.info("Initializing pipeline status...")
        from lightrag.kg.shared_storage import initialize_pipeline_status
        await initialize_pipeline_status()
        logger.info("Pipeline status initialized successfully!")
        
        # Test basic insert
        logger.info("Testing document insertion...")
        test_document = "This is a test document about artificial intelligence and machine learning."
        track_id = await rag.ainsert([test_document])
        logger.info(f"Document inserted with track_id: {track_id}")
        
        return True
        
    except Exception as e:
        logger.error(f"Test failed: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return False

async def main():
    """Main test execution"""
    logger.info("Starting LLM function debug test...")
    
    success = await test_llm_function_setup()
    
    if success:
        logger.info("LLM function test PASSED!")
        return 0
    else:
        logger.error("LLM function test FAILED!")
        return 1

if __name__ == "__main__":
    exit_code = asyncio.run(main())
    sys.exit(exit_code)