import * as vscode from "vscode";
import axios from "axios";

let VSConfig: vscode.WorkspaceConfiguration;
let apiEndpoint: string;
let apiAuthentication: string;
let apiModel: string;
let apiTemperature: number;
let numPredict: number;
let promptWindowSize: number;
let completionKeys: string;
let responsePreview: boolean | undefined;
let responsePreviewMaxTokens: number;
let responsePreviewDelay: number;
let continueInline: boolean | undefined;
let keepAlive: number | undefined;
let topP: number | undefined;

function updateVSConfig() {
    VSConfig = vscode.workspace.getConfiguration("fabelous-autocoder");
    apiEndpoint = VSConfig.get("endpoint") || "http://localhost:11434/api/generate";
    apiAuthentication = VSConfig.get("authentication") || "";
    apiModel = VSConfig.get("model") || "fabelous-coder:latest";
    numPredict = VSConfig.get("max tokens predicted") || 1000;
    promptWindowSize = VSConfig.get("prompt window size") || 2000;
    completionKeys = VSConfig.get("completion keys") || " ";
    responsePreview = VSConfig.get("response preview");
    responsePreviewMaxTokens = VSConfig.get("preview max tokens") || 50;
    responsePreviewDelay = VSConfig.get("preview delay") || 0;
    continueInline = VSConfig.get("continue inline");
    apiTemperature = VSConfig.get("temperature") || 0.7;
    keepAlive = VSConfig.get("keep alive") || 30;
    topP = VSConfig.get("top p") || 1;
}

updateVSConfig();
vscode.workspace.onDidChangeConfiguration(updateVSConfig);

function getContextLines(document: vscode.TextDocument, position: vscode.Position): string {
    const lines = [];
    const startLine = Math.max(0, position.line - 1);
    const endLine = position.line;

    for (let i = startLine; i <= endLine; i++) {
        lines.push(document.lineAt(i).text);
    }

    return lines.join("\n");
}


function createFIMPrompt(prefix: string, language: string): string {
    return `<fim_prefix>${prefix}<fim_middle><fim_suffix>${language}\n`;
}

const previewDecorationType = vscode.window.createTextEditorDecorationType({
    after: {
        color: '#888888', // Grayed-out preview text
        fontStyle: 'italic',
    },
    textDecoration: 'none; display: none;', // Hide the original text
});
// Generate extra lines for the preview
function generateExtraPreviewLines(document: vscode.TextDocument, position: vscode.Position, numLines: number): string[] {
    const extraLines = [];
    const startLine = position.line + 1;
    const endLine = Math.min(document.lineCount - 1, startLine + numLines - 1);

    for (let i = startLine; i <= endLine; i++) {
        extraLines.push(document.lineAt(i).text);
    }

    return extraLines;
}

