import os
import sys
import subprocess
import requests
import time

def check_cuda_11_8_setup():
    """Check if CUDA 11.8 is properly installed and configured"""
    print("=== CHECKING CUDA 11.8 SETUP ===")
    
    cuda_11_8_path = r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8'
    
    # Check if CUDA 11.8 directory exists
    if not os.path.exists(cuda_11_8_path):
        print(f"✗ CUDA 11.8 not found at: {cuda_11_8_path}")
        print("\nPlease install CUDA 11.8 from:")
        print("https://developer.nvidia.com/cuda-11-8-0-download-archive")
        return False
    
    print(f"✓ CUDA 11.8 found at: {cuda_11_8_path}")
    
    # Check for cuDNN
    cudnn_dll = os.path.join(cuda_11_8_path, 'bin', 'cudnn64_8.dll')
    if not os.path.exists(cudnn_dll):
        print(f"✗ cuDNN 8.x not found at: {cudnn_dll}")
        print("\nPlease install cuDNN 8.x for CUDA 11.8 from:")
        print("https://developer.nvidia.com/rdp/cudnn-archive")
        print("Then copy the DLLs to CUDA 11.8 bin directory")
        return False
    
    print(f"✓ cuDNN 8.x found at: {cudnn_dll}")
    
    # Set environment for CUDA 11.8
    os.environ['CUDA_PATH'] = cuda_11_8_path
    os.environ['CUDA_HOME'] = cuda_11_8_path
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    os.environ['LIGHTRAG_OCR_ENGINE'] = 'paddleocr'
    
    # Add CUDA 11.8 to PATH
    cuda_bin = os.path.join(cuda_11_8_path, 'bin')
    current_path = os.environ.get('PATH', '')
    if cuda_bin not in current_path:
        os.environ['PATH'] = cuda_bin + ';' + current_path
    
    print("✓ Environment configured for CUDA 11.8")
    return True

def test_paddlepaddle_gpu():
    """Test if PaddlePaddle can use GPU with CUDA 11.8"""
    print("\n=== TESTING PADDLEPADDLE GPU ===")
    
    try:
        import paddle
        print(f"✓ PaddlePaddle version: {paddle.__version__}")
        print(f"✓ GPU available: {paddle.is_compiled_with_cuda()}")
        
        if paddle.is_compiled_with_cuda():
            paddle.device.set_device('gpu')
            print("✓ PaddlePaddle GPU device set successfully")
            
            # Test a simple GPU operation
            x = paddle.to_tensor([1.0, 2.0, 3.0])
            print(f"✓ GPU tensor test: {x.numpy()}")
            return True
        else:
            print("✗ PaddlePaddle not compiled with CUDA")
            return False
            
    except Exception as e:
        print(f"✗ PaddlePaddle test failed: {e}")
        return False

def start_server_with_cuda11_8():
    """Start server with CUDA 11.8 environment"""
    print("\n=== STARTING SERVER WITH CUDA 11.8 ===")
    
    try:
        # Use the zrun_cuda11.8.bat approach
        cmd = [
            'lightrag-server', 
            '--port', '3015', 
            '--embedding-binding', 'ollama', 
            '--rerank-binding', 'null', 
            '--host', '0.0.0.0'
        ]
        
        print(f"Starting server: {' '.join(cmd)}")
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        # Wait for server to start
        print("Waiting for server to start...")
        for i in range(60):
            try:
                response = requests.get('http://localhost:3015/', timeout=5)
                if response.status_code == 200:
                    print("✓ Server started successfully with CUDA 11.8!")
                    return process
            except:
                pass
            time.sleep(1)
        
        print("✗ Server failed to start within timeout")
        return None
        
    except Exception as e:
        print(f"✗ Failed to start server: {e}")
        return None

