import gradio as gr
import PyPDF2
import re

def extrair_texto_pdf():
    with open("dados-cdu.pdf", "rb") as pdf_file:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

def consultar_cdu(palavra_chave):
    texto_base = extrair_texto_pdf()
    palavras_chave_encontradas = []
    linhas = texto_base.split("\n")
    for linha in linhas:
        if palavra_chave.lower() in linha.lower():
            palavras_chave_encontradas.append(linha)
    return palavras_chave_encontradas

def gerar_cdu(palavras_chave):
    palavras_chave = [palavra.strip() for palavra in palavras_chave.split(",")]
    resultados = {}
    for palavra_chave in palavras_chave:
        cdu_encontrada = consultar_cdu(palavra_chave)
        cdu_formatada = [re.search(r"\[(.*?)\]", cdu).group(1) for cdu in cdu_encontrada]
        resultados[palavra_chave] = cdu_formatada
    return resultados

# Interface Gradio
input_text = gr.inputs.Textbox(label="Inserir palavras-chave separadas por vírgulas")
button_label = "Gerar CDU"
output_text = gr.outputs.Textbox(label="Resultado")

def generate_cdu(palavras_chave):
    if palavras_chave:
        resultados = gerar_cdu(palavras_chave)
        formatted_resultados = {}
        for palavra_chave, cdu_encontrada in resultados.items():
            formatted_resultados[palavra_chave] = "\n".join(["CDU: " + cdu for cdu in cdu_encontrada])
        return formatted_resultados
    else:
        return {}

title = "Gerador de CDU"
description = "Insira palavras-chave separadas por vírgulas e clique em 'Gerar CDU' para obter a classificação relacionada."

gr.Interface(fn=generate_cdu, inputs=input_text, outputs=output_text, title=title, description=description).launch()