import gradio as gr
from docx import Document
import os

def split_by_headers(file_path, headers_per_chunk=1):
    doc = Document(file_path)
    chunks = []
    current_chunk = Document()
    header_count = 0
    
    for element in doc.element.body:
        if element.tag.endswith('p'):
            paragraph = element._element
            if any(style.val.startswith('Heading') for style in paragraph.xpath('.//w:pStyle')):
                header_count += 1
                if header_count > headers_per_chunk:
                    chunks.append(current_chunk)
                    current_chunk = Document()
                    header_count = 1
            
            current_chunk.element.body.append(element)
    
    if len(current_chunk.element.body):
        chunks.append(current_chunk)
    
    return chunks

def split_by_pages(file_path, pages_per_chunk=1):
    doc = Document(file_path)
    chunks = []
    current_chunk = Document()
    page_count = 0
    estimated_chars_per_page = 3000  # This is an estimation
    char_count = 0
    
    for element in doc.element.body:
        if element.tag.endswith('p'):
            text = element.text
            char_count += len(text)
            
            if char_count >= estimated_chars_per_page:
                page_count += 1
                char_count = 0
                
                if page_count >= pages_per_chunk:
                    chunks.append(current_chunk)
                    current_chunk = Document()
                    page_count = 0
            
            current_chunk.element.body.append(element)
    
    if len(current_chunk.element.body):
        chunks.append(current_chunk)
    
    return chunks

def save_chunks(chunks, original_filename):
    saved_files = []
    base_name = os.path.splitext(original_filename)[0]
    
    for i, chunk in enumerate(chunks, 1):
        output_path = f"{base_name}_part{i}.docx"
        chunk.save(output_path)
        saved_files.append(output_path)
    
    return saved_files

def process_document(file, split_type, headers_or_pages):
    if headers_or_pages < 1:
        return "Por favor, especifique un número positivo de encabezados o páginas por fragmento."
    
    if split_type == "Encabezados":
        chunks = split_by_headers(file.name, headers_or_pages)
    else:  # Páginas
        chunks = split_by_pages(file.name, headers_or_pages)
    
    saved_files = save_chunks(chunks, os.path.basename(file.name))
    return f"Documento dividido en {len(saved_files)} partes: {', '.join(saved_files)}"

# Interfaz Gradio
iface = gr.Interface(
    fn=process_document,
    inputs=[
        gr.File(label="Seleccione el archivo DOCX"),
        gr.Radio(["Encabezados", "Páginas"], label="Método de división"),
        gr.Number(value=1, label="Número de encabezados/páginas por fragmento", minimum=1)
    ],
    outputs=gr.Text(label="Resultado"),
    title="Divisor de Documentos DOCX",
    description="Divida documentos DOCX por encabezados o páginas estimadas"
)

if __name__ == "__main__":
    iface.launch()