document_creator.py

import os
import re
from pathlib import Path
from llama_index.core import Document

def create_single_document_with_filenames(directory_path: str) -> Document:
    """
    Percorre a pasta informada, organiza os arquivos em estrutura de anos (YYYY) e meses (MM),
    gera um texto descritivo semelhante ao código anterior (sem salvar em arquivo) e retorna
    esse texto dentro de um objeto Document.

    Uso:
        directory_path = "documentos"
        doc = create_single_document_with_filenames(directory_path)
        documents.append(doc)
    """

    # Dicionário {ano: {mes: [arquivos]}}
    estrutura_anos = {}
    # Lista para arquivos fora do padrão YYYY-MM
    docs_sem_data = []
    # Conjunto para todos os arquivos (para listagem geral no final)
    todos_arquivos = set()

    # Percorre o diretório de forma recursiva
    for root, dirs, files in os.walk(directory_path):
        match_ano_mes = re.search(r"(\d{4})-(\d{2})", root)
        if match_ano_mes:
            ano = match_ano_mes.group(1)
            mes = match_ano_mes.group(2)

            if ano not in estrutura_anos:
                estrutura_anos[ano] = {}
            if mes not in estrutura_anos[ano]:
                estrutura_anos[ano][mes] = []

            for nome_arq in files:
                estrutura_anos[ano][mes].append(nome_arq)
                todos_arquivos.add(nome_arq)
        else:
            # Se a pasta não segue o padrão YYYY-MM, consideramos esses arquivos "sem data"
            for nome_arq in files:
                docs_sem_data.append(nome_arq)
                todos_arquivos.add(nome_arq)

    # Montamos o texto final, seguindo a lógica dos "casos"
    descricao_final = []

    # Organiza a lista de anos e percorre
    for ano in sorted(estrutura_anos.keys()):
        meses_ordenados = sorted(estrutura_anos[ano].keys())
        qtd_meses = len(meses_ordenados)

        # Caso 1: Apenas um mês e um arquivo
        if qtd_meses == 1:
            unico_mes = meses_ordenados[0]
            arquivos_unico_mes = estrutura_anos[ano][unico_mes]
            if len(arquivos_unico_mes) == 1:
                nome_mes_extenso = mes_extenso(unico_mes)
                arquivo_unico = arquivos_unico_mes[0]
                descricao_final.append(
                    f"No ano de {ano}, temos somente o mês de {nome_mes_extenso} (mês {unico_mes}), "
                    f"outros meses não foram listados, e nessa pasta encontramos apenas "
                    f"um manual chamado {arquivo_unico}."
                )
                # pula para o próximo ano
                continue

        # Caso 2: Mais meses ou mais arquivos em algum mês
        frase_inicial = f"No ano de {ano}, temos os meses de "
        meses_descricao = [f"{mes_extenso(m)} (mês {m})" for m in meses_ordenados]
        frase_inicial += ", ".join(meses_descricao) + "."
        descricao_final.append(frase_inicial)

        for mes_ in meses_ordenados:
            arquivos_mes = estrutura_anos[ano][mes_]
            nome_mes_extenso = mes_extenso(mes_)
            qtd_arquivos = len(arquivos_mes)
            if qtd_arquivos == 1:
                descricao_final.append(
                    f"Em {nome_mes_extenso} temos somente o manual chamado {arquivos_mes[0]}."
                )
            else:
                descricao_final.append(
                    f"Em {nome_mes_extenso} temos {qtd_arquivos} manuais, chamados: {', '.join(arquivos_mes)}."
                )

    # Caso 3: Arquivos sem data
    if docs_sem_data:
        docs_sem_data_unicos = list(set(docs_sem_data))
        if len(docs_sem_data_unicos) == 1:
            descricao_final.append(
                f"Em nossos documentos, fora de pastas de data, temos somente este arquivo: {docs_sem_data_unicos[0]}."
            )
        else:
            descricao_final.append(
                "Em nossos documentos, fora de pastas de data, encontramos estes arquivos: "
                + ", ".join(docs_sem_data_unicos) + "."
            )

    # Lista global de todos os arquivos
    lista_geral_ordenada = sorted(todos_arquivos)
    descricao_final.append(
        "Essas são as listas de todos os documentos, manuais e organogramas que podemos "
        f"resolver, listar e usar para nossas respostas, esses documentos vão ser muito úteis: {', '.join(lista_geral_ordenada)}"
    )

    # Une o texto final
    document_text = "\n".join(descricao_final)

    # Cria e retorna um Document com esse texto
    document = Document(
        text=document_text,
        metadata={
            "description": "Lista de manuais, arquivos e documentos que podemos responder.",
            "file_name": "Lista de documentos",
            "summary": "Entre 2011 e 2024, há uma grande variedade de manuais e documentos organizados por ano e mês: alguns anos possuem apenas um mês e um único arquivo, enquanto outros registram diversos períodos, cada um com vários manuais. Adicionalmente, há alguns arquivos “soltos”, fora dessas pastas de data. No fim, existe uma listagem completa de todos os itens, cobrindo instruções bancárias, regulamentos internos e outros materiais de suporte."
            }
    )
    return document

def mes_extenso(mes_str: str) -> str:
    """
    Converte '07' em 'Julho', '09' em 'Setembro', etc.
    """
    meses_dict = {
        "01": "Janeiro",   "02": "Fevereiro", "03": "Março",
        "04": "Abril",     "05": "Maio",      "06": "Junho",
        "07": "Julho",     "08": "Agosto",    "09": "Setembro",
        "10": "Outubro",   "11": "Novembro",  "12": "Dezembro"
    }
    return meses_dict.get(mes_str, mes_str)