import * as vscode from 'vscode';
import * as ts from 'typescript';

import { CONSTANTS } from './constants'; 

export class CodeLensProvider implements vscode.CodeLensProvider {

    public async provideCodeLenses(document: vscode.TextDocument, _token: vscode.CancellationToken): Promise<vscode.CodeLens[]> {

        let functionRanges: { range: vscode.Range, code: string }[] = [];

        // 1. 使用 VS Code 内置的 DocumentSymbolProvider 获取符号
        const symbols = await vscode.commands.executeCommand<vscode.DocumentSymbol[]>(
            'vscode.executeDocumentSymbolProvider', document.uri
        );

        if (symbols && symbols.length > 0) {
            const flattenedSymbols = this.flattenSymbols(symbols);
            const providerFunctions = flattenedSymbols.filter(symbol =>
                symbol.kind === vscode.SymbolKind.Function || symbol.kind === vscode.SymbolKind.Method
            );

            providerFunctions.forEach(symbol => {
                const code = document.getText(symbol.range);
                functionRanges.push({ range: symbol.range, code });
            });
        }

        // 2. 针对 JavaScript/TypeScript，额外尝试 AST 解析
        if (document.languageId === 'javascript' || document.languageId === 'typescript') {
            const astFunctions = this.extractFunctionsFromAST(document);
            astFunctions.forEach(astFunc => {
                if (!this.isOverlapped(astFunc.range, functionRanges)) {
                    functionRanges.push(astFunc);
                }
            });
        }

        // 3. 如果上述方法未获取到函数（或某些语言符号解析失败），使用正则表达式进行兜底匹配
        if (functionRanges.length === 0) {
            for (let i = 0; i < document.lineCount; i++) {
                const line = document.lineAt(i);
                const lineText = line.text.trim();
                if (this.isFunctionDefinition(lineText)) {
                    const range = new vscode.Range(i, 0, i, line.text.length);
                    const functionCode = this.extractFunctionCode(document, i);
                    functionRanges.push({ range, code: functionCode });
                }
            }
        }

        // 根据每个函数范围创建 CodeLens（每个函数显示三个操作）
        const codeLenses: vscode.CodeLens[] = [];
        functionRanges.forEach(item => {
            codeLenses.push(new vscode.CodeLens(item.range, {
                title: CONSTANTS.CODE_INTERPRETATION,
                command: "extension.showExplanation",
                arguments: [item.code]
            }));
            codeLenses.push(new vscode.CodeLens(item.range, {
                title: CONSTANTS.FUNCTION_ANNOTATION,
                command: "extension.showComments",
                arguments: [item.code]
            }));
            codeLenses.push(new vscode.CodeLens(item.range, {
                title: CONSTANTS.TUNING_RECOMMENDATION,
                command: "extension.showOptimization",
                arguments: [item.code]
            }));
        });

        return codeLenses;
    }

    // 递归扁平化 DocumentSymbol 树
    private flattenSymbols(symbols: vscode.DocumentSymbol[]): vscode.DocumentSymbol[] {
        let result: vscode.DocumentSymbol[] = [];
        for (const symbol of symbols) {
            result.push(symbol);
            if (symbol.children && symbol.children.length > 0) {
                result = result.concat(this.flattenSymbols(symbol.children));
            }
        }
        return result;
    }

    // 判断一个范围是否与已有的范围有重叠（简单根据起始和结束行判断）
    private isOverlapped(range: vscode.Range, ranges: { range: vscode.Range, code: string }[]): boolean {
        return ranges.some(item =>
            item.range.start.line === range.start.line && item.range.end.line === range.end.line
        );
    }

