#!/usr/bin/env python
import os
import time
import json
import pika
import logging
import whisper
import torch
from datetime import datetime

# Configuración de logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('processor')

# Configuración de RabbitMQ
RABBITMQ_HOST = 'rabbitmq-app.whisper.svc.cluster.local'
RABBITMQ_USER = 'user'
RABBITMQ_PASS = 'password'
PROCESS_QUEUE = 'audio_process_queue'
UNIFY_QUEUE = 'text_unify_queue'

# Configuración de Whisper
WHISPER_MODEL = "base"  # Opciones: "tiny", "base", "small", "medium", "large"

# Directorios
SHARED_DIR = '/app/shared'

# ID del procesador para logs
PROCESSOR_ID = os.environ.get('PROCESSOR_ID', 'unknown')

def connect_to_rabbitmq():
    """Establece conexión con RabbitMQ"""
    tries = 0
    while True:
        try:
            credentials = pika.PlainCredentials(RABBITMQ_USER, RABBITMQ_PASS)
            connection = pika.BlockingConnection(
                pika.ConnectionParameters(
                    host=RABBITMQ_HOST,
                    credentials=credentials,
                    heartbeat=600  # 10 minutos
                )
            )
            return connection
        except pika.exceptions.AMQPConnectionError:
            tries += 1
            logger.warning(f"Intento {tries}: No se pudo conectar a RabbitMQ. Reintentando en 5 segundos...")
            time.sleep(5)

def load_whisper_model():
    """Carga el modelo de Whisper"""
    # Verificar si hay GPU disponible
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    if device == "cuda":
        gpu_name = torch.cuda.get_device_name(0)
        logger.info(f"Procesador {PROCESSOR_ID}: GPU detectada: {gpu_name}")
    else:
        logger.info(f"Procesador {PROCESSOR_ID}: No se detectó GPU, usando CPU")
    
    logger.info(f"Procesador {PROCESSOR_ID}: Cargando modelo Whisper '{WHISPER_MODEL}' en {device}...")
    model = whisper.load_model(WHISPER_MODEL, device=device)
    logger.info(f"Procesador {PROCESSOR_ID}: Modelo Whisper cargado correctamente")
    return model

def transcribe_audio(model, audio_path):
    """Transcribe un archivo de audio usando Whisper"""
    logger.info(f"Procesador {PROCESSOR_ID}: Transcribiendo {audio_path}")
    
    # Realizar transcripción
    result = model.transcribe(audio_path)
    
    logger.info(f"Procesador {PROCESSOR_ID}: Transcripción completada para {audio_path}")
    
    return result

def save_transcription(result, segment_info):
    """Guarda la transcripción en un archivo de texto"""
    segment_id = segment_info['segment_id']
    original_file_id = segment_info['original_file_id']
    segments_dir = segment_info['segments_dir']
    
    # Crear directorio para transcripciones
    transcriptions_dir = os.path.join(segments_dir, "transcriptions")
    os.makedirs(transcriptions_dir, exist_ok=True)
    
    # Generar nombre para archivo de transcripción
    transcription_filename = f"transcription_{segment_id:03d}_{original_file_id}.txt"
    transcription_path = os.path.join(transcriptions_dir, transcription_filename)
    
    # Guardar texto en archivo
    with open(transcription_path, 'w', encoding='utf-8') as f:
        f.write(result['text'])
    
    logger.info(f"Procesador {PROCESSOR_ID}: Transcripción guardada en {transcription_path}")
    
    return transcription_path

def send_to_unify_queue(channel, transcription_path, segment_info):
    """Envía la transcripción a la cola de unificación"""

    # Preparar mensaje
    message = {
        **segment_info,
        'transcription_path': transcription_path,
        'processor_id': PROCESSOR_ID,
        'processed_timestamp': datetime.now().isoformat()
    }
    
    # Publicar mensaje
    channel.basic_publish(
        exchange='',
        routing_key=UNIFY_QUEUE,
        body=json.dumps(message),
        properties=pika.BasicProperties(
            delivery_mode=2  # mensaje persistente
        )
    )
    
    logger.info(f"Procesador {PROCESSOR_ID}: Transcripción enviada a la cola de unificación")

def callback(ch, method, properties, body):
    """Callback para procesar mensajes de la cola de procesamiento"""
    try:
        # Decodificar mensaje
        segment_info = json.loads(body)
        segment_id = segment_info['segment_id']
        segment_path = segment_info['segment_path']
        original_file_id = segment_info['original_file_id']
        total_segments = segment_info['total_segments']
        
        logger.info(f"Procesador {PROCESSOR_ID}: Recibido segmento {segment_id+1}/{total_segments} para {original_file_id}")
        
        # Transcribir audio
        result = transcribe_audio(model, segment_path)
        
        # Guardar transcripción
        transcription_path = save_transcription(result, segment_info)
        
        # Enviar a cola de unificación
        send_to_unify_queue(ch, transcription_path, segment_info)
        
        # Confirmar procesamiento
        ch.basic_ack(delivery_tag=method.delivery_tag)
        
        logger.info(f"Procesador {PROCESSOR_ID}: Transcripción completada para segmento {segment_id+1}/{total_segments}")
        
    except Exception as e:
        logger.error(f"Procesador {PROCESSOR_ID}: Error procesando mensaje: {str(e)}")
        # Rechazar mensaje en caso de error para reintentarlo
        ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)

def main():
    """Función principal"""
    global model
    
    # Cargar modelo Whisper
    model = load_whisper_model()
    
    # Conexión a RabbitMQ
    connection = connect_to_rabbitmq()
    channel = connection.channel()
    
    # Declarar colas
    channel.queue_declare(queue=PROCESS_QUEUE, durable=True)
    channel.queue_declare(queue=UNIFY_QUEUE, durable=True)
    
    # Configurar prefetch count -> 1 mensaje a la vez
    channel.basic_qos(prefetch_count=1)
    
    # Configurar callback
    channel.basic_consume(queue=PROCESS_QUEUE, on_message_callback=callback)
    
    logger.info(f"Procesador {PROCESSOR_ID}: Servicio iniciado. Esperando segmentos...")
    
    # Iniciar consumo
    try:
        channel.start_consuming()
    except KeyboardInterrupt:
        logger.info(f"Procesador {PROCESSOR_ID}: Servicio detenido")
        channel.stop_consuming()
    
    connection.close()

if __name__ == "__main__":
    main()