import os
import logging
from pathlib import Path
from typing import Dict, Union, List
from watermark import embed_watermark, extract_watermark

# Configurar logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def batch_embed(
    input_dir: str,
    output_dir: str,
    message: bytes,
    password: str,
    strength: float = 1e-4
) -> Dict[str, Union[int, str]]:
    """
    Recorre WAV/MP3 en input_dir, inserta watermark en cada uno, guarda en output_dir.
    Retorna dict {filename: bits_insertados | error_message}.
    """
    try:
        # Validaciones
        if not os.path.exists(input_dir):
            raise FileNotFoundError(f"Directorio de entrada no existe: {input_dir}")
        
        if not message:
            raise ValueError("Mensaje no puede estar vacío")
        
        if not password:
            raise ValueError("Contraseña requerida")
        
        # Crear directorio de salida
        os.makedirs(output_dir, exist_ok=True)
        
        results = {}
        input_path = Path(input_dir)
        output_path = Path(output_dir)
        
        # Buscar archivos de audio
        audio_files = []
        for ext in [".wav", ".mp3", ".WAV", ".MP3"]:
            audio_files.extend(input_path.glob(f"*{ext}"))
        
        if not audio_files:
            logger.warning(f"No se encontraron archivos de audio en {input_dir}")
            return {"warning": "No se encontraron archivos de audio"}
        
        logger.info(f"Procesando {len(audio_files)} archivos de audio")
        
        for audio_file in audio_files:
            try:
                output_file = output_path / audio_file.name
                bits_inserted = embed_watermark(
                    str(audio_file), 
                    str(output_file), 
                    message, 
                    password,
                    strength=strength
                )
                results[audio_file.name] = bits_inserted
                logger.info(f"✅ {audio_file.name}: {bits_inserted} bits insertados")
                
            except Exception as e:
                error_msg = f"Error: {str(e)}"
                results[audio_file.name] = error_msg
                logger.error(f"❌ {audio_file.name}: {error_msg}")
        
        return results
        
    except Exception as e:
        logger.error(f"Error en batch_embed: {str(e)}")
        return {"error": f"Error crítico en batch_embed: {str(e)}"}

def batch_extract(
    input_dir: str,
    password: str
) -> Dict[str, Union[bytes, str]]:
    """
    Extrae watermark de cada archivo en input_dir. 
    Retorna {filename: payload_bytes | error_message}.
    """
    try:
        # Validaciones
        if not os.path.exists(input_dir):
            raise FileNotFoundError(f"Directorio de entrada no existe: {input_dir}")
        
        if not password:
            raise ValueError("Contraseña requerida")
        
        results = {}
        input_path = Path(input_dir)
        
        # Buscar archivos de audio
        audio_files = []
        for ext in [".wav", ".mp3", ".WAV", ".MP3"]:
            audio_files.extend(input_path.glob(f"*{ext}"))
        
        if not audio_files:
            logger.warning(f"No se encontraron archivos de audio en {input_dir}")
            return {"warning": "No se encontraron archivos de audio"}
        
        logger.info(f"Procesando {len(audio_files)} archivos de audio")
        
        for audio_file in audio_files:
            try:
                payload = extract_watermark(str(audio_file), password)
                results[audio_file.name] = payload
                logger.info(f"✅ {audio_file.name}: payload extraído")
                
            except Exception as e:
                error_msg = f"Error: {str(e)}"
                results[audio_file.name] = error_msg
                logger.error(f"❌ {audio_file.name}: {error_msg}")
        
        return results
        
    except Exception as e:
        logger.error(f"Error en batch_extract: {str(e)}")
        return {"error": f"Error crítico en batch_extract: {str(e)}"}