import { filterCodeBlockLines, filterEnglishLinesAtEnd, filterEnglishLinesAtStart, filterLeadingAndTrailingNewLineInsertion, skipLines, stopAtLines, } from "../autocomplete/lineStream.js";
import { streamDiff } from "../diff/streamDiff.js";
import { streamLines } from "../diff/util.js";
import { gptEditPrompt } from "../llm/templates/edit.js";
import { Telemetry } from "./posthog.js";
function constructPrompt(prefix, highlighted, suffix, llm, userInput, language) {
    const template = llm.promptTemplates?.edit ?? gptEditPrompt;
    return llm.renderPromptTemplate(template, [], {
        userInput,
        prefix,
        codeToEdit: highlighted,
        suffix,
        language: language ?? "",
    });
}
async function* addIndentation(diffLineGenerator, indentation) {
    for await (const diffLine of diffLineGenerator) {
        yield {
            ...diffLine,
            line: indentation + diffLine.line,
        };
    }
}
function modelIsInept(model) {
    return !(model.includes("gpt") || model.includes("claude"));
}
function isGpt4Trial(model, provider) {
    return provider === "free-trial" && model.startsWith("gpt-4");
}
export async function* streamDiffLines(prefix, highlighted, suffix, llm, input, language, onlyOneInsertion) {
    Telemetry.capture("inlineEdit", {
        model: llm.model,
        provider: llm.providerName,
    });
    // Strip common indentation for the LLM, then add back after generation
    let oldLines = highlighted.length > 0
        ? highlighted.split("\n")
        : // When highlighted is empty, we need to combine last line of prefix and first line of suffix to determine the line being edited
            [(prefix + suffix).split("\n")[prefix.split("\n").length - 1]];
    // But if that line is empty, we can assume we are insertion-only
    if (oldLines.length === 1 && oldLines[0].trim() === "") {
        oldLines = [];
    }
    const prompt = constructPrompt(prefix, highlighted, suffix, llm, input, language);
    const inept = modelIsInept(llm.model);
    const options = {};
    if (isGpt4Trial(llm.model, llm.providerName)) {
        options.maxTokens = 2048;
    }
    const completion = typeof prompt === "string"
        ? llm.streamComplete(prompt, { raw: true })
        : llm.streamChat(prompt);
    let lines = streamLines(completion);
    lines = filterEnglishLinesAtStart(lines);
    lines = filterCodeBlockLines(lines);
    lines = stopAtLines(lines);
    lines = skipLines(lines);
    if (inept) {
        // lines = fixCodeLlamaFirstLineIndentation(lines);
        lines = filterEnglishLinesAtEnd(lines);
    }
    let diffLines = streamDiff(oldLines, lines);
    diffLines = filterLeadingAndTrailingNewLineInsertion(diffLines);
    if (highlighted.length === 0) {
        const line = prefix.split("\n").slice(-1)[0];
        const indentation = line.slice(0, line.length - line.trimStart().length);
        diffLines = addIndentation(diffLines, indentation);
    }
    let seenGreen = false;
    for await (const diffLine of diffLines) {
        yield diffLine;
        if (diffLine.type === "new") {
            seenGreen = true;
        }
        else if (onlyOneInsertion && seenGreen && diffLine.type === "same") {
            break;
        }
    }
}
