import * as vscode from 'vscode';

export class TplDocumentFormatter implements vscode.DocumentFormattingEditProvider {
    provideDocumentFormattingEdits(document: vscode.TextDocument, options: vscode.FormattingOptions, token: vscode.CancellationToken): vscode.ProviderResult<vscode.TextEdit[]> {
        const text = document.getText();
        const isCRLF = text.includes('\r\n');
        const lines = text.split(/\r?\n/);
        
        const formattedLines: string[] = [];
        const edits: vscode.TextEdit[] = [];
        
        let indentLevel = 0;
        const indentString = '    ';
        let inTripleDouble = false;
        let inTripleSingle = false;
        let consecutiveEmptyLines = 0;
        
        const alignmentStack: number[] = [];

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

            const tripleDoubleCount = (rawLine.match(/"""/g) || []).length;
            const tripleSingleCount = (rawLine.match(/'''/g) || []).length;
            
            let isInternalStringLine = false;

            if (inTripleDouble) {
                isInternalStringLine = true;
                if (tripleDoubleCount % 2 !== 0) { inTripleDouble = false; isInternalStringLine = false; }
            } else if (tripleDoubleCount % 2 !== 0) { inTripleDouble = true; }

            if (inTripleSingle) {
                isInternalStringLine = true;
                if (tripleSingleCount % 2 !== 0) { inTripleSingle = false; isInternalStringLine = false; }
            } else if (tripleSingleCount % 2 !== 0) { inTripleSingle = true; }

            if (isInternalStringLine && (inTripleDouble || inTripleSingle)) {
                formattedLines.push(rawLine);
                consecutiveEmptyLines = 0;
                continue; 
            }

            if (line.length === 0) {
                consecutiveEmptyLines++;
                if (consecutiveEmptyLines <= 1) { formattedLines.push(""); }
                continue; 
            }

            consecutiveEmptyLines = 0;

            if (line.match(/^(end\s+(pattern|syncmapping|table|enumeration|configuration|definitions?|define|if|for|while|metadata|overview|constants|triggers|mapping|removal|body)|else|elif)\b/)) {
                indentLevel = Math.max(0, indentLevel - 1);
            }

            let newLineText = "";
            if (alignmentStack.length > 0) {
                let spaces = alignmentStack[alignmentStack.length - 1];
                if (line.match(/^[)\]]/)) { spaces = Math.max(0, spaces - 1); }
                newLineText = " ".repeat(spaces) + line;
            } else {
                newLineText = indentString.repeat(indentLevel) + line;
            }

            formattedLines.push(newLineText);

            let cleanLineForParens = newLineText.replace(/"(?:\\.|[^"\\])*"|'(?:\\.|[^'\\])*'/g, match => " ".repeat(match.length));
            const commentIdx = cleanLineForParens.indexOf('//');
            if (commentIdx !== -1) {
                cleanLineForParens = cleanLineForParens.substring(0, commentIdx) + " ".repeat(cleanLineForParens.length - commentIdx);
            }

            for (let charIdx = 0; charIdx < cleanLineForParens.length; charIdx++) {
                const char = cleanLineForParens[charIdx];
                if (char === '(' || char === '[') {
                    alignmentStack.push(charIdx + 1); 
                } else if (char === ')' || char === ']') {
                    if (alignmentStack.length > 0) { alignmentStack.pop(); }
                }
            }

            if (line.match(/^(pattern|syncmapping|table|enumeration|configuration|definitions?|define|metadata|overview|constants|triggers|mapping|removal|body|if|for|while|else|elif)\b/)) {
                if (!line.startsWith('end ')) {
                    indentLevel++;
                }
            }
        }

        const newText = formattedLines.join(isCRLF ? '\r\n' : '\n');
        
        if (newText !== text) {
            const fullRange = new vscode.Range(document.positionAt(0), document.positionAt(text.length));
            edits.push(vscode.TextEdit.replace(fullRange, newText));
        }

        return edits;
    }
}