import { LRUCache } from 'lru-cache';
import vscode, { TextDocumentChangeEvent, Uri } from 'vscode';
import { type Edit } from 'web-tree-sitter';

import { isSupportedLanguage } from 'base/common/languages/languages';
import { ILanguageServiceProvider } from 'base/common/languages/languageService';

import { TreeSitterFile } from '../../code-context/ast/TreeSitterFile';

// TODO Merge to LanguageServiceProvider?
export class TreeSitterFileManager implements vscode.Disposable {
	private documentUpdateListener: vscode.Disposable;
	private didOpenTextDocument: vscode.Disposable;
	private cache: LRUCache<Uri, TreeSitterFile> = new LRUCache({ max: 20 });

	constructor(private lsp: ILanguageServiceProvider) {
		this.documentUpdateListener = vscode.workspace.onDidChangeTextDocument(async event => {
			if (!isSupportedLanguage(event.document.languageId)) {
				return;
			}

			if (event.contentChanges.length <= 0) {
				return;
			}

			await this.updateCacheOnChange(event);
		});

		this.didOpenTextDocument = vscode.workspace.onDidOpenTextDocument(async document => {
			if (!isSupportedLanguage(document.languageId)) {
				return;
			}

			await this.recreate(document);
		});
	}

	// TODO: register all languages
	async init() {}

	private async updateCacheOnChange(event: TextDocumentChangeEvent) {
		const uri = event.document.uri;
		let tsfile = this.getDocument(uri);
		const tree = tsfile?.tree;
		if (!tree) {
			if (!this.cache.has(uri)) {
				const file = await this.create(event.document);
				this.setDocument(uri, file);
			}

			return;
		}

		for (const change of event.contentChanges) {
			const editParams = this.createEditParams(change, event.document);
			tree.edit(editParams);
		}

		tsfile!!.update(tree, event.document.getText());
		this.setDocument(uri, tsfile!!);
	}

	async create(document: vscode.TextDocument): Promise<TreeSitterFile> {
		const cached = this.getDocument(document.uri);
		if (cached) {
			return cached;
		}

		const file = await this.recreate(document);
		return file;
	}

	async recreate(document: vscode.TextDocument) {
		const src = document.getText();
		const langId = document.languageId;

		const file = await TreeSitterFile.create(src, langId, this.lsp, document.uri.fsPath);
		this.setDocument(document.uri, file);
		return file;
	}

	/// todo: this algorithm was generated by the gpt-4o model, it should be reviewed by a developer
	createEditParams(change: vscode.TextDocumentContentChangeEvent, document: vscode.TextDocument): Edit {
		const startPosition = {
			row: change.range.start.line,
			column: change.range.start.character,
		};

		const oldEndPosition = {
			row: change.range.end.line,
			column: change.range.end.character,
		};

		let strings = change.text.split('\n');

		const newEndPosition = {
			row: change.range.start.line + (strings.length - 1),
			column: change.text.endsWith('\n') ? 0 : change.range.start.character + strings.slice(-1)[0].length,
		};

		const startIndex = document.offsetAt(change.range.start);
		const oldEndIndex = document.offsetAt(change.range.end);
		const newEndIndex = startIndex + change.text.length;

		return {
			startIndex: startIndex,
			oldEndIndex: oldEndIndex,
			newEndIndex: newEndIndex,
			startPosition: startPosition,
			oldEndPosition: oldEndPosition,
			newEndPosition: newEndPosition,
		};
	}

	dispose() {
		this.documentUpdateListener?.dispose();
		this.didOpenTextDocument?.dispose();
	}

	public setDocument(uri: Uri, file: TreeSitterFile): void {
		this.cache.set(uri, file);
	}

	/**
	 * If you want to get doc with cache, please use `documentToTreeSitterFile` instead
	 * @param uri
	 */
	public getDocument(uri: Uri): TreeSitterFile | undefined {
		return this.cache.get(uri);
	}
}