# ==========================================
# utils.py - Fonctions communes CPU/GPU
# ==========================================

"""
Utilitaires partagés pour le traitement d'images OCR
Fonctions communes aux versions CPU et GPU
"""

from PIL import Image, ImageEnhance
import numpy as np
import base64
from io import BytesIO
import gc
import os
import time

def create_white_canvas(width: int = 300, height: int = 300) -> Image.Image:
    """Crée un canvas blanc pour le dessin de calculs"""
    return Image.new('RGB', (width, height), 'white')

def log_memory_usage(context: str = "") -> None:
    """Log l'usage mémoire actuel"""
    try:
        import psutil
        process = psutil.Process(os.getpid())
        memory_mb = process.memory_info().rss / 1024 / 1024
        print(f"🔍 Mémoire {context}: {memory_mb:.1f}MB")
    except:
        pass

def cleanup_memory() -> None:
    """Force le nettoyage mémoire"""
    gc.collect()

def optimize_image_for_ocr(image_dict: dict | np.ndarray | Image.Image | None, max_size: int = 300) -> Image.Image | None:
    """
    Optimisation image commune pour tous types d'OCR
    
    Args:
        image_dict: Image d'entrée (format Gradio, numpy ou PIL)
        max_size: Taille maximale pour le redimensionnement
        
    Returns:
        Image PIL optimisée ou None si erreur
    """
    if image_dict is None:
        return None

    try:
        # Gérer les formats Gradio
        if isinstance(image_dict, dict):
            if 'composite' in image_dict and image_dict['composite'] is not None:
                image = image_dict['composite']
            elif 'background' in image_dict and image_dict['background'] is not None:
                image = image_dict['background']
            else:
                return None
        elif isinstance(image_dict, np.ndarray):
            image = image_dict
        elif isinstance(image_dict, Image.Image):
            image = image_dict
        else:
            return None

        # Conversion vers PIL
        if isinstance(image, np.ndarray):
            pil_image = Image.fromarray(image).convert('RGB')
        else:
            pil_image = image.convert('RGB')

        # Redimensionnement si nécessaire
        if pil_image.size[0] > max_size or pil_image.size[1] > max_size:
            pil_image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)

        return pil_image

    except Exception as e:
        print(f"❌ Erreur optimisation image: {e}")
        return None

def prepare_image_for_dataset(image: Image.Image, max_size: tuple[int, int] = (100, 100), quality: int = 60) -> dict[str, str | int | float | tuple] | None:
    """
    Prépare une image pour l'inclusion dans le dataset
    
    Args:
        image: Image PIL à traiter
        max_size: Taille maximale (largeur, hauteur)
        quality: Qualité de compression PNG
        
    Returns:
        Dictionnaire avec image_base64, taille, etc. ou None
    """
    try:
        if image is None:
            return None
            
        # Copier et redimensionner
        dataset_image = image.copy()
        dataset_image.thumbnail(max_size, Image.Resampling.LANCZOS)
        compressed_size = dataset_image.size
        
        # Convertir en base64
        buffer = BytesIO()
        dataset_image.save(buffer, format='PNG', optimize=True, quality=quality)
        
        buffer_data = buffer.getvalue()
        image_base64 = base64.b64encode(buffer_data).decode()
        file_size_kb = len(image_base64) / 1024
        
        # Structure propre pour dataset
        result = {
            "image_base64": image_base64,
            "compressed_size": compressed_size,
            "file_size_kb": round(file_size_kb, 1),
            "format": "PNG",
            "quality": quality
        }
        
        # Nettoyage
        dataset_image.close()
        buffer.close()
        
        return result
        
    except Exception as e:
        print(f"❌ Erreur préparation image dataset: {e}")
        return None

def create_thumbnail_fast(optimized_image: Image.Image | None, size: tuple[int, int] = (40, 40)) -> str:
    """
    Création miniature rapide pour affichage dans les résultats
    
    Args:
        optimized_image: Image PIL source
        size: Taille de la miniature (largeur, hauteur)
        
    Returns:
        HTML img tag avec image base64 ou icône par défaut
    """
    try:
        if optimized_image is None:
            return "📝"
            
        thumbnail = optimized_image.copy()
        thumbnail.thumbnail(size, Image.Resampling.LANCZOS)
        
        buffer = BytesIO()
        thumbnail.save(buffer, format='PNG', optimize=True, quality=70)
        img_str = base64.b64encode(buffer.getvalue()).decode()
        
        thumbnail.close()
        buffer.close()
        
        return f'<img src="data:image/png;base64,{img_str}" width="{size[0]}" height="{size[1]}" style="border: 1px solid #ccc; border-radius: 3px;" alt="Réponse calcul">'
        
    except Exception:
        return "📝"

def decode_image_from_dataset(base64_string: str) -> Image.Image | None:
    """
    Décode une image depuis le dataset pour fine-tuning ou analyse
    
    Args:
        base64_string: String base64 de l'image
        
    Returns:
        Image PIL ou None si erreur
    """
    try:
        image_bytes = base64.b64decode(base64_string)
        image = Image.open(BytesIO(image_bytes))
        return image
    except Exception as e:
        print(f"❌ Erreur décodage image dataset: {e}")
        return None

def validate_ocr_result(raw_result: str, max_length: int = 4) -> str:
    """
    Valide et nettoie un résultat OCR
    
    Args:
        raw_result: Résultat brut de l'OCR
        max_length: Longueur maximale autorisée
        
    Returns:
        Résultat nettoyé (chiffres uniquement)
    """
    if not raw_result:
        return "0"
    
    # Extraire uniquement les chiffres
    cleaned_result = ''.join(filter(str.isdigit, str(raw_result)))
    
    # Valider la longueur
    if cleaned_result and len(cleaned_result) <= max_length:
        return cleaned_result
    elif cleaned_result:
        # Si trop long, prendre les premiers chiffres
        return cleaned_result[:max_length]
    else:
        return "0"

def analyze_calculation_complexity(operand_a: int, operand_b: int, operation: str) -> dict:
    """
    Analyse la complexité d'un calcul pour enrichir les métadonnées dataset
    
    Args:
        operand_a: Premier opérande
        operand_b: Deuxième opérande  
        operation: Type d'opération (×, +, -, ÷)
        
    Returns:
        Dictionnaire avec score de complexité et catégorie
    """
    complexity_score = 0
    
    if operation == "×":
        complexity_score = max(operand_a, operand_b)
    elif operation == "+":
        complexity_score = (operand_a + operand_b) / 20
    elif operation == "-":
        complexity_score = max(operand_a, operand_b) / 10
    elif operation == "÷":
        complexity_score = operand_a / 10
    
    # Catégorisation
    if complexity_score < 5:
        category = "easy"
    elif complexity_score < 10:
        category = "medium"
    else:
        category = "hard"
    
    return {
        "complexity_score": round(complexity_score, 2),
        "difficulty_category": category,
        "operation_type": operation
    }