import pandas as pd
import os
import logging
import json
import PyLibrary.funciones_comunes as fc
from openai import OpenAI
from openai_api_key import openai_api_key

# Definir el logger a nivel de módulo
logger = None

def corregir_texto_batch(text_pairs, logger):
    """
    Corrige errores de OCR en lotes de pares de texto usando GPT-4.
    Cada par consiste en dos versiones del mismo texto en diferentes idiomas.
    
    Args:
        text_pairs: Lista de tuplas (texto1, texto2) donde cada texto es una lectura OCR
        logger: Logger para registrar el proceso
    """
    client = OpenAI(api_key=openai_api_key())
    
    system_prompt = """You are an OCR correction specialist. For each pair of texts, which are OCR readings of the same text in different languages, analyze and correct any obvious OCR errors.
    Pay special attention to:
    - Incorrectly joined words (missing spaces)
    - Wrong character recognition (0 vs O, 1 vs I, etc.)
    - Extra or missing characters
    Return the corrected versions maintaining the general structure and meaning.
    
    Input format: List of text pairs
    Expected output format: List of corrected pairs in the same order
    Example input: [["PULSANTE DI STARTNASTRI", "START PUSHBUTTONTTOP"]]
    Example output: [["PULSANTE DI START NASTRI", "START PUSH BUTTON TOP"]]"""

    try:
        # Convertir los pares de texto a formato JSON
        request_payload = json.dumps({"pairs": text_pairs})
        logger.info(f"Solicitando corrección para el lote de textos:\n{request_payload}")

        response = client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": request_payload},
            ],
            max_tokens=2000,
            temperature=0.3,
        )

        # Procesar la respuesta
        corrections = json.loads(response.choices[0].message.content)
        logger.info(f"Correcciones recibidas:\n{corrections}")
        return corrections["pairs"]
    except Exception as e:
        logger.error(f"Error en la corrección por lotes: {str(e)}")
        return None

def procesar_archivo(config, archivo_entrada):
    """
    Procesa el archivo Excel con los textos OCR y genera un nuevo archivo con las correcciones.
    """
    logger.info(f"Iniciando procesamiento de {archivo_entrada}")
    
    try:
        # Leer el archivo de entrada
        df = fc.read_dataframe_with_cleanup_retries(archivo_entrada)
        
        if config.columna_origen1 not in df.columns or config.columna_origen2 not in df.columns:
            logger.error(f"Columnas requeridas no encontradas en el archivo")
            print(f"Error: Las columnas {config.columna_origen1} y/o {config.columna_origen2} no existen en el archivo")
            return
        
        # Filtrar filas donde ambas columnas tienen contenido
        df_filtered = df.dropna(subset=[config.columna_origen1, config.columna_origen2])
        df_filtered = df_filtered[
            (df_filtered[config.columna_origen1].astype(str).str.strip() != '') & 
            (df_filtered[config.columna_origen2].astype(str).str.strip() != '')
        ]
        
        # Crear columnas para textos corregidos
        df[f"{config.columna_origen1}_Corregido"] = df[config.columna_origen1]
        df[f"{config.columna_origen2}_Corregido"] = df[config.columna_origen2]
        
        # Procesar en lotes de 10 pares
        batch_size = 10
        total_rows = len(df_filtered)
        
        progress_bar = fc.ProgressBar(
            total_rows, prefix="Procesando textos:", suffix="Completado"
        )
        
        for i in range(0, total_rows, batch_size):
            batch_df = df_filtered.iloc[i:i+batch_size]
            
            # Preparar pares de textos para el lote
            text_pairs = [
                [str(row[config.columna_origen1]).strip(), str(row[config.columna_origen2]).strip()]
                for _, row in batch_df.iterrows()
            ]
            
            # Obtener correcciones para el lote
            corrections = corregir_texto_batch(text_pairs, logger)
            
            if corrections:
                # Aplicar correcciones al DataFrame original
                for j, correction in enumerate(corrections):
                    idx = batch_df.index[j]
                    df.at[idx, f"{config.columna_origen1}_Corregido"] = correction[0]
                    df.at[idx, f"{config.columna_origen2}_Corregido"] = correction[1]
            
            progress_bar.update(min(i + batch_size, total_rows))
        
        progress_bar.finish()
        
        # Guardar resultados
        output_path = config.get_output_path()
        fc.save_dataframe_with_retries(df, output_path)
        
        logger.info(f"Archivo procesado y guardado en: {output_path}")
        print(f"\nArchivo procesado y guardado en: {output_path}")
        print(f"Se procesaron {total_rows} pares de texto.")
        
    except Exception as e:
        logger.error(f"Error durante el procesamiento: {str(e)}")
        print(f"Error durante el procesamiento: {str(e)}")

def run(config):
    global logger
    logger = fc.configurar_logger(config.work_dir)
    script_name = os.path.basename(__file__)
    print(f"\rIniciando: {script_name}\r")

    # Solicitar archivo de entrada
    from tkinter import filedialog
    archivo_entrada = filedialog.askopenfilename(
        title="Seleccione el archivo con textos OCR",
        filetypes=[("Excel files", "*.xls*")]
    )

    if archivo_entrada:
        procesar_archivo(config, archivo_entrada)
    else:
        print("No se seleccionó ningún archivo para procesar.")

if __name__ == "__main__":
    import menu_ocr_correction
    menu_ocr_correction.main()