import torch
import time
import gc
from datetime import timedelta
import soundfile as sf
import sys
import os
import glob

# Obtener la ruta a la carpeta AUDIO (un nivel arriba)
audio_folder = os.path.join(os.path.dirname(os.path.dirname(__file__)), "AUDIO")

# Buscar todos los archivos MP3 en la carpeta AUDIO
mp3_files = glob.glob(os.path.join(audio_folder, "*.mp3"))

if not mp3_files:
    print("Error: No se encontraron archivos MP3 en la carpeta AUDIO")
    print(f"Buscando en: {audio_folder}")
    sys.exit(1)

# Si hay múltiples archivos MP3, usar el primero
mp3_path = mp3_files[0]
print(f"Se procesará el archivo: {mp3_path}")

def format_timestamp(seconds):
    """Convert seconds to SRT timestamp format"""
    td = timedelta(seconds=seconds)
    hours = td.seconds // 3600
    minutes = (td.seconds // 60) % 60
    seconds = td.seconds % 60
    milliseconds = td.microseconds // 1000
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"

def print_gpu_status():
    print(f"\nMemoria GPU usada: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
    print(f"Memoria GPU reservada: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

def get_audio_duration(audio_path):
    """Obtener la duración del archivo de audio"""
    with sf.SoundFile(audio_path) as audio_file:
        return len(audio_file) / audio_file.samplerate

def create_srt_content(chunks):
    """Create SRT formatted content from chunks"""
    srt_content = []
    for i, chunk in enumerate(chunks, start=1):
        # Extraer timestamps y texto del chunk
        start_time = format_timestamp(chunk['timestamp'][0])
        end_time = format_timestamp(chunk['timestamp'][1])
        text = chunk['text'].strip()
        
        # Mostrar en tiempo real el segmento que se está procesando
        print(f"\nSegmento {i}:")
        print(f"{start_time} --> {end_time}")
        print(f"{text}\n")
        
        srt_entry = f"{i}\n{start_time} --> {end_time}\n{text}\n"
        srt_content.append(srt_entry)
    
    return "\n".join(srt_content)

def main():
    try:
        # Definir el directorio de salida para los archivos SRT
        output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "DATA")
        
        # Crear el directorio si no existe
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # Definir la ruta completa del archivo SRT de salida
        output_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(mp3_path))[0]}.srt")
        
        print(f"\nProcesando archivo: {mp3_path}")
        print(f"La transcripción se guardará en: {output_path}")
        
        # Limpiar memoria GPU
        torch.cuda.empty_cache()
        gc.collect()
        
        # Configurar para mejor rendimiento
        torch.backends.cudnn.benchmark = True
        device = "cuda:0"
        
        print("\n=== Configuración Inicial ===")
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print_gpu_status()
        
        # Obtener duración del audio
        duration = get_audio_duration(mp3_path)
        print(f"\nDuración del audio: {duration:.2f} segundos")
        
        print("\nCargando modelo...")
        start_time = time.time()
        
        from transformers import pipeline
        from transformers.utils import is_flash_attn_2_available
        
        # Configurar el pipeline de Whisper
        pipe = pipeline(
            "automatic-speech-recognition",
            model="openai/whisper-large-v3",
            torch_dtype=torch.float32,
            device=device,
            model_kwargs={
                "attn_implementation": "flash_attention_2"
            } if is_flash_attn_2_available() else {
                "attn_implementation": "sdpa"
            },
        )
        
        load_time = time.time() - start_time
        print(f"Modelo cargado en {load_time:.2f} segundos")
        print_gpu_status()
        
        print("\nIniciando transcripción...")
        print("Los segmentos se mostrarán en tiempo real mientras se procesan...")
        start_time = time.time()
        
        # Realizar la transcripción
        outputs = pipe(
            mp3_path,
            chunk_length_s=30,
            batch_size=24,
            return_timestamps=True,
            generate_kwargs={"language": "es"}
        )
        
        transcription_time = time.time() - start_time
        print(f"\nTranscripción completada en {transcription_time:.2f} segundos")
        print_gpu_status()
        
        # Crear y guardar contenido SRT final
        print("\nGenerando archivo SRT final...")
        srt_content = create_srt_content(outputs["chunks"])
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(srt_content)
        
        # Confirmar la creación del archivo SRT
        if os.path.isfile(output_path):
            print(f"Archivo SRT creado exitosamente: {output_path}")
        else:
            print("Error: No se pudo crear el archivo SRT.")
        
        # Estadísticas finales
        print("\n=== Estadísticas de la transcripción ===")
        print(f"Archivo de entrada: '{mp3_path}'")
        print(f"Archivo de salida: '{output_path}'")
        print(f"Número total de segmentos: {len(outputs['chunks'])}")
        print(f"Velocidad de procesamiento: {duration/transcription_time:.2f}x")
        print(f"Tiempo total de procesamiento: {transcription_time:.2f} segundos")
        
        # Limpiar
        del pipe
        torch.cuda.empty_cache()
        gc.collect()
        print("\nLimpieza completada")
        print_gpu_status()

    except FileNotFoundError as e:
        print(f"\nError: {str(e)}")
        print("Asegúrate de que existe un archivo MP3 en la carpeta AUDIO")
    except Exception as e:
        print(f"\nError: {str(e)}")
        print("\nDetalles completos del error:")
        import traceback
        print(traceback.format_exc())

if __name__ == "__main__":
    main() tengo estos warnings en el cmd: Iniciando transcripción...
Los segmentos se mostrarán en tiempo real mientras se procesan...
C:\Users\Eurobloke\AppData\Local\Programs\Python\Python310\lib\site-packages\transformers\models\whisper\generation_whisper.py:509: FutureWarning: The input name `inputs` is deprecated. Please make sure to use `input_features` instead.
  warnings.warn(
You have passed language=es, but also have set `forced_decoder_ids` to [[1, None], [2, 50360]] which creates a conflict. `forced_decoder_ids` will be ignored in favor of language=es.
The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.