agent-v1

import {
    type FileHandle,
    type PredictionProcessStatusController,
} from "@lmstudio/sdk";
import { embedCustomParsedFiles } from "./embedCustomParsedFiles";
import { parseFile } from "../parser/parseFile";
import { type PluginCapableCtl } from "../../../utils/shared/pluginCtl";

function isLikelyBoilerplate(line: string): boolean {
    const trimmed = line.trim();
    if (trimmed.length === 0) return true;

    // Pure URL
    if (/^https?:\/\/\S+$/.test(trimmed)) return true;

    // Pure page counter, e.g. "4/8"
    if (/^\d{1,3}\/\d{1,3}$/.test(trimmed)) return true;

    // Date/time stamp lines, e.g. "6/29/26, 10:22 AM" or "29 Jun 2026 08:22AM"
    if (/\d{1,2}[\/\s]\w+[\/\s]\d{2,4}.{0,15}\d{1,2}:\d{2}\s*(AM|PM)?/i.test(trimmed) && trimmed.length < 80) return true;

    // Copyright / legal footer
    if (/copyright|all rights reserved|terms (and|&) conditions|privacy policy/i.test(trimmed)) return true;

    // Nav-menu-like lines: many capitalized words mashed together with no
    // normal sentence punctuation, e.g. "Top StoriesLatest NewsSingaporeWorld..."
    // Heuristic: short-ish line, no period at end, has 4+ capital-letter "word starts"
    // packed close together (avg word length low, capital density high).
    const words = trimmed.split(/\s+/);
    const capStarts = words.filter(w => /^[A-Z]/.test(w)).length;
    const hasSentencePunctuation = /[.?!]"?$/.test(trimmed);
    if (!hasSentencePunctuation && words.length >= 4 && capStarts / words.length > 0.6 && trimmed.length < 120) {
        return true;
    }

    // Very short fragment lines unlikely to be real sentence content on their own
    if (trimmed.length < 15 && !/[.?!]$/.test(trimmed)) return true;

    return false;
}

export function cleanCitationText(text: string): string {
    const lines = text.split(/\n+/);
    const kept = lines.filter(line => !isLikelyBoilerplate(line));
    return kept.join(" ").replace(/\s+/g, " ").trim();
}

export async function prepareRetrievalResultsContextInjection(
    ctl: PluginCapableCtl,
    originalUserPrompt: string,
    files: Array<FileHandle>,
): Promise<string> {
    const retrievalLimit = ctl.getConfigValue("retrievalLimit");
    const retrievalAffinityThreshold = ctl.getConfigValue("retrievalAffinityThreshold");
    const statusSteps = new Map<FileHandle, PredictionProcessStatusController>();

    const retrievingStatus = ctl.createStatus({
        status: "loading",
        text: `Loading an embedding model for retrieval...`,
    });
    const model = await ctl.client.embedding.model("nomic-ai/nomic-embed-text-v1.5-GGUF", {
        signal: ctl.abortSignal,
    });

    const normalFiles: Array<FileHandle> = [];
    const customFiles: Array<{ file: FileHandle; content: string }> = [];

    for (const file of files) {
        const parsed = await parseFile(ctl, file);
        if (parsed.customParsed) {
            customFiles.push({ file, content: parsed.content });
        } else {
            normalFiles.push(file);
        }
    }
    retrievingStatus.setState({
        status: "loading",
        text: `Retrieving relevant citations for user query...`,
    });

    type SdkResult = Awaited<ReturnType<typeof ctl.client.files.retrieve>>;
    type SdkEntry = SdkResult["entries"][number];

    let sdkResult: SdkResult | null = null;

    // Each merged entry carries content/score for sorting+display, plus an
    // optional back-reference to the original (untouched) SDK entry object
    // so we can hand it back to addCitations with its original type intact.
    type MergedEntry = { content: string; score: number; sdkRef?: SdkEntry; sourceFile?: FileHandle };
    let entries: MergedEntry[] = [];

    if (normalFiles.length > 0) {
        sdkResult = await ctl.client.files.retrieve(originalUserPrompt, normalFiles, {
            embeddingModel: model,
            limit: retrievalLimit,
            signal: ctl.abortSignal,
            onFileProcessList(filesToProcess) {
                for (const file of filesToProcess) {
                    statusSteps.set(
                        file,
                        retrievingStatus.addSubStatus({
                            status: "waiting",
                            text: `Process ${file.name} for retrieval`,
                        }),
                    );
                }
            },
            onFileProcessingStart(file) {
                statusSteps.get(file)!.setState({ status: "loading", text: `Processing ${file.name} for retrieval` });
            },
            onFileProcessingEnd(file) {
                statusSteps.get(file)!.setState({ status: "done", text: `Processed ${file.name} for retrieval` });
            },
            onFileProcessingStepProgress(file, step, progressInStep) {
                const verb = step === "loading" ? "Loading" : step === "chunking" ? "Chunking" : "Embedding";
                statusSteps.get(file)!.setState({
                    status: "loading",
                    text: `${verb} ${file.name} for retrieval (${(progressInStep * 100).toFixed(1)}%)`,
                });
            },
        });
        entries = sdkResult.entries.map(e => ({ content: e.content, score: e.score, sdkRef: e }));
    }

    if (customFiles.length > 0) {
        const customSubStatus = retrievingStatus.addSubStatus({
            status: "loading",
            text: `Embedding custom-parsed content (${customFiles.map(f => f.file.name).join(", ")})...`,
        });
        const customEntries = await embedCustomParsedFiles(ctl, originalUserPrompt, customFiles, model);
        entries = entries.concat(
            customEntries.map(e => ({ content: e.content, score: e.score, sourceFile: e.file }))
        );
        customSubStatus.setState({ status: "done", text: `Embedded custom-parsed content` });
    }

    entries = entries
        .filter(entry => entry.score > retrievalAffinityThreshold)
        .sort((a, b) => b.score - a.score)
        .slice(0, retrievalLimit);

    let processedContent = "";
    const numRetrievals = entries.length;

    if (numRetrievals > 0) {
        const prefix = `Citations from the provided file${files.length > 1 ? "s" : ""}:\n`;
        processedContent += prefix;
        entries.forEach((entry, i) => {
            const sourceName = entry.sdkRef?.source?.name ?? entry.sourceFile?.name ?? "unknown source";
            const cleaned = cleanCitationText(entry.content);
            if (cleaned.length === 0) return;
            processedContent += `\n----------------------------------------\n`;
            processedContent += `[${i + 1}] ${cleaned}\n`;
            processedContent += `— ${sourceName}\n`;
        });
        processedContent += `\n----------------------------------------\n\n`;
        processedContent +=
            `Use the citations above to respond to the user query, only if they are relevant. ` +
            `Otherwise, respond to the best of your ability without them.`;
    } else {
        retrievingStatus.setState({
            status: "canceled",
            text: `No relevant citations found for user query`,
        });
        // processedContent =
        //     `Important: No citations were found in the user files for the user query. ` +
        //     `In less than one sentence, inform the user of this. ` +
        //     `Then respond to the query to the best of your ability.`;

        processedContent = ""
    }
    ctl.debug("Processed content", processedContent);

    return processedContent; // no longer includes "User Query:" — preprocess.ts adds it once, at the end
}