import "dotenv/config";
import { PDFLoader } from "@langchain/community/document_loaders/fs/pdf";
import * as fs from "fs/promises";
import * as path from "path";
import { Pinecone } from "@pinecone-database/pinecone";
import { GoogleGenerativeAIEmbeddings } from "@langchain/google-genai";
import { PineconeStore } from "@langchain/pinecone";


async function ingest() {
    console.log("Loading pdf files...");
    const invoicesDir = path.join(process.cwd(), "invoices");

    const files = await fs.readdir(invoicesDir);
    const pdfFiles = files.filter((file) => file.toLowerCase().endsWith(".pdf"));

    console.log(`Found ${pdfFiles.length} PDF file(s) to process.`);

    const allDocs = [];

    for (const file of pdfFiles) {
        const fullPath = path.join(invoicesDir, file);
        console.log(fullPath);

        const loader = new PDFLoader(fullPath, {
            splitPages: false, // Combines all pages of the invoice into ONE document
        });

        const docs = await loader.load();

        // Attach metadata: which file name and invoice it belongs to
        docs.forEach((doc) => {
            doc.metadata.sourceFile = file;
        });

        allDocs.push(...docs);
        console.log(` Loaded: ${file} (${docs.length} page(s))`);
    }

    console.log("Initializing pinecone...");

    const pinecone = new Pinecone({
        apiKey: process.env.PINECONE_API_KEY!,
    });

    const pineconeIndex = pinecone.index(process.env.PINECONE_INDEX!);

    console.log("Clearing old data from Pinecone index...");
    try {
        await pineconeIndex.deleteAll();
    } catch {
        console.log("Index is newly created or already empty. Skipping delete...");
    }

    const embeddings = new GoogleGenerativeAIEmbeddings({
        model: process.env.AI_EMBEDDING_MODEL,
        apiKey: process.env.AI_API_KEY,
        outputDimensionality: parseInt(process.env.VECTOR_DIMENSIONS!, 10),
    });

    const vectorStore = new PineconeStore(embeddings, { pineconeIndex });

    const BATCH_SIZE = 25;
    for (let i = 0; i < allDocs.length; i += BATCH_SIZE) {
        const batch = allDocs.slice(i, i + BATCH_SIZE);
        const batchNum = Math.floor(i / BATCH_SIZE) + 1;
        const totalBatches = Math.ceil(allDocs.length / BATCH_SIZE);
        console.log(`Ingesting batch ${batchNum}/${totalBatches} (${batch.length} docs)...`);

        await vectorStore.addDocuments(batch);

        if (i + BATCH_SIZE < allDocs.length) {
            console.log("Waiting 30 seconds to stay within Google API quota...");
            await new Promise((resolve) => setTimeout(resolve, 30000));
        }
    }

    console.log("All pdf data has been ingested successfully!");
}

ingest().catch(console.log)