"""
GPU acceleration utilities for LightRAG embedding functions.

This module provides GPU-accelerated embedding computation with support for:
- CUDA (NVIDIA GPUs)
- ROCm (AMD GPUs) 
- CPU fallback
- Batch processing optimization
- Memory management
- Automatic device detection
"""

import os
import logging
import asyncio
from typing import List, Optional, Callable, Any
import numpy as np

try:
    import torch
    import torch.nn as nn
    from torch.cuda.amp import autocast
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

logger = logging.getLogger("lightrag")

# Global GPU configuration
GPU_CONFIG = {
    "enabled": os.getenv("LIGHTRAG_GPU_ENABLED", "true").lower() == "true",
    "batch_size": int(os.getenv("LIGHTRAG_GPU_BATCH_SIZE", "64")),
    "max_concurrent_batches": int(os.getenv("LIGHTRAG_GPU_MAX_CONCURRENT_BATCHES", "4")),
    "precision": os.getenv("LIGHTRAG_GPU_PRECISION", "fp16"),  # fp16, fp32, bf16
}


def detect_gpu_device() -> Optional[str]:
    """
    Detect available GPU device and return appropriate device string.
    
    Returns:
        Device string ("cuda", "mps", "cpu") or None if no GPU available
    """
    if not TORCH_AVAILABLE:
        return None
        
    try:
        if torch.cuda.is_available():
            # Check CUDA device count and memory
            if torch.cuda.device_count() > 0:
                device_name = torch.cuda.get_device_name(0)
                free_memory = torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated(0)
                logger.info(f"CUDA GPU detected: {device_name}, Free memory: {free_memory / 1024**3:.1f} GB")
                return "cuda"
                
        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            logger.info("Apple MPS (Metal Performance Shaders) detected")
            return "mps"
            
    except Exception as e:
        logger.warning(f"GPU detection failed: {e}")
        
    return "cpu"


def get_optimal_batch_size(embedding_dim: int, device: str) -> int:
    """
    Calculate optimal batch size based on embedding dimension and device.
    
    Args:
        embedding_dim: Dimension of the embeddings
        device: Device type ("cuda", "mps", "cpu")
        
    Returns:
        Optimal batch size
    """
    base_batch_size = GPU_CONFIG["batch_size"]
    
    if device == "cuda":
        # Adjust for GPU memory
        try:
            free_memory = torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated(0)
            # Rough estimate: 4 bytes per float * dimensions * batch size
            memory_per_batch = 4 * embedding_dim * base_batch_size
            if free_memory > 2 * memory_per_batch:  # Leave room for other operations
                return base_batch_size
            else:
                # Reduce batch size based on available memory
                reduced_batch = max(1, int((free_memory * 0.5) / (4 * embedding_dim)))
                logger.info(f"Reducing batch size from {base_batch_size} to {reduced_batch} due to memory constraints")
                return reduced_batch
        except:
            return base_batch_size
            
    elif device == "mps":
        # MPS typically has less memory than dedicated GPUs
        return min(base_batch_size, 32)
        
    else:
        # CPU - use smaller batches
        return min(base_batch_size, 16)


async def gpu_embedding_wrapper(
    embedding_func: Callable[[List[str]], np.ndarray],
    texts: List[str],
    device: Optional[str] = None,
    model: Optional[Any] = None,
    tokenizer: Optional[Any] = None,
    **kwargs
) -> np.ndarray:
    """
    GPU-accelerated embedding wrapper with batch processing optimization.
    
    Args:
        embedding_func: Original embedding function to wrap
        texts: List of texts to embed
        device: Target device ("cuda", "mps", "cpu")
        model: Optional model for local GPU inference
        tokenizer: Optional tokenizer for local models
        **kwargs: Additional arguments for the embedding function
        
    Returns:
        Numpy array of embeddings
    """
    if not GPU_CONFIG["enabled"] or not texts:
        return await embedding_func(texts, **kwargs)
    
    # Detect device if not specified
    target_device = device or detect_gpu_device() or "cpu"
    
    if target_device == "cpu":
        # Fall back to CPU processing
        return await embedding_func(texts, **kwargs)
    
    # For Hugging Face models with local GPU support
    if model is not None and tokenizer is not None and TORCH_AVAILABLE:
        return await _hf_gpu_embedding(texts, model, tokenizer, target_device)
    
    # For API-based embedding functions, use batch optimization
    return await _batch_optimized_embedding(embedding_func, texts, target_device, **kwargs)


