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

let VSConfig: vscode.WorkspaceConfiguration;
let apiEndpoint: string;
let apiAuthentication: string;
let apiModel: string;
let apiMessageHeader: 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"; // Updated to use FIM model
    apiMessageHeader = VSConfig.get("message header") || "";
    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 messageHeaderSub(document: vscode.TextDocument) {
    const sub = apiMessageHeader
        .replace("{LANG}", document.languageId)
        .replace("{FILE_NAME}", document.fileName)
        .replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled");
    return sub;
}

function getContextLines(document: vscode.TextDocument, position: vscode.Position): string {
    const lines = [];
    const lineCount = document.lineCount;
    
    // Get more context for FIM
    const startLine = Math.max(0, position.line - 10);
    const endLine = Math.min(lineCount - 1, position.line + 10);
    
    for (let i = startLine; i <= endLine; i++) {
        lines.push(document.lineAt(i).text);
    }
    
    return lines.join("\n");
}

async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) {
    const document = textEditor.document;
    const position = textEditor.selection.active;

    // Get the current context
    const context = getContextLines(document, position);

    // Split the context into prefix and suffix for FIM
    const lines = context.split("\n");
    const currentLineIndex = position.line - Math.max(0, position.line - 10);
    const prefix = lines.slice(0, currentLineIndex + 1).join("\n");
    const suffix = lines.slice(currentLineIndex + 1).join("\n");

    // Create FIM prompt
    const fimPrompt = `<fim_prefix>${prefix}<fim_suffix>${suffix}<fim_middle>`;

    // Replace {Prompt} with the FIM prompt
    const sub = messageHeaderSub(document).replace("{PROMPT}", fimPrompt);

    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) => {
                    const cancelPost = function () {
                        c("Autocompletion request terminated by user cancel");
                    };

                    axiosCancelPost = cancelPost;
                    if (cancellationToken) cancellationToken.onCancellationRequested(cancelPost);
                    progressCancellationToken.onCancellationRequested(cancelPost);
                    vscode.workspace.onDidCloseTextDocument(cancelPost);
                });

                const response = await axios.post(apiEndpoint, {
                    model: apiModel,
                    prompt: sub,
                    stream: true,
                    raw: true,
                    options: {
                        num_predict: numPredict,
                        temperature: apiTemperature,
                        stop: ["<fim_suffix>", "```"]
                    }
                }, {
                    cancelToken: axiosCancelToken,
                    responseType: 'stream',
                    headers: {
                        'Authorization': apiAuthentication
                    }
                });

                let currentPosition = position;
                let completionText = "";
                response.data.on('data', async (d: Uint8Array) => {
                    progress.report({ message: "Generating..." });
                    if (currentPosition.line != textEditor.selection.end.line || currentPosition.character != textEditor.selection.end.character) {
                        axiosCancelPost();
                        return;
                    }
                    const completion: string = JSON.parse(d.toString()).response;
                    
                    if (completion === "") {
                        return;
                    }

                    completionText += completion;
                    
                    const edit = new vscode.WorkspaceEdit();
                    edit.insert(document.uri, currentPosition, completion);
                    await vscode.workspace.applyEdit(edit);

                    const completionLines = completion.split("\n");
                    const newPosition = new vscode.Position(
                        currentPosition.line + completionLines.length - 1,
                        (completionLines.length > 1 ? 0 : currentPosition.character) + completionLines[completionLines.length - 1].length
                    );
                    const newSelection = new vscode.Selection(
                        position,
                        newPosition
                    );
                    currentPosition = newPosition;

                    progress.report({ message: "Generating...", increment: 1 / (numPredict / 100) });
                    textEditor.selection = newSelection;
                });

                const finished = new Promise((resolve) => {
                    response.data.on('end', () => {
                        progress.report({ message: "Fabelous completion finished." });
                        resolve(true);
                    });
                    axiosCancelToken.promise.finally(() => {
                        resolve(false);
                    });
                });
                await finished;

                // Remove any remaining FIM tokens from the completion
                completionText = completionText.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '');
                const finalEdit = new vscode.WorkspaceEdit();
                finalEdit.replace(document.uri, new vscode.Range(position, currentPosition), completionText);
                await vscode.workspace.applyEdit(finalEdit);

            } 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 ];
    }

    if (responsePreview) {
        const context = getContextLines(document, position);
        const lines = context.split("\n");
        const currentLineIndex = position.line - Math.max(0, position.line - 10);
        const prefix = lines.slice(0, currentLineIndex + 1).join("\n");
        const suffix = lines.slice(currentLineIndex + 1).join("\n");
        const fimPrompt = `<fim_prefix>${prefix}<fim_suffix>${suffix}<fim_middle>`;

        const response_preview = await axios.post(apiEndpoint, {
            model: apiModel,
            prompt: messageHeaderSub(document) + 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) => {
                const cancelPost = function () {
                    c("Autocompletion request terminated by completion cancel");
                };
                cancellationToken.onCancellationRequested(cancelPost);
            })
        });
        if (response_preview.data.response.trim() != "") {
            const previewText = response_preview.data.response.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trimStart();
            item.label = previewText;
            item.insertText = previewText;
        }
    }

    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,
};