import { injectable } from 'inversify';
import Parser, { Language } from 'web-tree-sitter';

import { BaseStructurerProvider } from "../base/StructurerProvider";
import { CodeFile, CodeStructure, StructureType } from "../../codemodel/CodeElement";
import { LanguageProfileUtil } from "../base/LanguageProfileUtil";
import { LanguageProfile } from "../base/LanguageProfile";

@injectable()
export class PythonStructurer extends BaseStructurerProvider {
	protected langId: string = 'python';
	protected config: LanguageProfile = LanguageProfileUtil.from(this.langId)!!;
	protected parser: Parser | undefined;
	protected language: Language | undefined;

	isApplicable(lang: string) {
		return lang === 'python';
	}

	parseFile(code: string, filepath: string): Promise<CodeFile | undefined> {
		const tree = this.parser!!.parse(code);
		const query = this.config.structureQuery.query(this.language!!);
		const captures = query!!.captures(tree.rootNode);

		let filename = filepath.split('/')[filepath.split('/').length - 1];
		const codeFile: CodeFile = {
			name: filename,
			filepath: filepath,
			language: this.langId,
			functions: [],
			package: '',
			imports: [],
			classes: [],
		};

		let classObj: CodeStructure = this.createEmptyClassStructure();

		// 存储已处理的函数节点，避免重复添加
		const processedFunctionNodes = new Set<Parser.SyntaxNode>();
		// 存储类方法节点，用于排除顶级函数
		const classMethodNodes = new Set<Parser.SyntaxNode>();

		// 第一遍：识别所有类和类方法
		for (const element of captures) {
			const capture: Parser.QueryCapture = element!!;

			switch (capture.name) {
				case 'class-name':
					// 标记类节点存在
					const classNode = capture.node?.parent;
					if (classNode) {
						// 处理类体中的内容
						const bodyNode = this.findChildNode(classNode, 'block');
						if (bodyNode) {
							// 将所有类中的函数定义节点添加到集合中
							this.collectFunctionNodesInBlock(bodyNode, classMethodNodes);
						}
					}
					break;
				case 'class-method-name':
					// 将类方法节点添加到集合
					const methodNode = capture.node?.parent;
					if (methodNode) {
						classMethodNodes.add(methodNode);
					}
					break;
			}
		}

		// 第二遍：正式处理所有节点
		for (const element of captures) {
			const capture: Parser.QueryCapture = element!!;
			const text = capture.node.text;

			switch (capture.name) {
				case 'import-name':
				case 'module-name':
					codeFile.imports.push(text);
					break;
				case 'class-name':
					// 创建新的类对象
					classObj = this.createEmptyClassStructure();
					classObj.name = text;
					classObj.type = StructureType.Class;
					const classNode: Parser.SyntaxNode | null = capture.node?.parent ?? null;
					if (classNode) {
						classObj.start = { row: classNode.startPosition.row, column: classNode.startPosition.column };
						classObj.end = { row: classNode.endPosition.row, column: classNode.endPosition.column };
						codeFile.classes.push(classObj);
					}
					break;
				case 'superclass-name':
					// 将父类名添加到当前类的extends数组中
					if (codeFile.classes.length > 0) {
						let currentClass = codeFile.classes[codeFile.classes.length - 1];
						currentClass.extends.push(text);
					}
					break;
				case 'class-method-name':
					// 将方法添加到当前类
					if (codeFile.classes.length > 0) {
						const methodNode = capture.node?.parent;
						if (methodNode && !processedFunctionNodes.has(methodNode)) {
							const methodFunc = this.createFunction(capture.node, text);
							codeFile.classes[codeFile.classes.length - 1].methods.push(methodFunc);
							processedFunctionNodes.add(methodNode);
						}
					}
					break;
				case 'class-attribute-name':
					// 将属性添加到当前类
					if (codeFile.classes.length > 0) {
						let lastClass = codeFile.classes[codeFile.classes.length - 1];
						lastClass.fields = lastClass.fields ?? [];
						lastClass.fields.push(this.createVariable(capture.node, text, ''));
					}
					break;
				// 使用明确标记的顶层函数
				case 'toplevel-function-name':
					const toplevelFunctionNode = capture.node?.parent;
					if (toplevelFunctionNode && !processedFunctionNodes.has(toplevelFunctionNode)) {
						codeFile.functions.push(this.createFunction(capture.node, text));
						processedFunctionNodes.add(toplevelFunctionNode);
					}
					break;
				// 保留原来的函数处理逻辑但不添加到顶层函数
				case 'function-name':
					// 不再将这类函数添加到顶层函数中，除非确认它不是类方法
					break;
			}
		}

		// 合并重复的类
		this.mergeClasses(codeFile);

		return Promise.resolve(codeFile);
	}

	// 查找指定类型的子节点
	private findChildNode(node: Parser.SyntaxNode, type: string): Parser.SyntaxNode | null {
		for (let i = 0; i < node.childCount; i++) {
			const child = node.child(i);
			if (child && child.type === type) {
				return child;
			}
		}
		return null;
	}

	// 收集块内的所有函数定义节点
	private collectFunctionNodesInBlock(blockNode: Parser.SyntaxNode, collection: Set<Parser.SyntaxNode>): void {
		for (let i = 0; i < blockNode.childCount; i++) {
			const child = blockNode.child(i);
			if (child && child.type === 'function_definition') {
				collection.add(child);
			} else if (child && child.childCount > 0) {
				// 递归处理嵌套块
				if (child.type === 'block') {
					this.collectFunctionNodesInBlock(child, collection);
				}
			}
		}
	}

	// 创建空的类结构对象
	private createEmptyClassStructure(): CodeStructure {
		return {
			type: StructureType.Class,
			canonicalName: '',
			constant: [],
			extends: [],
			methods: [],
			name: '',
			package: '',
			implements: [],
			start: { row: 0, column: 0 },
			end: { row: 0, column: 0 },
		};
	}

	// 合并重复的类
	private mergeClasses(codeFile: CodeFile): void {
		const uniqueClasses: CodeStructure[] = [];
		const classMap = new Map<string, CodeStructure>();

		for (const classObj of codeFile.classes) {
			const className = classObj.name;

			if (classMap.has(className)) {
				// 合并属性到现有类
				const existingClass = classMap.get(className)!;

				// 合并方法
				if (classObj.methods && classObj.methods.length > 0) {
					existingClass.methods = [...(existingClass.methods || []), ...classObj.methods];
				}

				// 合并字段
				if (classObj.fields && classObj.fields.length > 0) {
					existingClass.fields = [...(existingClass.fields || []), ...classObj.fields];
				}

				// 合并常量
				if (classObj.constant && classObj.constant.length > 0) {
					existingClass.constant = [...(existingClass.constant || []), ...classObj.constant];
				}

				// 更新其他可能需要合并的属性
				if (classObj.extends && classObj.extends.length > 0) {
					existingClass.extends = [...(existingClass.extends || []), ...classObj.extends];
				}
			} else {
				// 添加新类到映射表
				classMap.set(className, classObj);
				uniqueClasses.push(classObj);
			}
		}

		// 更新代码文件中的类数组
		codeFile.classes = uniqueClasses;
	}
}