#!/usr/bin/env python3
"""
Verify GPU Mode for Both PaddleOCR and OpenCLIP
"""

import sys
import os
import tempfile
import subprocess

# Add LightRAG to path
workspace_dir = os.getcwd()
lightrag_path = os.path.join(workspace_dir, 'LightRAG-main')
if lightrag_path not in sys.path:
    sys.path.insert(0, lightrag_path)

def test_paddleocr_gpu():
    """Test PaddleOCR GPU configuration"""
    print('1. Testing PaddleOCR GPU configuration...')
    
    test_script = """
import paddle
print(f'    Paddle compiled with CUDA: {paddle.is_compiled_with_cuda()}')
print(f'    Paddle CUDA device count: {paddle.device.cuda.device_count()}')
try:
    if paddle.device.cuda.device_count() > 0:
        print(f'    Paddle using GPU: True')
        # Test actual OCR initialization with GPU
        from paddleocr import PaddleOCR
        ocr = PaddleOCR(use_gpu=True, use_angle_cls=True, lang='en', show_log=False, gpu_mem=2000)
        print(f'    PaddleOCR GPU initialization: SUCCESS')
    else:
        print(f'    Paddle using GPU: False')
except Exception as e:
    print(f'    PaddleOCR GPU test failed: {e}')
"""
    
    temp_dir = tempfile.mkdtemp()
    script_path = os.path.join(temp_dir, 'test_paddle_gpu.py')
    with open(script_path, 'w') as f:
        f.write(test_script)
    
    result = subprocess.run([sys.executable, script_path], capture_output=True, text=True, timeout=30)
    print(result.stdout)
    if result.stderr:
        print(f'    Paddle stderr: {result.stderr}')
    
    return 'Paddle using GPU: True' in result.stdout

def test_openclip_gpu():
    """Test OpenCLIP GPU configuration"""
    print('2. Testing OpenCLIP GPU configuration...')
    
    test_script = """
import torch
print(f'    PyTorch CUDA available: {torch.cuda.is_available()}')
print(f'    PyTorch CUDA device count: {torch.cuda.device_count()}')
if torch.cuda.is_available():
    print(f'    PyTorch current device: {torch.cuda.current_device()}')
    print(f'    PyTorch device name: {torch.cuda.get_device_name()}')
    # Test OpenCLIP model loading on GPU
    import open_clip
    model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = model.to(device)
    print(f'    OpenCLIP model loaded on: {device}')
else:
    print(f'    OpenCLIP using CPU: True')
"""
    
    temp_dir = tempfile.mkdtemp()
    script_path = os.path.join(temp_dir, 'test_openclip_gpu.py')
    with open(script_path, 'w') as f:
        f.write(test_script)
    
    result = subprocess.run([sys.executable, script_path], capture_output=True, text=True, timeout=30)
    print(result.stdout)
    if result.stderr:
        print(f'    OpenCLIP stderr: {result.stderr}')
    
    return 'OpenCLIP model loaded on: cuda' in result.stdout

def test_actual_ocr_gpu():
    """Test actual OCR with GPU"""
    print('3. Testing actual OCR processing with GPU...')
    
    try:
        from lightrag.document_processor import get_document_processor
        processor = get_document_processor()
        print(f'    OCR processor GPU mode: {processor.ocr_processor.use_gpu}')
        print(f'    OCR available: {processor.ocr_processor.ocr_available}')
        
        # Test actual OCR with GPU
        from PIL import Image, ImageDraw
        import tempfile
        
        img = Image.new('RGB', (200, 50), color='white')
        d = ImageDraw.Draw(img)
        d.text((10, 10), 'GPU Test', fill='black')
        
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
            img.save(f.name)
            temp_path = f.name
        
        try:
            ocr_result = processor.ocr_processor.extract_text_from_image(temp_path)
            text = ocr_result.get('text', '').strip()
            confidence = ocr_result.get('confidence', 0)
            print(f'    OCR test result: "{text}"')
            print(f'    OCR confidence: {confidence}')
            print(f'    OCR GPU working: {confidence > 0}')
            return confidence > 0
        finally:
            os.unlink(temp_path)
            
    except Exception as e:
        print(f'    OCR test failed: {e}')
        return False

def main():
    print('🔍 VERIFYING GPU MODE FOR BOTH MODULES')
    print('=' * 50)
    
    results = {
        'paddleocr_gpu': False,
        'openclip_gpu': False,
        'ocr_working': False
    }
    
    # Test PaddleOCR GPU
    results['paddleocr_gpu'] = test_paddleocr_gpu()
    
    # Test OpenCLIP GPU
    results['openclip_gpu'] = test_openclip_gpu()
    
    # Test actual OCR
    results['ocr_working'] = test_actual_ocr_gpu()
    
    print('\n🎯 GPU MODE SUMMARY:')
    print('=' * 50)
    print(f'   - PaddleOCR GPU mode: {"✅ ENABLED" if results["paddleocr_gpu"] else "❌ DISABLED"}')
    print(f'   - OpenCLIP GPU mode: {"✅ ENABLED" if results["openclip_gpu"] else "❌ DISABLED"}')
    print(f'   - OCR working with GPU: {"✅ YES" if results["ocr_working"] else "❌ NO"}')
    
    all_gpu_enabled = all(results.values())
    if all_gpu_enabled:
        print('\n✅ SUCCESS: Both modules running in GPU mode!')
        print('   - PaddleOCR: GPU acceleration enabled')
        print('   - OpenCLIP: GPU acceleration enabled')
        print('   - OCR processing: Working with high confidence')
        return True
    else:
        print('\n⚠️  WARNING: Some modules not in GPU mode')
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)