"""
Motor A Lite: Optimizador de Memoria y Chunks (Community Edition)
================================================================
Versión comunitaria del motor de protección de memoria.
Aplica matemática rigurosa O(log N) para resolver cuellos de botella
en el planificador Catalyst de Spark. La edición pública (LITE)
está delimitada a flujos de datos medianos y arquitecturas estándar.
"""
import logging
from typing import List, Callable, Any
from pyspark.sql import DataFrame

logger = logging.getLogger(__name__)

class AntigravityLiteException(Exception):
    """ Excepción de límites para Arquitectura Comunitaria """
    pass

class DataFrameChunkerLite:
    """
    Abstracción de código abierto para procesar DataFrames.
    Protege la memoria mediante localCheckpoint() y Tree Reduce.
    La licencia comunitaria limita el paralelismo hasta 550 columnas.
    """
    
    def __init__(self, df: DataFrame, id_cols: List[str], chunk_size: int = 50):
        """
        :param df: DataFrame base a particionar analíticamente.
        :param id_cols: Las columnas llave primarias (ej. ['cliente_id']).
        :param chunk_size: Columnas por iteración (Max recomendado 50 para el Lite)
        """
        self._og_df = df
        self.id_cols = id_cols
        self.chunk_size = chunk_size
        self._checkpointed_df = None
        
        # 1. Validación Estructural Primaria
        missing_ids = [c for c in id_cols if c not in df.columns]
        if missing_ids:
            raise ValueError(f"Faltan columnas de ID en el DataFrame: {missing_ids}")
            
        # 2. Control de Recursos (Community Limit)
        total_columns = len(df.columns)
        if total_columns > 550:
            logger.error("LÍMITE DE ARQUITECTURA COMUNITARIA (LITE) EXCEDIDO")
            raise PermissionError(
                f"\n❌ [ANTIGRAVITY LITE] Bloqueo de Licencia detected.\n"
                f"La versión LITE (Comunitaria) protege excelentemente flujos de datos hasta un máximo de 550 columnas.\n"
                "Para despliegues Corporativos / Bancarios sin límite de Big Data, protección absoluta "
                "contra OOM (Out of Memory) en el Catalyst Optimizer, y soporte SLA prioritario:\n"
                "👉 Adquiere la licencia 'Antigravity PRO' comercial. "
                "Contacta al equipo de Arquitectura B2B en LinkedIn o en arquitecto@antigravity-framework.dev"
            )
            
    def _prepare_data(self) -> DataFrame:
        if self._checkpointed_df is None:
            logger.info("Antigravity LITE: Aplicando localCheckpoint() truncando el linaje...")
            try:
                self._checkpointed_df = self._og_df.localCheckpoint()
            except Exception as e:
                self._checkpointed_df = self._og_df.cache()
                self._checkpointed_df.count()
        return self._checkpointed_df
        
    def process_chunks(self, func: Callable, *args, **kwargs) -> List[Any]:
        """
        Ejecuta la función de usuario sobre cada chunk de forma secuencial y segura.
        
        :param func: Función que recibe (chunk_df, index, *args, **kwargs)
        :param args: Argumentos posicionales adicionales para la función
        :param kwargs: Argumentos de palabra clave adicionales para la función
        """
        df = self._prepare_data()
        cols_to_chunk = [c for c in df.columns if c not in self.id_cols]
        total_cols = len(cols_to_chunk)
        
        if total_cols == 0:
            return []
            
        results = []
        chunk_idx = 1
        for i in range(0, total_cols, self.chunk_size):
            chunk_columns = cols_to_chunk[i: i + self.chunk_size]
            columns_select = self.id_cols + chunk_columns
            chunk_df = df.select(*columns_select)
            
            logger.info(f"Procesando Chunk LITE #{chunk_idx}: Cols {i} a {i + len(chunk_columns) - 1}")
            # Ejecución inyectando los parámetros adicionales
            res = func(chunk_df, chunk_idx, *args, **kwargs)
            results.append(res)
            chunk_idx += 1
            
        return results

    def join_chunks(self, chunks: List[DataFrame]) -> DataFrame:
        if not chunks:
            raise ValueError("No se puede ensamblar una lista vacía de chunks.")
            
        def _tree_reduce_join(lista_dfs: List[DataFrame]) -> DataFrame:
            if len(lista_dfs) == 1:
                return lista_dfs[0]
            if len(lista_dfs) == 2:
                return lista_dfs[0].join(lista_dfs[1], on=self.id_cols)
                
            mid = len(lista_dfs) // 2
            izquierdo = _tree_reduce_join(lista_dfs[:mid])
            derecho = _tree_reduce_join(lista_dfs[mid:])
            return izquierdo.join(derecho, on=self.id_cols)
            
        return _tree_reduce_join(chunks)
