"""
Final Test with Complete GPU Dependency Isolation
Tests PaddleOCR and OpenCLIP with separate GPU environments
"""

import os
import sys
import asyncio
import subprocess
import tempfile
import zipfile

class GPUTestProcessor:
    """Test processor with complete GPU dependency isolation"""
    
    def __init__(self):
        self.ocr_gpu_available = False
        self.openclip_gpu_available = False
        self._check_gpu_capabilities()
    
    def _check_gpu_capabilities(self):
        """Check GPU capabilities for both PaddleOCR and OpenCLIP"""
        print("🎯 GPU CAPABILITY CHECK")
        print("=" * 40)
        
        # Check PaddleOCR GPU
        self.ocr_gpu_available = self._check_paddleocr_gpu()
        print(f"   PaddleOCR GPU: {'✅ AVAILABLE' if self.ocr_gpu_available else '❌ NOT AVAILABLE'}")
        
        # Check OpenCLIP GPU
        self.openclip_gpu_available = self._check_openclip_gpu()
        print(f"   OpenCLIP GPU: {'✅ AVAILABLE' if self.openclip_gpu_available else '❌ NOT AVAILABLE'}")
        
        print(f"   Dependency Isolation: {'✅ ACHIEVED' if self.ocr_gpu_available and self.openclip_gpu_available else '❌ FAILED'}")
    
    def _check_paddleocr_gpu(self):
        """Check if PaddleOCR can use GPU"""
        try:
            check_script = """
import paddle
print(f"PADDLE_CUDA_AVAILABLE:{paddle.is_compiled_with_cuda()}")
print(f"PADDLE_CUDA_COUNT:{paddle.device.cuda.device_count()}")
"""
            result = subprocess.run([sys.executable, "-c", check_script], 
                                  capture_output=True, text=True, timeout=10)
            return "PADDLE_CUDA_AVAILABLE:True" in result.stdout
        except:
            return False
    
    def _check_openclip_gpu(self):
        """Check if OpenCLIP can use GPU in isolated environment"""
        try:
            venv_python = "openclip_gpu_env\\Scripts\\python.exe"
            if not os.path.exists(venv_python):
                return False
            
            check_script = """
import torch
print(f"OPENCLIP_CUDA_AVAILABLE:{torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"OPENCLIP_CUDA_COUNT:{torch.cuda.device_count()}")
"""
            result = subprocess.run([venv_python, "-c", check_script], 
                                  capture_output=True, text=True, timeout=10)
            return "OPENCLIP_CUDA_AVAILABLE:True" in result.stdout
        except:
            return False
    
    def run_ocr_on_image(self, image_path):
        """Run OCR on image using isolated PaddleOCR"""
        if not self.ocr_gpu_available:
            return {"text": "", "confidence": 0.0, "line_count": 0}
        
        try:
            ocr_script = f"""
import sys
from paddleocr import PaddleOCR

try:
    ocr = PaddleOCR(use_gpu=True, cls=True)
    result = ocr.ocr(r"{image_path}")
    
    if not result or not result[0]:
        print("OCR_RESULT:EMPTY")
        sys.exit(0)
    
    extracted_text = []
    total_confidence = 0.0
    line_count = 0
    
    for line in result[0]:
        try:
            if len(line) == 2:
                bbox, (text, confidence) = line
            elif len(line) >= 1:
                bbox = line[0] if len(line) > 0 else []
                if len(line) > 1:
                    if isinstance(line[1], (list, tuple)) and len(line[1]) >= 2:
                        text, confidence = line[1][0], line[1][1]
                    else:
                        text, confidence = str(line[1]) if len(line) > 1 else "", 0.0
                else:
                    text, confidence = "", 0.0
            else:
                continue
            
            text_str = str(text) if text is not None else ""
            confidence_float = float(confidence) if isinstance(confidence, (int, float)) else 0.0
            
            extracted_text.append(text_str)
            total_confidence += confidence_float
            line_count += 1
            
        except Exception:
            extracted_text.append("")
            total_confidence += 0.0
            line_count += 1
    
    avg_confidence = total_confidence / line_count if line_count > 0 else 0.0
    full_text = " ".join(extracted_text)
    
    print(f"OCR_RESULT:TEXT={{full_text}}")
    print(f"OCR_RESULT:CONFIDENCE={{avg_confidence}}")
    print(f"OCR_RESULT:LINES={{line_count}}")
    
except Exception as e:
    print(f"OCR_ERROR:{{e}}")
"""
            result = subprocess.run([sys.executable, "-c", ocr_script], 
                                  capture_output=True, text=True, timeout=60)
            
            # Parse results
            text = ""
            confidence = 0.0
            line_count = 0
            
            for line in result.stdout.split('\n'):
                if line.startswith("OCR_RESULT:TEXT="):
                    text = line.replace("OCR_RESULT:TEXT=", "").strip()
                elif line.startswith("OCR_RESULT:CONFIDENCE="):
                    try:
                        confidence = float(line.replace("OCR_RESULT:CONFIDENCE=", "").strip())
                    except:
                        confidence = 0.0
                elif line.startswith("OCR_RESULT:LINES="):
                    try:
                        line_count = int(line.replace("OCR_RESULT:LINES=", "").strip())
                    except:
                        line_count = 0
            
            return {
                "text": text,
                "confidence": confidence,
                "line_count": line_count
            }
            
        except Exception as e:
            print(f"❌ OCR processing failed: {e}")
            return {"text": "", "confidence": 0.0, "line_count": 0}
    
    def classify_image(self, image_path, top_k=3):
        """Classify image using isolated OpenCLIP GPU environment"""
        if not self.openclip_gpu_available:
            return [{"label": "classification_unavailable", "confidence": 0.0}]
        
        try:
            venv_python = "openclip_gpu_env\\Scripts\\python.exe"
            
            classification_script = f"""
import open_clip
import torch
from PIL import Image

try:
    # Load model
    model, _, processor = open_clip.create_model_and_transforms(
        model_name="ViT-B-32",
        pretrained="laion2b_s34b_b79k"
    )
    
    # Move to GPU if available
    if torch.cuda.is_available():
        model = model.cuda()
        print("GPU_USED:YES")
    else:
        print("GPU_USED:NO")
    
    # Load and process image
    image = Image.open(r"{image_path}").convert("RGB")
    image_tensor = processor(image).unsqueeze(0)
    
    if torch.cuda.is_available():
        image_tensor = image_tensor.cuda()
    
    # Get predictions
    with torch.no_grad():
        image_features = model.encode_image(image_tensor)
        image_features /= image_features.norm(dim=-1, keepdim=True)
        
        # Common labels
        text_labels = [
            "a photo of a bee", "a photo of a flower", "a photo of a person", 
            "a photo of a document", "a photo of a chart", "a photo of a diagram",
            "a photo of a table", "a photo of a graph", "a photo of a logo",
            "a photo of a signature", "a photo of a stamp", "a photo of a barcode",
            "a photo of a QR code", "a photo of a screenshot", "a photo of a landscape"
        ]
        
        # Encode text labels
        text_tokens = open_clip.tokenize(text_labels)
        if torch.cuda.is_available():
            text_tokens = text_tokens.cuda()
        
        text_features = model.encode_text(text_tokens)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        
        # Calculate similarity
        similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
        values, indices = similarity[0].topk({top_k})
        
        for value, index in zip(values, indices):
            label = text_labels[index]
            confidence = float(value)
            print(f"CLASSIFICATION_RESULT:{{label}}|{{confidence}}")
            
except Exception as e:
    print(f"CLASSIFICATION_ERROR:{{e}}")
"""
            result = subprocess.run([venv_python, "-c", classification_script], 
                                  capture_output=True, text=True, timeout=30)
            
            results = []
            gpu_used = False
            
            for line in result.stdout.split('\n'):
                if line.startswith("GPU_USED:YES"):
                    gpu_used = True
                if line.startswith("CLASSIFICATION_RESULT:"):
                    parts = line.replace("CLASSIFICATION_RESULT:", "").split("|")
                    if len(parts) == 2:
                        try:
                            label = parts[0]
                            confidence = float(parts[1])
                            results.append({"label": label, "confidence": confidence, "gpu_used": gpu_used})
                        except:
                            continue
            
            if results:
                return results
            else:
                return [{"label": "classification_failed", "confidence": 0.0, "gpu_used": False}]
                
        except Exception as e:
            print(f"❌ Classification failed: {e}")
            return [{"label": "classification_error", "confidence": 0.0, "gpu_used": False}]


