import { ExtensionState } from "./config";
import * as vscode from 'vscode';
import { flattenTrie, trieInsert, trieLookup, TrieNode, triePrune } from "./trie";

const getModelSupportsSuffix = async (extension: ExtensionState, model: string) => {
	// TODO: get if model supports suffixes and use that if available

	// const response = await ollama.show({
	// 	model: model
	// })

	// model.capabilities.includes('suffix')
	return false;
};

const getPrompt = (extension: ExtensionState, document: vscode.TextDocument, position: vscode.Position, prefix: string) => {
	const messageHeader = `In an english code base with the file.\nfile:\nproject {PROJECT_NAME}\nfile {FILE_NAME}\nlanguage {LANG}`
		.replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled")
		.replace("{FILE_NAME}", document.fileName)
		.replace("{LANG}", document.languageId);

	const message = `File:\n${extension.configuration.inlineCompletion.prefixStart}`;


	const prompt = `${messageHeader}\n${message}\n${prefix}`;

	return prompt;
};

const getPromptWithSuffix = (extension: ExtensionState, document: vscode.TextDocument, position: vscode.Position, prefix: string) => {
	const suffix = document.getText(new vscode.Range(position.line, position.character, document.lineCount - 1, document.lineAt(document.lineCount - 1).text.length));

	const messageSuffix = `End of the file:\n${extension.configuration.inlineCompletion.suffixStart}\n${suffix}\n${extension.configuration.inlineCompletion.suffixEnd}\n`;
	const messagePrefix = `Start of the file:\n${extension.configuration.inlineCompletion.prefixStart}`;

	const messageHeader = `In an english code base with the file.\nfile:\nproject {PROJECT_NAME}\nfile {FILE_NAME}\nlanguage {LANG}\n.`
		.replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled")
		.replace("{FILE_NAME}", document.fileName)
		.replace("{LANG}", document.languageId);

	const prompt = `${messageHeader}\n${messageSuffix}\n${messagePrefix}\n${prefix}`;

	return prompt;
};

const getSuffix = (extension: ExtensionState, document: vscode.TextDocument, position: vscode.Position) => {
	const suffix = document.getText(new vscode.Range(position.line, position.character, document.lineCount - 1, document.lineAt(document.lineCount - 1).text.length));

	return suffix;
};

let trieRoot: TrieNode = {
	isLeaf: true,
	value: '',
};

const trieRootInsert = (text: string) => {
	trieRoot = trieInsert(trieRoot, text);
};

const trieRootLookup = (text: string) => {
	return trieLookup(trieRoot, text);
};

const trieRootPrune = (text: string) => {
	return triePrune(trieRoot, text);
};

const tokenProvider = async (
	extension: ExtensionState,
	document: vscode.TextDocument,
	position: vscode.Position,
	_context: vscode.InlineCompletionContext,
	token: vscode.CancellationToken,
) => {
	const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character));

	const model = extension.configuration.inlineCompletion.model;

	const modelSupportsSuffix = await getModelSupportsSuffix(extension, model);
	const prompt = modelSupportsSuffix ? getPrompt(extension, document, position, prefix) : getPromptWithSuffix(extension, document, position, prefix);

	const suffix = modelSupportsSuffix ? getSuffix(extension, document, position) : undefined;

	const result = trieRootLookup(prefix);

	if (result !== null) {
		return flattenTrie(result);
	}

	if (token.isCancellationRequested) {
		return [];
	}

	const response = await extension.ollama.generate({
		model,
		prompt,
		suffix,
		raw: true,
		stream: true,
		options: {
			num_predict: extension.configuration.inlineCompletion.maxTokens,
			stop: extension.configuration.inlineCompletion.prefixEnds,
		},
	});

	const pruneTimeout = setTimeout(() => {
		trieRootPrune(prompt);
	}, extension.configuration.inlineCompletion.triePruneTimeout);

	token.onCancellationRequested(() => {
		clearTimeout(pruneTimeout);
		try {
			response.abort();
		} catch { }
	});

	const resultBuffer: string[] = await new Promise(async (resolve, reject) => {
		const buffer: string[] = [];
		const timeout = setTimeout(() => {
			resolve(buffer);
		}, extension.configuration.inlineCompletion.generationTimeout);

		try {
			for await (const part of response) {
				buffer.push(part.response);
				trieRootInsert(prefix + buffer.join(''));
			}
			resolve(buffer);
		} catch (err) {
			reject(err);
		} finally {
			response.abort();
			clearTimeout(timeout);
		};
	});

	return [
		resultBuffer.join('')
	];
};

export const getAutoCompleteProvider = (extension: ExtensionState) => {
	const provider: vscode.InlineCompletionItemProvider = {
		async provideInlineCompletionItems(document, position, context, token) {
			if (!extension.configuration.inlineCompletion.enable) {
				return [];
			}
			try {
				const completions = await tokenProvider(extension, document, position, context, token);

				return completions.map((text) => ({
					insertText: text,
					range: new vscode.Range(position, position),
				}));
			} catch (err) {
				console.log(err);
			}

			return [];
		},
	};

	return provider;
};