import * as vscode from 'vscode';

// 文档格式化提供者类
export class DocumentFormattingProvider implements vscode.DocumentFormattingEditProvider {
    // 提供整个文档的代码格式化功能
    async provideDocumentFormattingEdits(
        document: vscode.TextDocument,
        options: vscode.FormattingOptions,
        token: vscode.CancellationToken
    ): Promise<vscode.TextEdit[]> {
        const edits: vscode.TextEdit[] = [];
        const lines: string[] = [];

        // 逐行处理文档内容
        for (let i = 0; i < document.lineCount; i++) {
            const line = document.lineAt(i);
            lines.push(line.text);
        }

        // 格式化文档
        const formattedLines = this.formatDocument(lines, options);

        // 创建编辑操作
        const fullRange = new vscode.Range(
            document.positionAt(0),
            document.positionAt(document.getText().length)
        );

        edits.push(vscode.TextEdit.replace(fullRange, formattedLines.join('\n')));

        return edits;
    }

    // 格式化文档
    private formatDocument(lines: string[], options: vscode.FormattingOptions): string[] {
        const formattedLines: string[] = [];
        let indentLevel = 0;

        for (let i = 0; i < lines.length; i++) {
            let line = lines[i];
            const trimmedLine = line.trim();

            if (trimmedLine.length === 0) {
                // 空行保持原样或删除
                formattedLines.push('');
                continue;
            }

            // 检查是否需要减少缩进（遇到右大括号等）
            if (trimmedLine.startsWith('}') || trimmedLine.startsWith(')') || trimmedLine.startsWith(']')) {
                indentLevel = Math.max(0, indentLevel - 1);
            }

            // 添加适当的缩进
            const indent = this.createIndent(indentLevel, options);
            const formattedLine = indent + trimmedLine;
            formattedLines.push(formattedLine);

            // 检查是否需要增加缩进（遇到左大括号等）
            if (trimmedLine.endsWith('{') || trimmedLine.endsWith('(') || trimmedLine.endsWith('[')) {
                indentLevel++;
            }

            // 特殊处理else、catch等关键字，它们应该与前面的右大括号在同一行
            if (trimmedLine.startsWith('else') || trimmedLine.startsWith('catch') || trimmedLine.startsWith('elif')) {
                if (formattedLines.length > 1) {
                    const prevLine = formattedLines[formattedLines.length - 2];
                    if (prevLine.trim().endsWith('}')) {
                        // 将else等关键字与前面的右大括号放在同一行
                        formattedLines[formattedLines.length - 2] = prevLine.trim() + ' ' + trimmedLine;
                        formattedLines.pop(); // 移除当前行
                        indentLevel--; // 修正缩进级别
                    }
                }
            }
        }

        // 规范化空格（在操作符周围添加适当空格）
        return formattedLines.map(line => this.normalizeSpaces(line));
    }

    // 在操作符周围添加适当空格
    private normalizeSpaces(line: string): string {
        // 在常见的操作符周围添加空格
        return line
            .replace(/([^\s])(==|!=|<=|>=|<|>|=|\+|-|\*|\/|%)([^\s])/g, '$1 $2 $3')
            .replace(/([^\s])(,)([^\s])/g, '$1$2 $3')
            .replace(/\s+/g, ' ') // 将多个空格替换为单个空格
            .trim();
    }

    // 创建指定级别的缩进
    private createIndent(level: number, options: vscode.FormattingOptions): string {
        if (options.insertSpaces) {
            return ' '.repeat(level * options.tabSize);
        } else {
            return '\t'.repeat(level);
        }
    }
}