import gradio as gr
from docx import Document
import os
import shutil
import zipfile

def split_by_headers(file_path, headers_per_chunk=1):
    doc = Document(file_path)
    chunks = []
    current_chunk = Document()
    header_count = 0
    
    for paragraph in doc.paragraphs:
        # Verificar si el párrafo es un encabezado
        if paragraph.style.name.startswith('Heading'):
            header_count += 1
            if header_count > headers_per_chunk:
                chunks.append(current_chunk)
                current_chunk = Document()
                header_count = 1
        
        # Copiar el párrafo al chunk actual
        current_chunk.add_paragraph(paragraph.text, style=paragraph.style.name)
    
    if len(current_chunk.paragraphs):
        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  # Esta es una estimación
    char_count = 0
    
    for paragraph in doc.paragraphs:
        text = paragraph.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.add_paragraph(text, style=paragraph.style.name)
    
    if len(current_chunk.paragraphs):
        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 zip_files(files, zip_name="document_parts.zip"):
    with zipfile.ZipFile(zip_name, 'w') as zipf:
        for file in files:
            zipf.write(file, os.path.basename(file))
    return zip_name

def process_document(file, split_type, headers_or_pages, download_type):
    if headers_or_pages < 1:
        return "Por favor, especifique un número positivo de encabezados o páginas por fragmento."
    
    try:
        # Dividir el documento según el tipo seleccionado
        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))
        
        # Si el usuario selecciona descargar en ZIP, comprimir los archivos
        if download_type == "ZIP":
            zip_path = zip_files(saved_files)
            return zip_path  # Devolver el archivo zip para descargar
        else:
            return saved_files  # Devolver los archivos separados para descargar
        
    except Exception as e:
        return f"Error al procesar el documento: {str(e)}"

# 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),
        gr.Radio(["Separados", "ZIP"], label="Tipo de descarga")
    ],
    outputs=gr.File(label="Descargar archivo(s)"),
    title="Divisor de Documentos DOCX",
    description="Divida documentos DOCX por encabezados o páginas estimadas y descárguelos como archivos separados o en un archivo ZIP."
)

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