def test_ocr_pdf_indexing():
    """Test OCR PDF indexing with CUDA 11.8"""
    print("\n=== TESTING OCR PDF INDEXING WITH CUDA 11.8 ===")
    base_url = 'http://localhost:3015'
    
    try:
        # Login
        login_data = {'username': 'jleu3482', 'password': 'jleu1212'}
        login_response = requests.post(f'{base_url}/login', data=login_data, timeout=30)
        
        if login_response.status_code != 200:
            print(f"✗ Login failed: {login_response.text}")
            return False
            
        token = login_response.json().get('access_token')
        headers = {'Authorization': f'Bearer {token}'}
        print("✓ Login successful")
        
        # Clear existing documents
        print("Clearing existing documents...")
        clear_response = requests.delete(f'{base_url}/documents', headers=headers, timeout=30)
        print(f"Clear status: {clear_response.status_code}")
        
        # Upload OCR PDF
        print(f"\nUploading OCR PDF: ocr.pdf ({os.path.getsize('ocr.pdf')} bytes)")
        with open('ocr.pdf', 'rb') as f:
            files = {'file': ('ocr.pdf', f, 'application/pdf')}
            upload_response = requests.post(f'{base_url}/documents/upload', files=files, headers=headers, timeout=60)
        
        print(f"Upload status: {upload_response.status_code}")
        if upload_response.status_code != 200:
            print(f"✗ Upload failed: {upload_response.text}")
            return False
            
        upload_data = upload_response.json()
        print(f"Upload response: {upload_data}")
        track_id = upload_data.get('track_id')
        
        if not track_id:
            print("✗ No track ID returned")
            return False
        
        # Monitor processing
        print(f"\n=== MONITORING OCR PROCESSING (CUDA 11.8) ===")
        print("OCR processing with GPU acceleration...")
        
        max_wait = 300  # 5 minutes for GPU processing
        start_time = time.time()
        
        while time.time() - start_time < max_wait:
            try:
                # Check document status
                docs_response = requests.get(f'{base_url}/documents', headers=headers, timeout=30)
                if docs_response.status_code == 200:
                    docs_data = docs_response.json()
                    statuses = docs_data.get('statuses', {})
                    
                    completed = statuses.get('completed', [])
                    processing = statuses.get('processing', [])
                    failed = statuses.get('failed', [])
                    
                    elapsed = int(time.time() - start_time)
                    
                    # Check for our file
                    for doc in completed:
                        if doc.get('file_path') == 'ocr.pdf':
                            print(f"\n🎉 OCR PROCESSING COMPLETED in {elapsed} seconds!")
                            print(f"  File: {doc.get('file_path')}")
                            print(f"  Size: {doc.get('file_size')}")
                            print(f"  Chunks: {doc.get('chunk_count')}")
                            print(f"  Processing time: {doc.get('processing_time', 'N/A')}")
                            
                            # Test search
                            print("\n=== TESTING SEARCH ===")
                            search_queries = [
                                "table", "data", "information", "content",
                                "scanned", "document", "text", "analysis"
                            ]
                            
                            total_results = 0
                            for query in search_queries:
                                print(f"Searching: '{query}'")
                                search_data = {'query': query, 'top_k': 3}
                                try:
                                    search_response = requests.post(f'{base_url}/search', json=search_data, headers=headers, timeout=30)
                                    if search_response.status_code == 200:
                                        results = search_response.json().get('results', [])
                                        total_results += len(results)
                                        print(f"  Found {len(results)} results")
                                        if results:
                                            print(f"    Best score: {results[0].get('score'):.3f}")
                                    else:
                                        print(f"  Search failed: {search_response.text}")
                                except Exception as e:
                                    print(f"  Search error: {e}")
                            
                            print(f"\n=== SUMMARY ===")
                            print(f"Total search results: {total_results}")
                            print("OCR PDF with scanned table successfully indexed and searchable!")
                            return True
                    
                    # Check if failed
                    for doc in failed:
                        if doc.get('file_path') == 'ocr.pdf':
                            print(f"✗ OCR processing failed: {doc.get('error_msg', 'Unknown error')}")
                            return False
                    
                    # Still processing
                    if elapsed % 30 == 0:
                        print(f"  Processing... ({elapsed}s elapsed)")
                
                time.sleep(5)
                
            except requests.exceptions.RequestException as e:
                print(f"  Connection error: {e}")
                time.sleep(10)
        
        print(f"✗ OCR processing timed out after {max_wait} seconds")
        return False
        
    except Exception as e:
        print(f"✗ Error during OCR test: {e}")
        return False

def main():
    """Main function to test CUDA 11.8 OCR setup"""
    print("CUDA 11.8 OCR TEST SCRIPT")
    print("=" * 50)
    
    # Step 1: Check CUDA 11.8 setup
    if not check_cuda_11_8_setup():
        print("\n❌ CUDA 11.8 setup incomplete")
        return
    
    # Step 2: Test PaddlePaddle GPU
    if not test_paddlepaddle_gpu():
        print("\n❌ PaddlePaddle GPU test failed")
        return
    
    # Step 3: Start server with CUDA 11.8
    server_process = start_server_with_cuda11_8()
    if not server_process:
        print("\n❌ Failed to start server")
        return
    
    try:
        # Step 4: Test OCR PDF indexing
        success = test_ocr_pdf_indexing()
        
        if success:
            print("\n🎉 SUCCESS: OCR PDF indexing with CUDA 11.8 completed!")
            print("   GPU-accelerated OCR is working correctly.")
            print("   The scanned table document has been processed and is searchable.")
        else:
            print("\n❌ OCR PDF indexing failed")
            
    finally:
        # Clean up
        print("\nStopping server...")
        server_process.terminate()
        try:
            server_process.wait(timeout=10)
        except:
            server_process.kill()

if __name__ == "__main__":
    main()