"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.isSemanticTokenReference = exports.LSPServer = void 0;
/* --------------------------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See License.txt in the project root for license information.
 * ------------------------------------------------------------------------------------------ */
const antlr4_c3_1 = require("antlr4-c3");
const tree_1 = require("antlr4ts/tree");
const lsp = __importStar(require("vscode-languageserver"));
const flinkdb_schema_json_1 = __importDefault(require("./assets/flinkdb.schema.json"));
const constants_1 = require("./constants");
const FlinkSQLLexer_1 = require("./lib/FlinkSQLLexer");
const FlinkSQLParser_1 = require("./lib/FlinkSQLParser");
const schema_listener_1 = require("./listeners/schema.listener");
const semantic_tokens_listener_1 = require("./listeners/semantic-tokens.listener");
const statement_listener_1 = require("./listeners/statement.listener");
const structure_listener_1 = require("./listeners/structure.listener");
const parsing_warehouse_1 = require("./parsing-warehouse");
const protocol_translation_1 = require("./protocol-translation");
const schema_registry_1 = require("./schema-registry");
const utils_1 = require("./utils");
const format_visitor_1 = require("./visitors/format.visitor");
const preview_visitor_1 = require("./visitors/preview.visitor");
class LSPServer {
    constructor() {
        this._warehouse = new parsing_warehouse_1.ParsingWarehouse();
        this._schemaRegistry = new schema_registry_1.SchemaRegistry();
    }
    get warehouse() {
        return this._warehouse;
    }
    get schemaRegistry() {
        return this._schemaRegistry;
    }
    getTokenByPosition(document, position) {
        const parsingValue = this._warehouse.getParsingValue(document);
        const offset = document.offsetAt(position);
        return parsingValue.tokenStream.getTokens().find(t => t.startIndex <= offset && offset <= t.stopIndex + 1);
    }
    doValidation(document) {
        const parsingValue = this._warehouse.getParsingValue(document);
        return parsingValue.errors;
    }
    doFormatting(document, options) {
        const parsingValue = this._warehouse.getParsingValue(document);
        if (parsingValue.errors.length > 0) {
            return document.getText();
        }
        const visitor = new format_visitor_1.FormatVisitor(parsingValue.tokenStream, options);
        let formatted = visitor.visit(parsingValue.tree);
        // handle unvisited comments at the end of the code
        const unvisitedComments = parsingValue.tokenStream
            .getTokens()
            .filter(t => t.channel === 2 && !visitor.visitedCommentTokens.has(t));
        for (const comment of unvisitedComments) {
            formatted += `\n${comment.text}`;
        }
        return formatted;
    }
    doFoldingRanges(document) {
        const parsingValue = this._warehouse.getParsingValue(document);
        // fold statements for now
        const listener = new statement_listener_1.StatementListener(document);
        tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
        const ranges = listener.getRanges();
        return ranges.map(range => (0, protocol_translation_1.toFoldingRange)(range));
    }
    doReferences(document, position, skipSelf = true) {
        const parsingValue = this._warehouse.getParsingValue(document);
        // find current token based on parameter position
        const tokens = parsingValue.tokenStream.getTokens();
        const currentOffset = document.offsetAt(position);
        const currentToken = tokens.find(token => {
            return token.startIndex <= currentOffset && currentOffset <= token.stopIndex;
        });
        if (currentToken?.text === undefined) {
            return [];
        }
        const listener = new semantic_tokens_listener_1.SemanticTokensListener(document);
        tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
        const semanticTokens = listener.getSemanticTokens();
        const references = [];
        const compareText = (str1, str2) => str1.replace(/['"`]/g, '').localeCompare(str2.replace(/['"`]/g, '')) === 0;
        const findReferences = (token) => {
            if (currentToken.text && compareText(currentToken.text, token.name)) {
                references.push(token.range);
            }
        };
        semanticTokens.forEach(findReferences);
        if (skipSelf) {
            return references.filter(range => !(0, protocol_translation_1.positionInRange)(document, position, range));
        }
        return references;
    }
    parseForSemanticTokens(document) {
        const parsingValue = this._warehouse.getParsingValue(document);
        const listener = new semantic_tokens_listener_1.SemanticTokensListener(document);
        tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
        return listener.getSemanticTokens();
    }
    doCompletion(document, position) {
        const parsingValue = this._warehouse.getParsingValue(document);
        const offset = document.offsetAt(position);
        const token = parsingValue.tokenStream.getTokens().find(t => t.startIndex <= offset && offset <= t.stopIndex + 1);
        if (!token) {
            return { completionItems: [], identifiers: [] };
        }
        const core = new antlr4_c3_1.CodeCompletionCore(parsingValue.parser);
        core.ignoredTokens = new Set(constants_1.ignoredTokens);
        core.preferredRules = new Set(constants_1.preferredRules);
        let index = token.tokenIndex;
        if (token.type === FlinkSQLLexer_1.FlinkSQLLexer.DOT) {
            index += 1;
        }
        const candidates = core.collectCandidates(index);
        const completionItems = [];
        const identifiers = [];
        if (candidates.rules.size > 0) {
            const listener = new schema_listener_1.SchemaListener();
            tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
            const schemaContexts = listener.getSchemaContexts();
            const context = schemaContexts.find(c => c.startOffset <= offset && c.endOffset >= offset);
            // e.g. `catalog`.`database`.`table`
            const schemaParts = context?.text.replace(/[`'"]/g, '').split('.');
            if (candidates.rules.has(FlinkSQLParser_1.FlinkSQLParser.RULE_catalogName)) {
                const catalogsInDocument = (0, utils_1.unique)(schemaContexts.filter(c => c !== context && c.kind === schema_listener_1.SchemaContextKindEnum.CATALOG).map(c => c.catalogName));
                completionItems.push(...catalogsInDocument.map(c => (0, protocol_translation_1.toCompletionItem)('catalog', c, `\`${c}\``, 'Catalog in document')), ...this._schemaRegistry
                    .getSchemaCatalogs()
                    .map(c => (0, protocol_translation_1.toCompletionItem)('catalog', c.label, `\`${c.label}\``, 'Schema catalog', c.description)));
                identifiers.push('catalog');
            }
            if (candidates.rules.has(FlinkSQLParser_1.FlinkSQLParser.RULE_databaseName)) {
                if (!schemaParts || schemaParts.length <= 1) {
                    const databasesInDocument = (0, utils_1.unique)(schemaContexts.filter(d => d !== context && d.kind === schema_listener_1.SchemaContextKindEnum.DATABASE).map(d => d.databaseName));
                    completionItems.push(...databasesInDocument.map(d => (0, protocol_translation_1.toCompletionItem)('database', d, `\`${d}\``, 'Database in document')), ...this._schemaRegistry
                        .getSchemaDatabases({})
                        .map(d => (0, protocol_translation_1.toCompletionItem)('database', d.label, `\`${d.catalog}\`.\`${d.label}\``, `Schema database: ${d.catalog}.*`, d.description)));
                }
                else {
                    completionItems.push(...this._schemaRegistry
                        .getSchemaDatabases({ catalog: schemaParts[0] })
                        .map(d => (0, protocol_translation_1.toCompletionItem)('database', d.label, `\`${d.label}\``, `Schema database: ${d.catalog}.*`, d.description)));
                }
                identifiers.push('database');
            }
            if (candidates.rules.has(FlinkSQLParser_1.FlinkSQLParser.RULE_viewName)) {
                if (!schemaParts || schemaParts.length <= 1) {
                    const viewsInDocument = (0, utils_1.unique)(schemaContexts.filter(v => v !== context && v.kind === schema_listener_1.SchemaContextKindEnum.VIEW).map(v => v.viewName));
                    completionItems.push(...viewsInDocument.map(v => (0, protocol_translation_1.toCompletionItem)('view', v, v, 'View in document')), ...this._schemaRegistry
                        .getSchemaViews({})
                        .map(v => (0, protocol_translation_1.toCompletionItem)('view', v.label, `\`${v.catalog}\`.\`${v.database}\`.${v.label}`, `Schema view: ${v.catalog}.${v.database}.*`, v.description)));
                }
                else if (schemaParts.length === 2) {
                    completionItems.push(...this._schemaRegistry
                        .getSchemaViews({ database: schemaParts[0] })
                        .map(v => (0, protocol_translation_1.toCompletionItem)('view', v.label, v.label, `Schema view: ${v.catalog}.${v.database}.*`, v.description)));
                }
                else {
                    completionItems.push(...this._schemaRegistry
                        .getSchemaViews({ catalog: schemaParts[0], database: schemaParts[1] })
                        .map(v => (0, protocol_translation_1.toCompletionItem)('view', v.label, v.label, `Schema view: ${v.catalog}.${v.database}.*`, v.description)));
                }
                identifiers.push('view');
            }
            if (candidates.rules.has(FlinkSQLParser_1.FlinkSQLParser.RULE_tableName)) {
                if (!schemaParts || schemaParts.length <= 1) {
                    const tablesInDocument = (0, utils_1.unique)(schemaContexts.filter(t => t !== context && t.kind === schema_listener_1.SchemaContextKindEnum.TABLE).map(t => t.tableName));
                    completionItems.push(...tablesInDocument.map(t => (0, protocol_translation_1.toCompletionItem)('table', t, t, 'Table in document')), ...this._schemaRegistry
                        .getSchemaTables({})
                        .map(t => (0, protocol_translation_1.toCompletionItem)('table', t.label, `\`${t.catalog}\`.\`${t.database}\`.${t.label}`, `Schema table: ${t.catalog}.${t.database}.*`, t.description)));
                }
                else if (schemaParts.length === 2) {
                    completionItems.push(...this._schemaRegistry
                        .getSchemaTables({ database: schemaParts[0] })
                        .map(t => (0, protocol_translation_1.toCompletionItem)('table', t.label, t.label, `Schema table: ${t.catalog}.${t.database}.*`, t.description)));
                }
                else {
                    completionItems.push(...this._schemaRegistry
                        .getSchemaTables({ catalog: schemaParts[0], database: schemaParts[1] })
                        .map(t => (0, protocol_translation_1.toCompletionItem)('table', t.label, t.label, `Schema table: ${t.catalog}.${t.database}.*`, t.description)));
                }
                identifiers.push('table');
            }
            if (candidates.rules.has(FlinkSQLParser_1.FlinkSQLParser.RULE_columnName)) {
                let columns;
                if (!schemaParts || schemaParts.length <= 1) {
                    columns = this._schemaRegistry.getSchemaColumns({});
                }
                else if (schemaParts.length === 2) {
                    columns = this._schemaRegistry.getSchemaColumns({ table: schemaParts[0] });
                }
                else if (schemaParts.length === 3) {
                    columns = this._schemaRegistry.getSchemaColumns({ database: schemaParts[0], table: schemaParts[1] });
                }
                else {
                    columns = this._schemaRegistry.getSchemaColumns({
                        catalog: schemaParts[0],
                        database: schemaParts[1],
                        table: schemaParts[2]
                    });
                }
                if (columns.length) {
                    completionItems.push(...columns.map(c => (0, protocol_translation_1.toCompletionItem)('column', c.label, c.label, `Schema column: ${c.catalog}.${c.database}.${c.table}.*`, c.description)));
                }
                identifiers.push('column');
            }
            if (candidates.rules.has(FlinkSQLParser_1.FlinkSQLParser.RULE_functionName)) {
                completionItems.push(...flinkdb_schema_json_1.default.functions.map(f => (0, protocol_translation_1.toCompletionItem)('function', f.name, f.name, 'Built-in Function', f.description)));
                identifiers.push('function');
            }
        }
        for (const candidate of Array.from(candidates.tokens)) {
            const keyword = parsingValue.parser.vocabulary.getDisplayName(candidate[0]).replace(/'/g, '');
            completionItems.push({
                kind: lsp.CompletionItemKind.Keyword,
                label: keyword,
                insertText: keyword,
                insertTextFormat: lsp.InsertTextFormat.Snippet,
                sortText: `F${keyword}`
            });
        }
        return { completionItems, identifiers };
    }
    parseStructure(document, options = { ignoreInsertSource: false, addCreateContext: false }) {
        const parsingValue = this._warehouse.getParsingValue(document);
        const listener = new structure_listener_1.FlinkStructureListener(document);
        tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
        let nodes = listener.getNodes();
        let edges = listener.getEdges();
        if (options.ignoreInsertSource) {
            const tmp = (0, structure_listener_1.ignoreInsertSource)(nodes, edges);
            nodes = tmp.nodes;
            edges = tmp.edges;
        }
        if (options.addCreateContext) {
            const tmp = (0, structure_listener_1.addFromCreateContext)(nodes, edges);
            nodes = tmp.nodes;
            edges = tmp.edges;
        }
        return { nodes, edges };
    }
    generatePreviewScript(document, context) {
        const parsingValue = this._warehouse.getParsingValue(document);
        if (parsingValue.errors.length) {
            console.error(`Document '${document.uri}' conatins syntax errors, preview script failed to generate.`);
            return document.getText();
        }
        const visitor = new preview_visitor_1.PreviewVisitor(parsingValue.tokenStream, context);
        const script = visitor.visit(parsingValue.tree);
        return (0, preview_visitor_1.appendCreateNonTemporaryTables)(script, context);
    }
    extractSchemaContexts(document) {
        const parsingValue = this._warehouse.getParsingValue(document);
        const listener = new schema_listener_1.SchemaListener();
        try {
            tree_1.ParseTreeWalker.DEFAULT.walk(listener, parsingValue.tree);
            return listener.getSchemaContexts();
        }
        catch (e) {
            return [];
        }
    }
}
exports.LSPServer = LSPServer;
function isSemanticTokenReference(token1, token2) {
    return (token1.name.replace(/['"`]/g, '').localeCompare(token2.name.replace(/['"`]/g, '')) === 0 &&
        token1.type === token2.type);
}
exports.isSemanticTokenReference = isSemanticTokenReference;
//# sourceMappingURL=lsp-server.js.map