# process_docs.py
import os
from dotenv import load_dotenv

# Importações do LangChain
from langchain_chroma import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

load_dotenv()

# --- Configuração ---
UPLOAD_FOLDER = 'uploads'
VECTOR_STORE_PATH = 'vector_store'
PROCESSED_FILES_LOG = 'processed_files.log'
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')

def get_processed_files():
    """Lê o log de arquivos já processados."""
    if not os.path.exists(PROCESSED_FILES_LOG):
        return set()
    with open(PROCESSED_FILES_LOG, 'r') as f:
        return set(f.read().splitlines())

def log_processed_file(filename):
    """Adiciona um arquivo ao log de processados."""
    with open(PROCESSED_FILES_LOG, 'a') as f:
        f.write(filename + '\n')

def process_new_documents():
    """
    Verifica a pasta de uploads por novos arquivos, processa-os e os adiciona ao Vector Store.
    """
    print("Iniciando verificação de novos documentos...")
    
    # --- Inicialização dos componentes LangChain ---
    embedding = OpenAIEmbeddings()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

    # Carrega ou cria o Vector Store
    if os.path.exists(VECTOR_STORE_PATH):
        print("Carregando Vector Store existente...")
        vector_store = Chroma(persist_directory=VECTOR_STORE_PATH, embedding_function=embedding)
    else:
        print("Criando novo Vector Store...")
        # Precisa de um documento inicial para ser criado, então o criamos vazio
        # A biblioteca pode exigir uma criação inicial com dados, então podemos ajustar se necessário
        # A abordagem mais segura é garantir que ele seja criado com o primeiro documento.
        vector_store = None

    processed_files = get_processed_files()
    new_files_found = False

    for filename in os.listdir(UPLOAD_FOLDER):
        if filename.endswith('.pdf') and filename not in processed_files:
            new_files_found = True
            print(f"Novo arquivo encontrado: {filename}")
            file_path = os.path.join(UPLOAD_FOLDER, filename)
            
            try:
                # 1. Carregar o novo documento
                loader = PyPDFLoader(file_path)
                documents = loader.load()

                # 2. Dividir em chunks
                chunks = text_splitter.split_documents(documents=documents)
                
                # 3. Adicionar ao Vector Store
                if vector_store is None:
                    # Primeira execução, cria o store
                    vector_store = Chroma.from_documents(
                        documents=chunks,
                        embedding=embedding,
                        persist_directory=VECTOR_STORE_PATH
                    )
                else:
                    # Adiciona os novos chunks ao store existente
                    vector_store.add_documents(chunks)
                
                print(f"Arquivo '{filename}' processado e adicionado à base de conhecimento.")
                log_processed_file(filename)

            except Exception as e:
                print(f"Erro ao processar o arquivo {filename}: {e}")

    if not new_files_found:
        print("Nenhum documento novo para processar.")
    else:
        print("Processamento concluído.")

if __name__ == '__main__':
    process_new_documents()