"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.tablePrimaryFromMultipart = exports.QueryStructureVisitor = exports.QueryRelation = exports.TableRelation = exports.Relation = exports.Column = exports.ROOT_QUERY_NAME = void 0;
const tree_1 = require("antlr4ts/tree");
const FlinkSQLParser_1 = require("../lib/FlinkSQLParser");
const utils_1 = require("../utils");
const ROOT_QUERY_ID = 'result_1';
exports.ROOT_QUERY_NAME = '[final result]';
class Column {
    constructor(id, label, range, data, isAssumed) {
        this.id = id;
        this.label = label;
        this.range = range;
        this.data = data;
        this.isAssumed = isAssumed;
        this.columnReferences = [];
    }
}
exports.Column = Column;
class Relation {
    constructor(id, columns, parent, range) {
        this.id = id;
        this.columns = columns;
        this.parent = parent;
        this.range = range;
    }
    findColumn(columnName) {
        return this.columns.find(c => c.label == columnName);
    }
    resolveColumn(columnName) {
        const col = this.findColumn(columnName);
        return col !== undefined
            ? {
                tableId: this.id,
                columnId: col.id,
                isAssumed: false
            }
            : undefined;
    }
}
exports.Relation = Relation;
class TableRelation extends Relation {
    constructor(id, tablePrimary, columns, isFetched, parent, range, data) {
        super(id, columns, parent, range);
        this.tablePrimary = tablePrimary;
        this.isFetched = isFetched;
        this.data = data;
    }
    addAssumedColumn(columnName, range) {
        const column = new Column(`column_${this.columns.length + 1}`, columnName, range, undefined, true);
        this.columns.push(column);
        return { tableId: this.id, columnId: column.id, isAssumed: true };
    }
}
exports.TableRelation = TableRelation;
class QueryRelation extends Relation {
    constructor(id, parent, range) {
        super(id, [], parent, range);
        this.ctes = new Map();
        this.relations = new Map();
        this.columnIdSeq = 0;
        this.columnReferences = [];
    }
    getNextColumnId() {
        this.columnIdSeq++;
        return `column_${this.columnIdSeq}`;
    }
    findLocalRelation(tableName) {
        for (const rel of this.relations) {
            if (tableName.localeCompare(rel[0], undefined, {
                sensitivity: 'accent'
            }) == 0)
                return rel[1];
        }
        return undefined;
    }
    findRelation(tableName) {
        return this.findLocalRelation(tableName) ?? this.parent?.findRelation(tableName);
    }
    findCTE(tableName) {
        for (const rel of this.ctes) {
            if (tableName.localeCompare(rel[0], undefined, {
                sensitivity: 'accent'
            }) == 0)
                return rel[1];
        }
        return this.parent?.findCTE(tableName);
    }
    getCTENames() {
        const localCTEs = Array.from(this.ctes.keys());
        if (this.parent !== undefined) {
            return localCTEs.concat(this.parent.getCTENames());
        }
        else {
            return localCTEs;
        }
    }
    resolveOrAssumeRelationColumn(columnName, range, tableName) {
        if (tableName !== undefined) {
            const rel = this.findRelation(tableName);
            const col = rel?.resolveColumn(columnName);
            if (col === undefined && rel != undefined && rel instanceof TableRelation) {
                return rel.addAssumedColumn(columnName, range);
            }
            return col;
        }
        else {
            const tables = [];
            for (const r of this.relations) {
                const rel = r[1];
                const col = rel.resolveColumn(columnName);
                if (col) {
                    return col;
                }
                if (rel instanceof TableRelation) {
                    tables.push(rel);
                }
            }
            const assumed = tables.filter(t => !t.isFetched);
            if (assumed.length == 1) {
                return assumed[0].addAssumedColumn(columnName, range);
            }
            else if (tables.length == 1) {
                return tables[0].addAssumedColumn(columnName, range);
            }
            return undefined;
        }
    }
}
exports.QueryRelation = QueryRelation;
class QueryStructureVisitor extends tree_1.AbstractParseTreeVisitor {
    constructor(getTable, options) {
        super();
        this.getTable = getTable;
        this.options = options;
        this.relationSeq = 0;
        this.currentRelation = new QueryRelation(this.getNextRelationId());
    }
    getNextRelationId() {
        return `result_${this.relationSeq++}`;
    }
    rangeFromContext(ctx) {
        const stop = ctx.stop ?? ctx.start;
        return {
            start: {
                line: ctx.start.line,
                character: ctx.start.charPositionInLine
            },
            end: {
                line: stop.line,
                character: stop.charPositionInLine + (stop.stopIndex - stop.startIndex + 1)
            }
        };
    }
    extractTableAndColumn(ctx) {
        if (ctx instanceof FlinkSQLParser_1.ColumnReferenceContext) {
            return { columnName: (0, utils_1.sanitizeText)(ctx.columnName().text) };
        }
        if (ctx instanceof FlinkSQLParser_1.DereferenceContext) {
            return {
                tableName: (0, utils_1.sanitizeText)(ctx.table().tableName().text),
                columnName: (0, utils_1.sanitizeText)(ctx.columnName().text)
            };
        }
        return undefined;
    }
    asPrimaryExpression(ctx) {
        const boolExpr = ctx.booleanExpr();
        if (boolExpr.predicate()) {
            const valExpr = boolExpr.valueExpr();
            return valExpr?.primaryExpr();
        }
        return undefined;
    }
    deriveColumnName(ctx) {
        const primExpr = this.asPrimaryExpression(ctx);
        if (primExpr?.column()) {
            return this.extractTableAndColumn(primExpr.column())?.columnName;
        }
        return undefined;
    }
    reportTableReferences() {
        for (const [alias, relation] of this.currentRelation.relations) {
            if (relation instanceof TableRelation) {
                this.onRelation(relation, alias !== relation.id ? alias : undefined);
            }
        }
    }
    isStar(ctx) {
        const boolExpr = ctx.booleanExpr();
        if (boolExpr.predicate()) {
            const valExpr = boolExpr.valueExpr();
            if (valExpr?.primaryExpr()?.column() instanceof FlinkSQLParser_1.StarContext) {
                return valExpr.primaryExpr().column();
            }
        }
        return undefined;
    }
    processStar(ctx) {
        const range = this.rangeFromContext(ctx);
        if (ctx.table() !== undefined) {
            const tableName = (0, utils_1.sanitizeText)(ctx.table().tableName().text);
            const rel = this.currentRelation.findLocalRelation(tableName);
            if (rel !== undefined) {
                this.addRelationColumns(rel, range);
            }
        }
        else {
            for (const r of this.currentRelation.relations) {
                this.addRelationColumns(r[1], range);
            }
        }
        return this.visitChildren(ctx);
    }
    addRelationColumns(rel, range) {
        rel.columns.forEach(c => {
            const columnId = this.currentRelation.getNextColumnId();
            const col = new Column(columnId, c.label, range);
            this.currentRelation.columns.push(col);
            this.currentRelation.currentColumnId = columnId;
            this.onColumnReference(rel.id, c.id);
            this.currentRelation.currentColumnId = undefined;
        });
    }
    processClause(clause, ctx) {
        this.currentRelation.currentClause = clause;
        const result = this.visitChildren(ctx);
        this.currentRelation.currentClause = undefined;
        return result;
    }
    visitQuery(ctx) {
        this.currentRelation = new QueryRelation(this.getNextRelationId(), this.currentRelation, this.rangeFromContext(ctx));
        const result = this.visitChildren(ctx);
        this.reportTableReferences();
        this.lastRelation = this.currentRelation;
        if (this.currentRelation.id == ROOT_QUERY_ID) {
            this.onRelation(this.currentRelation, exports.ROOT_QUERY_NAME);
        }
        this.currentRelation = this.currentRelation.parent ?? new QueryRelation(this.getNextRelationId());
        return result;
    }
    visitQueryPrimary(ctx) {
        this.currentRelation.columnIdSeq = 0;
        this.reportTableReferences();
        this.currentRelation.relations = new Map();
        return this.visitChildren(ctx);
    }
    visitSelect(ctx) {
        let result = ctx.fromTable()?.accept(this) ?? this.defaultResult();
        ctx.children?.forEach(c => {
            if (!(c instanceof FlinkSQLParser_1.FromTableContext)) {
                result = this.aggregateResult(result, c.accept(this));
            }
        });
        return result;
    }
    visitJoinCondition(ctx) {
        if (ctx.USING() && ctx.columns()) {
            const columns = ctx
                .columns()
                .columnName()
                .map(c => (0, utils_1.sanitizeText)(c.text));
            const size = this.currentRelation.relations.size;
            if (size >= 2) {
                let i = 0;
                const foundLeftCol = new Set();
                this.currentRelation.relations.forEach(r => {
                    if (i == size - 1) {
                        columns.forEach(c => {
                            const col = r.resolveColumn(c);
                            if (col)
                                this.onColumnReference(col.tableId, col.columnId);
                        });
                    }
                    else {
                        columns.forEach((c, j) => {
                            if (!foundLeftCol.has(j)) {
                                const col = r.resolveColumn(c);
                                if (col)
                                    this.onColumnReference(col.tableId, col.columnId);
                                foundLeftCol.add(j);
                            }
                        });
                    }
                    i++;
                });
            }
        }
        return this.visitChildren(ctx);
    }
    visitTableReference(ctx) {
        if (ctx.tablePrimary().tableSource()) {
            const tableName = ctx
                .tablePrimary()
                .tableSource()
                .table()
                .text.split('.')
                .map(t => (0, utils_1.sanitizeText)(t));
            const alias = ctx.alias() !== undefined ? (0, utils_1.sanitizeText)(ctx.alias().text) : tableName[tableName.length - 1];
            if (tableName.length === 1) {
                const cte = this.currentRelation.findCTE(tableName[0]);
                if (cte !== undefined) {
                    this.currentRelation.relations.set(alias, cte);
                    return this.defaultResult();
                }
                const rel = this.currentRelation.findRelation(tableName[0]);
                if (rel !== undefined) {
                    if (tableName[0] !== alias) {
                        this.currentRelation.relations.set(alias, rel);
                    }
                    return this.defaultResult();
                }
            }
            const tablePrimary = tablePrimaryFromMultipart(tableName);
            const metadata = this.getTable(tablePrimary);
            const columns = metadata?.columns.map(c => new Column(c.id, c.id, undefined, c.data)) || [];
            const relation = new TableRelation(this.getNextRelationId(), tablePrimary, columns, metadata !== undefined, this.currentRelation, this.rangeFromContext(ctx), metadata?.table.data);
            this.currentRelation.relations.set(alias, relation);
        }
        else if (ctx.tablePrimary().tableQuery()) {
            const result = this.visitChildren(ctx);
            const relation = this.lastRelation;
            if (relation !== undefined) {
                const alias = ctx.alias() !== undefined ? (0, utils_1.sanitizeText)(ctx.alias().text) : relation.id;
                const columnAliases = ctx
                    .columns()
                    ?.columnName()
                    .map(c => (0, utils_1.sanitizeText)(c.text));
                if (columnAliases !== undefined) {
                    relation.columns.forEach((c, i) => {
                        c.label = columnAliases[i] ?? c.label;
                    });
                }
                this.currentRelation.relations.set(alias, relation);
                this.onRelation(relation, alias);
                return result;
            }
            else {
                throw new Error('Expecting sub-query relation to be in stack');
            }
        }
        return this.defaultResult();
    }
    visitWithItem(ctx) {
        const result = this.visitChildren(ctx);
        const relation = this.lastRelation;
        if (relation !== undefined) {
            const alias = (0, utils_1.sanitizeText)(ctx.tableName().text);
            const columnAliases = ctx
                .columns()
                ?.columnName()
                .map(t => (0, utils_1.sanitizeText)(t.text));
            if (columnAliases !== undefined) {
                relation.columns.forEach((c, i) => {
                    c.label = columnAliases[i] ?? c.label;
                });
            }
            this.onRelation(relation, alias);
            this.currentRelation.ctes.set(alias, relation);
            return result;
        }
        else {
            throw new Error('Expecting CTE query relation to be in stack');
        }
    }
    visitSelectClause(ctx) {
        return this.processClause('select', ctx);
    }
    visitFromTable(ctx) {
        return this.processClause('from', ctx);
    }
    visitWhere(ctx) {
        return this.processClause('where', ctx);
    }
    visitGroupBy(ctx) {
        return this.processClause('group by', ctx);
    }
    visitHaving(ctx) {
        return this.processClause('having', ctx);
    }
    visitOrderBy(ctx) {
        return this.processClause('order by', ctx);
    }
    visitExists(ctx) {
        const result = this.visitChildren(ctx);
        const rel = this.lastRelation;
        if (rel !== undefined) {
            this.onRelation(rel);
            this.onColumnReference(rel.id);
        }
        return result;
    }
    visitPredicate(ctx) {
        const result = this.visitChildren(ctx);
        if (ctx.parenthesisQuery()) {
            const rel = this.lastRelation;
            if (rel !== undefined) {
                this.onRelation(rel);
                this.onColumnReference(rel.id);
            }
        }
        return result;
    }
    visitNamedExpression(ctx) {
        if (ctx.alias() === undefined) {
            const star = this.isStar(ctx.expr());
            if (star !== undefined) {
                return this.processStar(star);
            }
        }
        const columnId = this.currentRelation.getNextColumnId();
        let column = this.currentRelation.columns.find(c => c.id == columnId);
        if (column === undefined) {
            const label = ctx.alias() !== undefined ? (0, utils_1.sanitizeText)(ctx.alias().text) : this.deriveColumnName(ctx.expr()) ?? columnId;
            const range = this.rangeFromContext(ctx);
            column = new Column(columnId, label, range);
            this.currentRelation.columns.push(column);
        }
        this.currentRelation.currentColumnId = columnId;
        this.currentRelation.columnReferences = [];
        const result = this.visitChildren(ctx);
        column.columnReferences.push(...this.currentRelation.columnReferences);
        this.currentRelation.currentColumnId = undefined;
        return result;
    }
    visitSubQuery(ctx) {
        const result = this.visitChildren(ctx);
        const rel = this.lastRelation;
        if (rel !== undefined) {
            this.onRelation(rel);
            this.onColumnReference(rel.id);
        }
        return result;
    }
    visitFunctionCall(ctx) {
        if ((0, utils_1.sanitizeText)(ctx.functionName().text.toLowerCase()) == 'count' &&
            ctx.expr().length == 1 &&
            this.isStar(ctx.expr()[0]) !== undefined) {
            for (const r of this.currentRelation.relations) {
                this.onColumnReference(r[1].id);
            }
        }
        return this.visitChildren(ctx);
    }
    visitOrderItem(ctx) {
        if (this.options?.positionalRefsEnabled) {
            const primExp = this.asPrimaryExpression(ctx.expr());
            if (primExp?.literalValue()?.REAL_LITERAL() !== undefined) {
                const idx = Number(primExp.text) - 1;
                if (idx in this.currentRelation.columns) {
                    const col = this.currentRelation.columns[idx];
                    col.columnReferences.forEach(cr => this.onColumnReference(cr.tableId, cr.columnId));
                    return this.defaultResult();
                }
            }
        }
        return this.visitChildren(ctx);
    }
    processColumnReference(ctx) {
        const tableCol = this.extractTableAndColumn(ctx);
        if (tableCol !== undefined) {
            if (tableCol.tableName === undefined &&
                this.currentRelation.currentClause !== undefined &&
                ['group by', 'order by'].includes(this.currentRelation.currentClause)) {
                const selfCol = this.currentRelation.findColumn(tableCol.columnName);
                if (selfCol) {
                    selfCol.columnReferences.forEach(cr => this.onColumnReference(cr.tableId, cr.columnId));
                    return this.defaultResult();
                }
            }
            const range = this.rangeFromContext(ctx);
            const col = this.currentRelation.resolveOrAssumeRelationColumn(tableCol.columnName, range, tableCol.tableName);
            if (col !== undefined) {
                this.currentRelation.columnReferences.push(col);
                this.onColumnReference(col.tableId, col.columnId);
            }
            return this.defaultResult();
        }
        else {
            return this.visitChildren(ctx);
        }
    }
    visitColumnReference(ctx) {
        return this.processColumnReference(ctx);
    }
    visitDereference(ctx) {
        return this.processColumnReference(ctx);
    }
}
exports.QueryStructureVisitor = QueryStructureVisitor;
function tablePrimaryFromMultipart(names) {
    const len = names.length;
    if (len == 0) {
        throw new Error('Unexpected empty array for table name');
    }
    else if (len == 1) {
        return {
            tableName: names[0]
        };
    }
    else if (len == 2) {
        return {
            schemaName: names[0],
            tableName: names[1]
        };
    }
    else {
        return {
            catalogName: names[len - 3],
            schemaName: names[len - 2],
            tableName: names[len - 1]
        };
    }
}
exports.tablePrimaryFromMultipart = tablePrimaryFromMultipart;
//# sourceMappingURL=query-structure.visitor.js.map