import Parser, { SyntaxNode } from 'web-tree-sitter';

import { countTokens } from '../../token/TokenCounter';
import { ChunkWithoutID } from './Chunk';

/**
 * The `CollapsedCodeChunker` class provides methods to parse and chunk code into smaller pieces. It also provides
 * methods to collapse certain parts of the code, such as statement blocks and other syntax nodes, to make the
 * code more readable and manageable.
 */
export class CollapsedCodeChunker {
	async *parsedCodeChunker(
		parser: Parser,
		contents: string,
		maxChunkSize: number,
		language: string,
	): AsyncGenerator<ChunkWithoutID> {
		const tree = parser.parse(contents);
		yield* this.getSmartCollapsedChunks(tree.rootNode, contents, maxChunkSize, language);
	}

	collapsedReplacement(node: SyntaxNode): string {
		if (node.type === 'statement_block') {
			return '{ ... }';
		} else {
			return '...';
		}
	}

	firstChild(node: SyntaxNode, grammarName: string | string[]): SyntaxNode | null {
		if (Array.isArray(grammarName)) {
			return node.children.find(child => grammarName.includes(child.type)) || null;
		} else {
			return node.children.find(child => child.type === grammarName) || null;
		}
	}

	collapseChildren(
		node: SyntaxNode,
		code: string,
		blockTypes: string[],
		collapseTypes: string[],
		collapseBlockTypes: string[],
		maxChunkSize: number,
	): string {
		code = code.slice(0, node.endIndex);
		const block = this.firstChild(node, blockTypes);
		const collapsedChildren = [];

		if (block) {
			const childrenToCollapse = block.children.filter(child => collapseTypes.includes(child.type));
			for (const child of childrenToCollapse.reverse()) {
				const grandChild = this.firstChild(child, collapseBlockTypes);
				if (grandChild) {
					const start = grandChild.startIndex;
					const end = grandChild.endIndex;
					const collapsedChild = code.slice(child.startIndex, start) + this.collapsedReplacement(grandChild);
					code = code.slice(0, start) + this.collapsedReplacement(grandChild) + code.slice(end);

					collapsedChildren.unshift(collapsedChild);
				}
			}
		}
		code = code.slice(node.startIndex);
		let removedChild = false;
		while (countTokens(code) > maxChunkSize && collapsedChildren.length > 0) {
			removedChild = true;
			// Remove children starting at the end - TODO: Add multiple chunks so no children are missing
			const childCode = collapsedChildren.pop()!;
			const index = code.lastIndexOf(childCode);
			if (index > 0) {
				code = code.slice(0, index) + code.slice(index + childCode.length);
			}
		}

		if (removedChild) {
			// Remove the extra blank lines
			let lines = code.split('\n');
			let firstWhiteSpaceInGroup = -1;
			for (let i = lines.length - 1; i >= 0; i--) {
				if (lines[i].trim() === '') {
					if (firstWhiteSpaceInGroup < 0) {
						firstWhiteSpaceInGroup = i;
					}
				} else {
					if (firstWhiteSpaceInGroup - i > 1) {
						// Remove the lines
						lines = [...lines.slice(0, i + 1), ...lines.slice(firstWhiteSpaceInGroup + 1)];
					}
					firstWhiteSpaceInGroup = -1;
				}
			}

			code = lines.join('\n');
		}

		return code;
	}

	// TODO: chunk in different languages
	constructClassDefinitionChunk(node: SyntaxNode, code: string, maxChunkSize: number): string {
		return this.collapseChildren(
			node,
			code,
			['block', 'class_body', 'declaration_list'],
			['method_definition', 'function_definition', 'function_item'],
			['block', 'statement_block'],
			maxChunkSize,
		);
	}

	// TODO: chunk in different languages
	constructFunctionDefinitionChunk(node: SyntaxNode, code: string, maxChunkSize: number): string {
		const bodyNode = node.children[node.children.length - 1];
		const funcText = code.slice(node.startIndex, bodyNode.startIndex) + this.collapsedReplacement(bodyNode);

		if (
			node.parent &&
			['block', 'declaration_list'].includes(node.parent.type) &&
			node.parent.parent &&
			['class_definition', 'impl_item'].includes(node.parent.parent.type)
		) {
			// If inside a class, include the class header
			const classNode = node.parent.parent;
			const classBlock = node.parent;
			return (
				code.slice(classNode.startIndex, classBlock.startIndex) +
				'...\n\n' +
				' '.repeat(node.startPosition.column) + // ...
				funcText
			);
		}
		return funcText;
	}

	collapsedNodeConstructors: {
		[key: string]: (node: SyntaxNode, code: string, maxChunkSize: number) => string;
	} = {
		// Classes, structs, etc
		class_definition: this.constructClassDefinitionChunk.bind(this),
		class_declaration: this.constructClassDefinitionChunk.bind(this),
		impl_item: this.constructClassDefinitionChunk.bind(this),
		// Functions
		function_definition: this.constructFunctionDefinitionChunk.bind(this),
		function_declaration: this.constructFunctionDefinitionChunk.bind(this),
		function_item: this.constructFunctionDefinitionChunk.bind(this),
	};

	*getSmartCollapsedChunks(
		node: SyntaxNode,
		code: string,
		maxChunkSize: number,
		language: string,
		root = true,
	): Generator<ChunkWithoutID> {
		// Keep entire text if not over size
		if ((root || node.type in this.collapsedNodeConstructors) && countTokens(node.text) < maxChunkSize) {
			yield {
				content: node.text,
				startLine: node.startPosition.row,
				endLine: node.endPosition.row,
				language,
			};

			return;
		}

		// If a collapsed form is defined, use that
		if (node.type in this.collapsedNodeConstructors) {
			yield {
				content: this.collapsedNodeConstructors[node.type](node, code, maxChunkSize),
				startLine: node.startPosition.row,
				endLine: node.endPosition.row,
				language,
			};
		}

		// Recurse (because even if collapsed version was shown, want to show the children in full somewhere)
		for (const child of node.children) {
			yield* this.getSmartCollapsedChunks(child, code, maxChunkSize, language, false);
		}
	}
}