async def test_complete_gpu_isolation():
    """Test complete GPU dependency isolation"""
    print("🧪 COMPLETE GPU DEPENDENCY ISOLATION TEST")
    print("=" * 50)
    
    processor = GPUTestProcessor()
    
    if not processor.ocr_gpu_available or not processor.openclip_gpu_available:
        print("❌ GPU dependency isolation not achieved")
        return False
    
    # Test with test.docx
    test_file = "test.docx"
    if not os.path.exists(test_file):
        print(f"❌ Test file not found: {test_file}")
        return False
    
    print(f"\n📄 PROCESSING: {test_file}")
    
    try:
        with tempfile.TemporaryDirectory() as temp_dir:
            # Extract images from docx
            with zipfile.ZipFile(test_file, 'r') as zip_ref:
                image_files = []
                for file_info in zip_ref.filelist:
                    if file_info.filename.startswith('word/media/'):
                        image_filename = os.path.basename(file_info.filename)
                        image_path = os.path.join(temp_dir, image_filename)
                        
                        with zip_ref.open(file_info.filename) as source, open(image_path, 'wb') as target:
                            target.write(source.read())
                        
                        image_files.append(image_path)
                        print(f"📸 Extracted image: {image_path}")
                
                print(f"Found {len(image_files)} images in document")
            
            # Process each image
            ocr_results = []
            classification_results = []
            bee_detected = False
            
            for i, image_path in enumerate(image_files):
                print(f"\n🖼️ Processing image {i+1}...")
                
                # OCR processing
                if processor.ocr_gpu_available:
                    ocr_result = processor.run_ocr_on_image(image_path)
                    if ocr_result["text"].strip():
                        ocr_results.append({
                            "image": i+1,
                            "text_length": len(ocr_result["text"]),
                            "confidence": ocr_result["confidence"]
                        })
                        print(f"   ✅ OCR: {len(ocr_result['text'])} chars, confidence: {ocr_result['confidence']:.3f}")
                
                # Image classification
                if processor.openclip_gpu_available:
                    classification_result = processor.classify_image(image_path)
                    if classification_result and classification_result[0]["confidence"] > 0:
                        top_result = classification_result[0]
                        classification_results.append({
                            "image": i+1,
                            "label": top_result["label"],
                            "confidence": top_result["confidence"],
                            "gpu_used": top_result.get("gpu_used", False)
                        })
                        print(f"   ✅ Classification: {top_result['label']} ({top_result['confidence']:.3f})")
                        
                        if "bee" in top_result["label"].lower():
                            bee_detected = True
                            print(f"   🎯 BEE DETECTED!")
            
            # Summary
            print(f"\n📊 PROCESSING SUMMARY:")
            print(f"   OCR Images Processed: {len(ocr_results)}")
            print(f"   Classification Images Processed: {len(classification_results)}")
            print(f"   Bee Detection: {'✅ SUCCESS' if bee_detected else '❌ NOT FOUND'}")
            print(f"   PaddleOCR GPU: {'✅ USED' if processor.ocr_gpu_available else '❌ NOT USED'}")
            print(f"   OpenCLIP GPU: {'✅ USED' if processor.openclip_gpu_available else '❌ NOT USED'}")
            print(f"   Dependency Isolation: {'✅ ACHIEVED' if processor.ocr_gpu_available and processor.openclip_gpu_available else '❌ FAILED'}")
            
            return True
            
    except Exception as e:
        print(f"❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


async def main():
    """Main test function"""
    print("🚀 TESTING COMPLETE GPU DEPENDENCY ISOLATION")
    print("This test verifies:")
    print("  - PaddleOCR using GPU in main environment")
    print("  - OpenCLIP using GPU in isolated virtual environment")
    print("  - Complete dependency isolation between components")
    print("  - Bee image detection with GPU acceleration")
    print()
    
    success = await test_complete_gpu_isolation()
    
    if success:
        print("\n🎉 GPU DEPENDENCY ISOLATION SUCCESSFUL! 🎉")
        print("Both PaddleOCR and OpenCLIP are now:")
        print("  ✅ Running with GPU acceleration")
        print("  ✅ Completely isolated from each other")
        print("  ✅ Using separate PyTorch versions")
        print("  ✅ Successfully processing documents")
    else:
        print("\n❌ GPU DEPENDENCY ISOLATION FAILED")
        print("Please ensure:")
        print("  - NVIDIA GPU drivers are installed")
        print("  - CUDA is properly configured")
        print("  - Virtual environment 'openclip_gpu_env' exists")
        print("  - Both PaddleOCR and OpenCLIP are installed correctly")


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