"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.CompletionVisitor = void 0;
const cursor_1 = require("../cursor");
const FlinkSQLParser_1 = require("../lib/FlinkSQLParser");
const utils_1 = require("../utils");
const flinksql_relation_visitor_1 = require("./flinksql-relation.visitor");
const getRelationName = (relationPrimaryString) => {
    const splits = relationPrimaryString.split(flinksql_relation_visitor_1.RELATION_SEPARATOR);
    return splits[splits.length - 1];
};
function availableColumns(relation) {
    const columns = [];
    if (relation instanceof flinksql_relation_visitor_1.QueryRelation) {
        relation.relations.forEach((rel, name) => {
            const relationName = name !== rel.id ? name : undefined;
            rel.columns.forEach(col => {
                if (col.label !== flinksql_relation_visitor_1.STAR_LABEL && !col.label.includes(cursor_1.defaultCursorPlaceholder)) {
                    columns.push({
                        relation: relationName !== undefined ? { category: 'table', label: relationName } : undefined,
                        name: col.label,
                        desc: col.data !== undefined ? col.data : undefined
                    });
                }
            });
        });
    }
    else {
        if (relation.sinkTableRelation !== undefined) {
            relation.sinkTableRelation.columns.forEach(col => {
                if (col.label !== flinksql_relation_visitor_1.STAR_LABEL && !col.label.includes(cursor_1.defaultCursorPlaceholder)) {
                    columns.push({
                        relation: { category: 'table', label: relation.sinkTableRelation.tablePrimary.relationName },
                        name: col.label,
                        desc: col.data !== undefined ? col.data : undefined
                    });
                }
            });
        }
    }
    return columns;
}
function getTableMetadataFnFactory(getColumns) {
    return (relationPrimary) => {
        if (relationPrimary.type !== 'table') {
            return undefined;
        }
        const columns = getColumns({
            table: relationPrimary.relationName,
            catalog: relationPrimary.catalogName,
            database: relationPrimary.databaseName
        });
        return {
            columns: columns.map(c => {
                return {
                    id: `${(0, flinksql_relation_visitor_1.relationPrimaryToString)(relationPrimary)}_${c.label}`,
                    label: c.label,
                    data: c
                };
            })
        };
    };
}
const SQL_SNIPPETS = [
    {
        label: 'BEGIN STATEMENT SET',
        filterText: 'BEGIN',
        template: `BEGIN STATEMENT SET
;

$0`
    },
    {
        label: 'CREATE TEMPORARY TABLE ?',
        filterText: 'CREATE',
        template: `CREATE TEMPORARY TABLE \${1:__tableName}
(
    \${2:__columnDef} \${3:__columnType}
    -- ...
)
COMMENT \${4:'__comment'}
WITH (
    \${5:'__config_key'} = \${6:'__config_value'}
    -- ...
)
;`
    },
    {
        label: 'CREATE TEMPORARY VIEW ?',
        filterText: 'CREATE',
        template: `CREATE TEMPORARY VIEW \${1:__viewName}
AS SELECT \${2:__projectItems} FROM \${3:__tableReferences}
;`
    },
    {
        label: 'CREATE MATERIALIZED TABLE ?',
        filterText: 'CREATE',
        template: `CREATE MATERIALIZED TABLE \${1:__tableName}
    FRESHNESS = INTERVAL '\${2:__positiveInteger}' \${3:__timeUnit}
    AS
    SELECT \${4:__projectItems} FROM \${5:__tableReferences}
;`
    },
    {
        label: 'INSERT INTO ?',
        filterText: 'INSERT',
        template: `INSERT INTO \${1:__tableName}
SELECT \${2:__projectItems} FROM \${3:__tableReferences}
;`
    },
    {
        label: 'SELECT ? FROM ?',
        filterText: 'SELECT',
        template: `SELECT \${1:__projectItems} FROM \${2:__tableReferences}`
    }
];
class CompletionVisitor extends flinksql_relation_visitor_1.FlinkSQLRelationVisitor {
    constructor(cursor, getColumns) {
        super(getColumns && getTableMetadataFnFactory(getColumns));
        this.cursor = cursor;
        this.getColumns = getColumns;
        this.hasCompletions = false;
        this.completions = { type: 'other', snippets: [] };
    }
    defaultResult() {
        return;
    }
    aggregateResult() {
        return;
    }
    onColumnReference() {
        return;
    }
    updateQueryCompletionItems(relation) {
        if (!this.hasCompletions && this.caretScope !== undefined) {
            this.hasCompletions = true;
            switch (this.caretScope.type) {
                case 'select-column':
                case 'spec-column': {
                    const columns = availableColumns(relation);
                    this.completions = {
                        type: 'column',
                        columns,
                        relations: [...new Set(columns.filter(c => c.relation !== undefined).map(c => c.relation))],
                        snippets: this.completions.snippets
                    };
                    break;
                }
                case 'scoped-column': {
                    let newCompletions = [];
                    const relationName = this.caretScope.relation;
                    const relationFound = this.relationDDLColumns.get(relationName);
                    if (relationFound) {
                        newCompletions = relationFound.map(c => ({
                            relation: { category: 'table', label: getRelationName(relationName) },
                            name: c.label
                        }));
                    }
                    else {
                        const cteFound = relation.findCTE(getRelationName(relationName));
                        if (cteFound) {
                            newCompletions = cteFound.columns.map(c => ({
                                relation: { category: 'CTE', label: getRelationName(relationName) },
                                name: c.label
                            }));
                        }
                        else {
                            newCompletions =
                                relation
                                    .findLocalRelation(relationName)
                                    ?.columns.filter(c => c.label !== flinksql_relation_visitor_1.STAR_LABEL && !c.label.includes(cursor_1.defaultCursorPlaceholder))
                                    .map(c => ({
                                    relation: { category: 'table', label: getRelationName(relationName) },
                                    name: c.label
                                })) ?? [];
                        }
                    }
                    this.completions = {
                        type: 'column',
                        relations: [],
                        columns: newCompletions,
                        snippets: this.completions.snippets
                    };
                    break;
                }
                case 'relation': {
                    const ctes = relation.getCTENames().map(label => ({
                        label,
                        category: 'CTE'
                    }));
                    const tableDDLNames = [...this.tableDDLs.keys()].map(label => ({
                        label: getRelationName(label),
                        category: 'table'
                    }));
                    const viewDDLNames = [...this.viewDDLs.keys()].map(label => ({
                        label: getRelationName(label),
                        category: 'view'
                    }));
                    this.completions = {
                        type: 'relation',
                        relations: [...ctes, ...tableDDLNames, ...viewDDLNames],
                        snippets: this.completions.snippets,
                        incompleteReference: this.caretScope.prefix.length > 0 ? this.caretScope.prefix : undefined
                    };
                    break;
                }
                default:
                    break;
            }
        }
    }
    handleRelationContext(ctx) {
        const multipartRelationName = ctx.text.split('.').map(t => (0, utils_1.sanitizeText)(t));
        const lastPart = multipartRelationName[multipartRelationName.length - 1];
        const category = ctx instanceof FlinkSQLParser_1.TableContext ? 'table' : 'view';
        this.relationInStatementContext = this.relationPrimaryFromMultipart(multipartRelationName, category);
        if (!this.hasCompletions && this.cursor.isIn(lastPart)) {
            this.hasCompletions = true;
            this.caretScope = { type: 'relation', prefix: multipartRelationName.slice(0, -1) };
            this.completions = {
                type: 'relation',
                relations: [...(category === 'table' ? this.tableDDLs : this.viewDDLs).keys()].map(label => ({
                    label: getRelationName(label),
                    category
                })),
                snippets: this.completions.snippets,
                incompleteReference: this.caretScope.prefix.length > 0 ? this.caretScope.prefix : undefined
            };
        }
    }
    handleScopedColumn() {
        if (!this.hasCompletions && this.relationInStatementContext) {
            this.hasCompletions = true;
            this.caretScope = {
                type: 'scoped-column',
                relation: (0, flinksql_relation_visitor_1.relationPrimaryToString)(this.relationInStatementContext)
            };
            let newCompletions = [];
            const relationName = this.relationInStatementContext.relationName;
            const relationDDL = this.relationDDLColumns.get(relationName);
            if (relationDDL) {
                newCompletions = relationDDL.map(c => ({
                    relation: { category: 'table', label: relationName },
                    name: c.label
                }));
            }
            else if (this.getColumns) {
                newCompletions = this.getColumns({
                    catalog: this.relationInStatementContext.catalogName,
                    database: this.relationInStatementContext.databaseName,
                    table: this.relationInStatementContext.relationName
                }).map(c => ({
                    relation: { category: 'table', label: relationName },
                    name: c.label
                }));
            }
            this.completions = {
                type: 'column',
                relations: [],
                columns: newCompletions,
                snippets: this.completions.snippets
            };
        }
    }
    handleDatabaseContext(ctx) {
        const multipartDatabaseName = ctx.text.split('.').map(t => (0, utils_1.sanitizeText)(t));
        const lastPart = multipartDatabaseName[multipartDatabaseName.length - 1];
        if (this.cursor.isIn(lastPart)) {
            this.hasCompletions = true;
            const incompleteReference = multipartDatabaseName.length > 1 ? multipartDatabaseName[0] : undefined;
            this.completions = { type: 'database', incompleteReference, snippets: [] };
        }
    }
    getSuggestions() {
        return this.completions;
    }
    onRelation(relation) {
        if (relation instanceof flinksql_relation_visitor_1.TableRelation || relation instanceof flinksql_relation_visitor_1.InsertRelation) {
            return;
        }
        this.updateQueryCompletionItems(relation);
    }
    visitErrorNode(node) {
        super.visitErrorNode(node);
        if (this.cursor.isIn(node.text)) {
            if (node.parent instanceof FlinkSQLParser_1.StatementContext || node.parent instanceof FlinkSQLParser_1.ProgramContext) {
                this.completions.snippets.push(...SQL_SNIPPETS.filter(s => s.filterText.toUpperCase().indexOf(this.cursor.revert(node.text).toUpperCase()) === 0));
                this.caretScope = { type: 'other' };
            }
        }
    }
    visitStatement(ctx) {
        this.relationInStatementContext = undefined;
        super.visitStatement(ctx);
    }
    visitUse(ctx) {
        super.visitUse(ctx);
        if (ctx.CATALOG() && ctx.catalog()) {
            if (this.cursor.isIn(ctx.catalog().text)) {
                this.hasCompletions = true;
                this.completions = { type: 'catalog', snippets: [] };
            }
        }
        else if (ctx.database()) {
            this.handleDatabaseContext(ctx.database());
        }
    }
    visitAnalyze(ctx) {
        if (ctx.TABLE().text.toUpperCase() === 'TABLE') {
            this.handleRelationContext(ctx.table());
        }
        if (ctx.partition()) {
            const cursorCol = ctx
                .partition()
                .columnAssignments()
                .columnAssignment()
                .find(assign => this.cursor.isIn(assign.columnName().text));
            if (cursorCol !== undefined) {
                return this.handleScopedColumn();
            }
        }
        if (ctx.columnsWithoutParenthesis()) {
            const cursorCol = ctx
                .columnsWithoutParenthesis()
                .columnName()
                .find(col => this.cursor.isIn(col.text));
            if (cursorCol !== undefined) {
                return this.handleScopedColumn();
            }
        }
    }
    visitCreateTable(ctx) {
        super.visitCreateTable(ctx);
        if (ctx.likeTable()) {
            super.visit(ctx.likeTable());
        }
    }
    visitLikeTable(ctx) {
        this.handleRelationContext(ctx.table());
        super.visitChildren(ctx);
    }
    visitAsDatabase(ctx) {
        this.handleDatabaseContext(ctx.database());
        super.visitChildren(ctx);
    }
    visitDropDatabase(ctx) {
        this.handleDatabaseContext(ctx.database());
    }
    visitDropTable(ctx) {
        this.handleRelationContext(ctx.table());
        super.visitChildren(ctx);
    }
    visitAutoOptimizeTable(ctx) {
        this.handleRelationContext(ctx.table());
        super.visitChildren(ctx);
    }
    visitDescribeCatalog(ctx) {
        if (this.cursor.isIn(ctx.catalog().text)) {
            this.hasCompletions = true;
            this.completions = { type: 'catalog', snippets: [] };
        }
    }
    visitDescribeDatabase(ctx) {
        this.handleDatabaseContext(ctx.database());
    }
    visitDescribeTable(ctx) {
        this.handleRelationContext(ctx.table());
        super.visitChildren(ctx);
    }
    visitAlterDatabase(ctx) {
        this.handleDatabaseContext(ctx.database());
        super.visitChildren(ctx);
    }
    visitAlterTable(ctx) {
        this.handleRelationContext(ctx.table());
        super.visitChildren(ctx);
    }
    visitAlterView(ctx) {
        this.handleRelationContext(ctx.view());
        super.visitChildren(ctx);
    }
    visitInsert(ctx) {
        super.visitInsert(ctx);
        if (ctx.INTO()?.text.toUpperCase() === 'INTO' || ctx.OVERWRITE()?.text.toUpperCase() === 'OVERWRITE') {
            this.handleRelationContext(ctx.table());
        }
        if (this.relationInStatementContext) {
            if (ctx.partition()) {
                const cursorCol = ctx
                    .partition()
                    .columnAssignments()
                    .columnAssignment()
                    .find(assign => this.cursor.isIn(assign.columnName().text));
                if (cursorCol !== undefined) {
                    return this.handleScopedColumn();
                }
            }
            if (ctx.columns()) {
                const cursorCol = ctx
                    .columns()
                    .columnName()
                    .find(col => this.cursor.isIn(col.text));
                if (cursorCol !== undefined) {
                    return this.handleScopedColumn();
                }
            }
        }
    }
    visitQueryPrimary(ctx) {
        super.visitQueryPrimary(ctx);
        if (this.currentQueryRelation !== undefined) {
            this.updateQueryCompletionItems(this.currentQueryRelation);
        }
    }
    visitTableReference(ctx) {
        if (ctx.tablePrimary().tableSource()?.table() !== undefined) {
            const multipartTableName = ctx
                .tablePrimary()
                .tableSource()
                .table()
                .text.split('.')
                .map(t => (0, utils_1.sanitizeText)(t));
            const lastPart = multipartTableName[multipartTableName.length - 1];
            if (this.cursor.isIn(lastPart)) {
                this.caretScope = { type: 'relation', prefix: multipartTableName.slice(0, -1) };
            }
        }
        super.visitTableReference(ctx);
    }
    visitColumnReference(ctx) {
        super.visitColumnReference(ctx);
        if (this.currentQueryRelation !== undefined && this.cursor.isIn(ctx.text)) {
            if (this.currentQueryRelation.currentClause === 'select') {
                this.caretScope = { type: 'select-column' };
            }
            else {
                this.caretScope = { type: 'spec-column' };
            }
        }
    }
    visitDereference(ctx) {
        super.visitDereference(ctx);
        if (this.currentQueryRelation !== undefined && this.cursor.isIn(ctx.columnName().text)) {
            this.caretScope = { type: 'scoped-column', relation: (0, utils_1.sanitizeText)(ctx.table().tableName().text) };
        }
    }
    visitJoinCondition(ctx) {
        super.visitJoinCondition(ctx);
        if (ctx.columns()) {
            this.handleColumnsContext(ctx.columns());
        }
    }
    visitPartitionedBy(ctx) {
        super.visitChildren(ctx);
        if (ctx.columns()) {
            this.handleColumnsContext(ctx.columns());
        }
    }
    handleColumnsContext(ctx) {
        if (this.currentQueryRelation !== undefined) {
            const cursorCol = ctx?.columnName().find(col => this.cursor.isIn(col.text));
            if (cursorCol !== undefined) {
                if (this.currentQueryRelation.currentClause === 'select') {
                    this.caretScope = { type: 'select-column' };
                }
                else {
                    this.caretScope = { type: 'spec-column' };
                }
            }
        }
    }
}
exports.CompletionVisitor = CompletionVisitor;
