"""
Motor IO: Smart S3 Exporter (S3Finalizer)
=========================================
Maneja la "Última Milla" del procesamiento de PySpark en AWS.
Convierte las salidas nativas de Spark (ej. part-0000...) en archivos limpios, legibles y secuenciados.
"""
import time
import logging
import concurrent.futures
from typing import List, Optional

try:
    import boto3
    from botocore.exceptions import ClientError
except ImportError:
    boto3 = None

logger = logging.getLogger(__name__)

class S3Finalizer:
    """
    Abstracción inteligente para manejar las operaciones en S3.
    Evita el molesto comportamiento por defecto de Spark limpiando archivos 
    intermedios y renombrando la salida.
    """
    
    def __init__(self, bucket_name: str, aws_region: str = 'us-east-1', mock_s3_client=None):
        self.bucket = bucket_name
        self.region = aws_region
        
        if mock_s3_client:
            self.s3 = mock_s3_client
        elif boto3:
            self.s3 = boto3.client('s3', region_name=self.region)
        else:
            raise ImportError("La librería 'boto3' no está instalada.")
            
    def list_keys(self, prefix: str) -> List[str]:
        keys = []
        paginator = self.s3.get_paginator('list_objects_v2')
        pages = paginator.paginate(Bucket=self.bucket, Prefix=prefix)
        for page in pages:
            for obj in page.get('Contents', []):
                keys.append(obj['Key'])
        return keys

    def _matches_rule(self, basename: str, starts_with: str, ends_with: str, contains: str) -> bool:
        """Filtro lógico genérico interno."""
        if basename.startswith("_SUCCESS"): return False # Basura de Big Data ignorable siempre
        
        # Filtros estrictos si el usuario los provee
        if starts_with and not basename.startswith(starts_with): return False
        if ends_with and not basename.endswith(ends_with): return False
        if contains and contains not in basename: return False
        
        # Si al menos un filtro de los 3 se cumplió o no proporcionó filtros, lo aprobamos
        if not starts_with and not ends_with and not contains:
            # Precaución: si no puso reglas, y no es un folder base, devolver True
            return True if basename else False
            
        return True

    def rename_file(self, s3_prefix: str, final_filename: str, starts_with: str = "", ends_with: str = "", contains: str = "", max_retries: int = 3) -> bool:
        """
        Archivo único: Encuentra un solo archivo que cumpla las reglas y lo renombra fijamente.
        """
        if not s3_prefix.endswith('/'): s3_prefix += '/'

        all_keys = []
        for attempt in range(max_retries):
            all_keys = self.list_keys(s3_prefix)
            if any(self._matches_rule(k.split("/")[-1], starts_with, ends_with, contains) for k in all_keys):
                break
            time.sleep(2)
            
        target_key = None
        keys_to_delete = []
        
        for key in all_keys:
            basename = key.split("/")[-1]
            if self._matches_rule(basename, starts_with, ends_with, contains) and not target_key:
                target_key = key
            elif basename != final_filename and basename: # Nunca borrar el prefijo "/"
                keys_to_delete.append({'Key': key})
                
        if not target_key:
            logger.warning(f"No se encontró ningún archivo bajo los criterios dados en {self.bucket}/{s3_prefix}")
            return False
            
        new_key = f"{s3_prefix}{final_filename}"
        try:
            self.s3.copy_object(Bucket=self.bucket, CopySource={'Bucket': self.bucket, 'Key': target_key}, Key=new_key)
            keys_to_delete.append({'Key': target_key})
            logger.info(f"Éxito: Archivo único renombrado a s3://{self.bucket}/{new_key}")
            self._batch_delete(keys_to_delete)
            return True
        except Exception as e:
            logger.error(f"Fallo al renombrar {s3_prefix}: {str(e)}")
            return False

    def sequence_files(self, s3_prefix: str, pattern: str, starts_with: str = "", ends_with: str = "", contains: str = "", max_retries: int = 3, max_workers: int = 20) -> int:
        """
        Secuenciador Universal (Archivos Múltiples): Escanea tu bucket e higieniza S3 renombando masivamente.
        """
        if not s3_prefix.endswith('/'): s3_prefix += '/'

        all_keys = []
        for attempt in range(max_retries):
            all_keys = self.list_keys(s3_prefix)
            if any(self._matches_rule(k.split("/")[-1], starts_with, ends_with, contains) for k in all_keys):
                break
            time.sleep(2)

        target_keys = []
        keys_to_delete = []

        # Filtrar exactamente lo que pidió el ingeniero de datos
        for key in all_keys:
            basename = key.split("/")[-1]
            if not basename: continue # ignora el directorio root
            
            if self._matches_rule(basename, starts_with, ends_with, contains):
                target_keys.append(key)
            else:
                keys_to_delete.append({'Key': key})

        if not target_keys:
            logger.warning(f"No hay archivos coincidentes bajo los filtros actuales en {s3_prefix}")
            return 0

        # Ordenar lexicológicamente para asegurar una secuenciación limpia y determinista
        target_keys.sort()
        
        rename_tasks = []
        for i, old_target_key in enumerate(target_keys, start=1):
            try:
                new_basename = pattern.format(seq=i)
            except KeyError:
                new_basename = pattern.replace("{seq}", str(i))
                
            new_key = f"{s3_prefix}{new_basename}"
            rename_tasks.append((old_target_key, new_key))
            keys_to_delete.append({'Key': old_target_key})

        successful_renames = 0

        def _copy(task):
            src, dst = task
            try:
                self.s3.copy_object(Bucket=self.bucket, CopySource={'Bucket': self.bucket, 'Key': src}, Key=dst)
                return True
            except Exception as e:
                logger.error(f"Fallo copiando {src} a {dst}: {e}")
                return False

        # Procesamiento Paralelo en AWS
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            results = executor.map(_copy, rename_tasks)
            successful_renames = sum(1 for r in results if r)

        logger.info(f"Éxito: {successful_renames}/{len(target_keys)} archivos secuenciados orgánicamente.")
        self._batch_delete(keys_to_delete)
        return successful_renames
        
    def _batch_delete(self, keys: List[dict]):
        if not keys: return
        for i in range(0, len(keys), 1000):
            try: # maximoAWS permite 1000 borrados por batch
                self.s3.delete_objects(Bucket=self.bucket, Delete={'Objects': keys[i:i+1000], 'Quiet': True})
            except Exception as e:
                logger.error(f"Fallo limpiando basura S3: {e}")