    // 针对 JavaScript/TypeScript，使用 TypeScript AST 解析函数、方法、箭头函数等
    private extractFunctionsFromAST(document: vscode.TextDocument): { range: vscode.Range, code: string }[] {
        const results: { range: vscode.Range, code: string }[] = [];
        const sourceCode = document.getText();
        const sourceFile = ts.createSourceFile(document.fileName, sourceCode, ts.ScriptTarget.Latest, true);

        const visit = (node: ts.Node) => {
            // 检查函数声明、方法声明、箭头函数等
            if (ts.isFunctionDeclaration(node) || ts.isMethodDeclaration(node) ||
                ts.isArrowFunction(node) || ts.isFunctionExpression(node)) {

                const start = document.positionAt(node.getStart());
                const end = document.positionAt(node.getEnd());
                const range = new vscode.Range(start, end);
                const code = document.getText(range);

                // 确保我们抓取的是完整的函数定义，而不仅仅是函数体
                results.push({ range, code });
            }
            ts.forEachChild(node, visit);
        };
        visit(sourceFile);
        return results;
    }

    // 兜底正则匹配函数定义，支持多种语言（包括 Java 的泛型、数组等）
    private isFunctionDefinition(lineText: string): boolean {
        // Python
        if (lineText.startsWith("def ")) return true;

        // JavaScript/TypeScript的函数定义（包括箭头函数和方法）
        if (lineText.match(/^\s*(async\s+)?function\s+\w+\s*\(.*\)\s*\{?/)) return true;
        if (lineText.match(/^\s*(const|let|var)\s+\w+\s*=\s*\(?\s*\w*\s*\)?\s*=>\s*\{?/)) return true;
        if (lineText.match(/^\s*\w+\s*\(.*\)\s*\{?$/)) return true;

        // Java：支持 public/private/protected/static, 泛型、数组等
        if (lineText.match(/^\s*(public|private|protected|static|final|synchronized|abstract|native|strictfp)(\s+\S+)*\s+(\w+)\s*\(/m)) return true;

        // C#、C++：可能带返回类型
        if (lineText.match(/^\s*(public|private|protected|static|virtual|override|\w+)\s+\w+\s*\(.*\)\s*\{?$/)) return true;

        // PHP
        if (lineText.match(/^\s*function\s+\w+\s*\(.*\)\s*\{?$/)) return true;

        // Swift
        if (lineText.match(/^\s*func\s+\w+\s*\(.*\)\s*\{?$/)) return true;

        // Go
        if (lineText.match(/^\s*func\s+\w+\s*\(.*\)\s*\{?$/)) return true;

        return false;
    }

    // 简单基于缩进提取函数体（正则兜底方案）
    // 针对所有语言的提取函数，根据语言类型选择合适的提取策略
    private extractFunctionCode(document: vscode.TextDocument, startLine: number): string {
        if (document.languageId === 'python') {
            return this.extractPythonFunctionCode(document, startLine);
        } else {
            // 针对大括号语言（如 Java、JavaScript 等）的提取逻辑
            let code = '';
            let openBraces = 0;
            let started = false;
            for (let i = startLine; i < document.lineCount; i++) {
                const line = document.lineAt(i).text;
                // 如果还没开始并且当前行存在大括号，则标记开始
                if (!started && line.indexOf('{') !== -1) {
                    started = true;
                }
                // 累加当前行代码
                code += line + "\n";
                if (started) {
                    // 统计当前行中 { 和 } 的数量
                    const opens = (line.match(/{/g) || []).length;
                    const closes = (line.match(/}/g) || []).length;
                    openBraces += opens - closes;
                    // 当大括号全部匹配，则认为函数体结束
                    if (openBraces <= 0) {
                        break;
                    }
                } else {
                    // 如果函数体还未开始（即没有 {），则继续读取下一行
                    continue;
                }
            }
            return code;
        }
    }

    // 专门针对 Python 的函数体提取（基于缩进判断）
    private extractPythonFunctionCode(document: vscode.TextDocument, startLine: number): string {
        let code = '';
        // 获取定义行的缩进（空格数）
        const defIndent = document.lineAt(startLine).firstNonWhitespaceCharacterIndex;
        for (let i = startLine; i < document.lineCount; i++) {
            const line = document.lineAt(i).text;
            // 如果不是空行，且缩进小于等于定义行缩进，则认为函数体结束
            if (i > startLine && line.trim() !== '') {
                const currentIndent = document.lineAt(i).firstNonWhitespaceCharacterIndex;
                if (currentIndent <= defIndent) break;
            }
            code += line + "\n";
        }
        return code;
    }
}