from typing import Dict, Any

# Memory footprint multipliers (bytes per parameter)
QUANTIZATION_MULTIPLIERS = {
    "FP16 / BF16 (No Quantization)": 2.0,
    "8-bit (INT8)": 1.0,
    "4-bit (INT4)": 0.6  # approx 4-5 bits overhead
}

# KV Cache & Context window overhead multiplier
# A standard 20% overhead is typical for inference KV caching and frameworks
OVERHEAD_MULTIPLIER = 1.20

def calculate_kv_cache_memory(config: dict, context_length: int) -> float:
    """Calculates exactly how much memory the KV Cache needs in GB based on architecture."""
    if not config:
        return 0.0
    layers = config.get("num_hidden_layers", config.get("n_layer", 32))
    hidden_size = config.get("hidden_size", config.get("n_embd", 4096))
    num_heads = config.get("num_attention_heads", config.get("n_head", 32))
    num_kv_heads = config.get("num_key_value_heads", num_heads)
    
    # Bytes per token = 2 (fp16) * 2 (K & V) * layers * hidden_size * (kv_heads / attention_heads)
    bytes_per_token = 2 * 2 * layers * hidden_size * (num_kv_heads / num_heads)
    total_bytes = bytes_per_token * context_length
    return total_bytes / (1024 ** 3)  # Convert to GB

def get_max_model_sizes(hardware_profile: Dict[str, Any]) -> Dict[str, float]:
    """Returns the maximum supported parameter size (in billions) for the hardware natively."""
    if hardware_profile["is_apple_silicon"]:
        max_mem = hardware_profile["unified_memory_gb"]
    elif hardware_profile["has_nvidia_gpu"]:
        max_mem = hardware_profile["nvidia_vram_gb"] # Only count VRAM for strict "fits comfortably"
    else:
        max_mem = hardware_profile["system_ram_gb"]
        
    usable_mem = max(0.0, max_mem - 2.0) # Leave 2GB for OS buffer
    
    sizes = {}
    for quant_name, multiplier in QUANTIZATION_MULTIPLIERS.items():
        sizes[quant_name] = usable_mem / (multiplier * OVERHEAD_MULTIPLIER)
        
    return sizes

def calculate_training_memory(parameters_b: float) -> Dict[str, float]:
    """Calculates VRAM required for training/fine-tuning."""
    reqs = {}
    reqs["Full Fine-Tuning (AdamW, 16-bit)"] = parameters_b * 16.0
    reqs["LoRA (16-bit)"] = parameters_b * 4.0
    reqs["QLoRA (4-bit)"] = parameters_b * 1.5
    return reqs

def calculate_memory_requirements(
    parameters_b: float, 
    kv_cache_gb: float = 0.0,
    is_training: bool = False
) -> Dict[str, float]:
    """
    Calculates the memory required to load and infer the model at various quantizations.
    Returns memory requirements in gigabytes (GB).
    
    Args:
        parameters_b: Model size in billions of parameters.
    """
    requirements = {}
    
    # 1 Billion parameters = 1e9 parameters.
    # So sizes are (Params in B) * multiplier = Size in GB.
    # e.g., 8B params * 2 bytes = 16GB
    
    if is_training:
        return calculate_training_memory(parameters_b)
        
    for quant_name, multiplier in QUANTIZATION_MULTIPLIERS.items():
        base_memory = parameters_b * multiplier
        if kv_cache_gb > 0:
            total_memory = base_memory + kv_cache_gb
        else:
            total_memory = base_memory * OVERHEAD_MULTIPLIER
        requirements[quant_name] = total_memory
        
    return requirements

def evaluate_hardware_compatibility(
    memory_requirements: Dict[str, float],
    hardware_profile: Dict[str, Any]
) -> Dict[str, str]:
    """
    Evaluates if the detected hardware can support the model at different quantizations.
    Returns a dictionary mapping quantization type to a Go/No-Go status.
    """
    results = {}
    
    # Priority 1: Do they have Apple Silicon? (Unified Memory)
    if hardware_profile["is_apple_silicon"]:
        available_mem = hardware_profile["unified_memory_gb"]
        
        for quant_name, req_mem in memory_requirements.items():
            if available_mem >= req_mem:
                results[quant_name] = "🟢 Pass (Unified Memory)"
            elif available_mem > 0 and available_mem * 1.5 >= req_mem:
                results[quant_name] = "🟡 Warning (May require heavy swap)"
            else:
                results[quant_name] = "🔴 Fail (OOM)"
                
    # Priority 2: Do they have an NVIDIA GPU? (VRAM first, then offload)
    elif hardware_profile["has_nvidia_gpu"]:
        vram = hardware_profile["nvidia_vram_gb"]
        sys_ram = hardware_profile["system_ram_gb"]
        
        for quant_name, req_mem in memory_requirements.items():
            if vram >= req_mem:
                results[quant_name] = "🟢 Pass (Fits in VRAM)"
            elif (vram + sys_ram) >= req_mem:
                results[quant_name] = "🟡 Warning (Requires CPU Offloading)"
            else:
                results[quant_name] = "🔴 Fail (OOM GPU+CPU)"
                
    # Priority 3: CPU Only
    else:
        sys_ram = hardware_profile["system_ram_gb"]
        for quant_name, req_mem in memory_requirements.items():
            if sys_ram >= req_mem:
                results[quant_name] = "🟡 Warning (CPU Only - Very Slow)"
            else:
                results[quant_name] = "🔴 Fail (OOM CPU)"
                
    return results