async def _hf_gpu_embedding(
    texts: List[str],
    model: Any,
    tokenizer: Any,
    device: str
) -> np.ndarray:
    """
    GPU-accelerated embedding for Hugging Face models.
    
    Args:
        texts: List of texts to embed
        model: Hugging Face model
        tokenizer: Hugging Face tokenizer
        device: Target device
        
    Returns:
        Numpy array of embeddings
    """
    if not TORCH_AVAILABLE:
        raise ImportError("PyTorch is required for GPU-accelerated Hugging Face embeddings")
    
    try:
        # Move model to target device
        model = model.to(device)
        
        # Determine optimal batch size
        embedding_dim = model.config.hidden_size
        batch_size = get_optimal_batch_size(embedding_dim, device)
        
        embeddings = []
        
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            
            # Tokenize batch
            encoded = tokenizer(
                batch_texts, 
                return_tensors="pt", 
                padding=True, 
                truncation=True,
                max_length=512
            ).to(device)
            
            # Generate embeddings with appropriate precision
            with torch.no_grad():
                if GPU_CONFIG["precision"] == "fp16":
                    with autocast():
                        outputs = model(**encoded)
                else:
                    outputs = model(**encoded)
                
                # Use mean pooling of last hidden states
                batch_embeddings = outputs.last_hidden_state.mean(dim=1)
                
                # Convert to appropriate precision
                if GPU_CONFIG["precision"] == "fp16":
                    batch_embeddings = batch_embeddings.half()
                elif GPU_CONFIG["precision"] == "bf16":
                    batch_embeddings = batch_embeddings.bfloat16()
                
                embeddings.append(batch_embeddings.cpu().numpy())
        
        return np.concatenate(embeddings)
        
    except Exception as e:
        logger.error(f"GPU embedding failed: {e}")
        # Fall back to CPU
        from lightrag.llm.hf import hf_embed
        return await hf_embed(texts, tokenizer, model)


async def _batch_optimized_embedding(
    embedding_func: Callable[[List[str]], np.ndarray],
    texts: List[str],
    device: str,
    **kwargs
) -> np.ndarray:
    """
    Optimize batch processing for API-based embedding functions.
    
    Args:
        embedding_func: API embedding function
        texts: List of texts to embed
        device: Target device (for logging)
        **kwargs: Additional arguments
        
    Returns:
        Numpy array of embeddings
    """
    batch_size = GPU_CONFIG["batch_size"]
    max_concurrent = GPU_CONFIG["max_concurrent_batches"]
    
    logger.debug(f"GPU-optimized batch processing: {len(texts)} texts, batch_size={batch_size}, device={device}")
    
    # Create batches
    batches = [
        texts[i:i + batch_size]
        for i in range(0, len(texts), batch_size)
    ]
    
    # Process batches with concurrency control
    semaphore = asyncio.Semaphore(max_concurrent)
    
    async def process_batch(batch):
        async with semaphore:
            return await embedding_func(batch, **kwargs)
    
    # Process all batches concurrently
    embedding_tasks = [process_batch(batch) for batch in batches]
    embeddings_list = await asyncio.gather(*embedding_tasks)
    
    return np.concatenate(embeddings_list)


def wrap_with_gpu_acceleration(
    embedding_func: Callable[[List[str]], np.ndarray],
    embedding_dim: int,
    model: Optional[Any] = None,
    tokenizer: Optional[Any] = None
) -> Callable[[List[str]], np.ndarray]:
    """
    Wrap an embedding function with GPU acceleration.
    
    Args:
        embedding_func: Original embedding function
        embedding_dim: Dimension of the embeddings
        model: Optional model for local inference
        tokenizer: Optional tokenizer for local models
        
    Returns:
        Wrapped embedding function with GPU acceleration
    """
    from lightrag.utils import wrap_embedding_func_with_attrs
    
    async def gpu_accelerated_embedding(texts: List[str], **kwargs) -> np.ndarray:
        return await gpu_embedding_wrapper(
            embedding_func, texts, model=model, tokenizer=tokenizer, **kwargs
        )
    
    # Preserve original function attributes
    gpu_accelerated_embedding.__name__ = f"gpu_accelerated_{embedding_func.__name__}"
    gpu_accelerated_embedding.__doc__ = f"GPU-accelerated version of {embedding_func.__name__}"
    
    # Wrap with attributes
    return wrap_embedding_func_with_attrs(
        embedding_dim=embedding_dim,
        func=gpu_accelerated_embedding
    )


# Example usage for existing embedding functions
def enable_gpu_for_embedding_functions():
    """
    Enable GPU acceleration for all embedding functions in the system.
    This should be called during application initialization.
    """
    if not GPU_CONFIG["enabled"]:
        logger.info("GPU acceleration is disabled via configuration")
        return
        
    device = detect_gpu_device()
    if device and device != "cpu":
        logger.info(f"GPU acceleration enabled using {device.upper()}")
    else:
        logger.info("No GPU detected, using CPU fallback")
        GPU_CONFIG["enabled"] = False


# Initialize GPU configuration on import
enable_gpu_for_embedding_functions()