async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) {
    const document = textEditor.document;
    const position = textEditor.selection.active;
    const contextLines = 2;
    const startLine = Math.max(0, position.line - contextLines);
    const context = getContextLines(document, position);

    const fimPrompt = createFIMPrompt(context, document.languageId);

    vscode.window.withProgress(
        {
            location: vscode.ProgressLocation.Notification,
            title: "Fabelous Autocoder",
            cancellable: true,
        },
        async (progress, progressCancellationToken) => {
            try {
                progress.report({ message: "Starting model..." });

                let axiosCancelPost: () => void;
                const axiosCancelToken = new axios.CancelToken((c) => {
                    axiosCancelPost = () => {
                        c("Autocompletion request terminated by user cancel");
                    };
                    if (cancellationToken) cancellationToken.onCancellationRequested(axiosCancelPost);
                    progressCancellationToken.onCancellationRequested(axiosCancelPost);
                    vscode.workspace.onDidCloseTextDocument(axiosCancelPost);
                });

                // Increase the number of tokens to predict
                const extendedNumPredict = numPredict * 2; // Adjust this multiplier as needed

                // Make the API request
                const response = await axios.post(apiEndpoint, {
                    model: apiModel,
                    prompt: fimPrompt,
                    stream: false,
                    raw: true,
                    options: {
                        num_predict: extendedNumPredict,
                        temperature: apiTemperature,
                        stop: ["<fim_suffix>", "```"]
                    }
                }, {
                    cancelToken: axiosCancelToken,
                    headers: {
                        'Authorization': apiAuthentication
                    }
                });

                progress.report({ message: "Generating..." });

                let completionText = response.data.response;
                completionText = completionText.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();

                // Split the completion text by new lines
                const newLines = completionText.split('\n');

                // Calculate the number of new lines in the completion
                const completionLineCount = newLines.length;

                // Ensure we have at least as many new lines as the completion, plus some extra
                const extraLines = 1; // You can adjust this number
                const totalNewLines = Math.max(completionLineCount + extraLines, position.line - startLine + 1);

                // Create preview decorations
                const previewRanges: vscode.DecorationOptions[] = [];
                for (let i = 0; i < totalNewLines; i++) {
                    const lineContent = i < newLines.length ? newLines[i] : '';
                    const range = new vscode.Range(startLine + i, 0, startLine + i + 1, 0);
                    previewRanges.push({
                        range,
                        renderOptions: {
                            after: {
                                contentText: lineContent,
                            }
                        }
                    });
                }

                textEditor.setDecorations(previewDecorationType, previewRanges);

                let previewInserted = true;

                // Handle preview acceptance or dismissal
                const disposable = vscode.workspace.onDidChangeTextDocument(async (event) => {
                    if (event.document.uri.toString() === document.uri.toString()) {
                        const change = event.contentChanges[0];

                        // Dismiss preview with Backspace
                        if (change && change.text === '' && change.rangeLength === 1) {
                            textEditor.setDecorations(previewDecorationType, []);
                            disposable.dispose();
                            previewInserted = false;
                        }

                        // Accept preview with Tab key
                        if (change && change.text === '\t' && previewInserted) {
                            textEditor.setDecorations(previewDecorationType, []);
                            const edit = new vscode.WorkspaceEdit();

                            // Create the text to insert: completion text plus extra newlines
                            const insertText = completionText + '\n'.repeat(Math.max(0, totalNewLines - completionLineCount));

                            // Replace the context with the new text
                            const replaceRange = new vscode.Range(startLine, 0, position.line, 0);
                            edit.replace(document.uri, replaceRange, insertText);

                            await vscode.workspace.applyEdit(edit);
                            await document.save(); // Optionally save after inserting

                            // Move the cursor to the end of the inserted completion
                            const newPosition = new vscode.Position(startLine + totalNewLines, 0);
                            textEditor.selection = new vscode.Selection(newPosition, newPosition);

                            disposable.dispose();
                            previewInserted = false;
                        }
                    }
                });

            } catch (err: any) {
                vscode.window.showErrorMessage(
                    "Fabelous Autocoder encountered an error: " + err.message
                );
                console.log(err);
            }
        }
    );
}



async function provideCompletionItems(document: vscode.TextDocument, position: vscode.Position, cancellationToken: vscode.CancellationToken) {
    const item = new vscode.CompletionItem("Fabelous autocompletion");
    item.insertText = new vscode.SnippetString('${1:}');

    if (responsePreview) {
        await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000));
        if (cancellationToken.isCancellationRequested) {
            return [item];
        }

        const context = getContextLines(document, position);
        const fimPrompt = createFIMPrompt(context, document.languageId);

        try {
            const response_preview = await axios.post(apiEndpoint, {
                model: apiModel,
                prompt: fimPrompt,
                stream: false,
                raw: true,
                options: {
                    num_predict: responsePreviewMaxTokens,
                    temperature: apiTemperature,
                    stop: ['<fim_suffix>', '\n', '```'],
                    ...(keepAlive && { keep_alive: keepAlive }),
                    ...(topP && { top_p: topP }),
                }
            }, {
                cancelToken: new axios.CancelToken((c) => {
                    cancellationToken.onCancellationRequested(() => c("Autocompletion request terminated by completion cancel"));
                })
            });
        } catch (error) {
            console.error("Error fetching preview:", error);
        }
    }

    item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Fabelous Autocoder');
    if (continueInline || !responsePreview) {
        item.command = {
            command: 'fabelous-autocoder.autocomplete',
            title: 'Fabelous Autocomplete',
            arguments: [cancellationToken]
        };
    }
    return [item];
}

function activate(context: vscode.ExtensionContext) {
    const completionProvider = vscode.languages.registerCompletionItemProvider("*", {
        provideCompletionItems
    },
        ...completionKeys.split("")
    );
    const externalAutocompleteCommand = vscode.commands.registerTextEditorCommand(
        "fabelous-autocoder.autocomplete",
        (textEditor, _, cancellationToken?) => {
            autocompleteCommand(textEditor, cancellationToken);
        }
    );
    context.subscriptions.push(completionProvider);
    context.subscriptions.push(externalAutocompleteCommand);
}

function deactivate() { }

module.exports = {
    activate,
    deactivate,
};