import gradio as gr
import spacy
import pandas as pd
from docx import Document
from io import BytesIO
import tempfile
import os

# Cargar el modelo de SpaCy en español
nlp = spacy.load('zh_core_web_trf')
nlp.max_length = 15000000  # Aumenta el límite a 3 millones de caracteres

# Función para procesar bloques de texto y extraer nombres de personas
def extract_names_from_text(text):
    doc = nlp(text)
    # Extraer las entidades de tipo PERSON
    persons = [ent.text for ent in doc.ents if ent.label_ == 'PERSON']
    return persons

# Función para dividir el texto en fragmentos más pequeños
def split_text(text, max_length=15000000):
    paragraphs = text.split('\n\n')  # Divide el texto en párrafos
    result = []
    current_chunk = []

    for paragraph in paragraphs:
        # Si agregar el próximo párrafo no excede el límite, agrégalo al chunk actual
        if sum(len(p) for p in current_chunk) + len(paragraph) + (len(current_chunk) > 0) * 2 <= max_length:
            current_chunk.append(paragraph)
        else:
            # Si agregar el próximo párrafo excede el límite, guarda el chunk actual y comienza uno nuevo
            result.append('\n\n'.join(current_chunk))
            current_chunk = [paragraph]

    # Añade el último chunk si no está vacío
    if current_chunk:
        result.append('\n\n'.join(current_chunk))

    return result

# Función principal para extraer nombres de personas desde un archivo DOCX
def extract_names_from_docx(docx_file):
    # Cargar el archivo DOCX
    document = Document(docx_file)
    full_text = []
    for para in document.paragraphs:
        full_text.append(para.text)
    
    # Unir todo el texto
    text = ' '.join(full_text)
    
    # Dividir el texto en fragmentos si es necesario
    text_fragments = split_text(text)
    
    # Extraer los nombres de cada fragmento y combinarlos
    all_persons = []
    for fragment in text_fragments:
        persons = extract_names_from_text(fragment)
        all_persons.extend(persons)
    
    # Eliminar duplicados
    all_persons = list(set(all_persons))
    
    # Crear un DataFrame
    df = pd.DataFrame(all_persons, columns=['Nombres'])
    
    # Crear un archivo temporal para guardar el Excel
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, "nombres_personas.xlsx")
    
    # Guardar el DataFrame en un archivo Excel
    with pd.ExcelWriter(temp_file_path, engine='openpyxl') as writer:
        df.to_excel(writer, index=False)
    
    return temp_file_path  # Devolver la ruta del archivo temporal

# Interfaz de Gradio
iface = gr.Interface(
    fn=extract_names_from_docx,
    inputs=gr.File(file_types=[".docx"]),
    outputs=gr.File(),
    title="Extractor de Nombres",
    description="Sube un archivo .docx y extrae los nombres de las personas usando NLP con SpaCy. Descarga el resultado en un archivo Excel."
)

# Iniciar la aplicación
iface.launch()