"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.InsertQueryVisitor = exports.InsertRelation = exports.QueryRelation = exports.TableRelation = exports.Relation = exports.Column = void 0;
const tree_1 = require("antlr4ts/tree");
const FlinkSQLParser_1 = require("../lib/FlinkSQLParser");
const utils_1 = require("../utils");
const query_structure_visitor_1 = require("./query-structure.visitor");
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, parent, range, data) {
        super(id, columns, parent, range);
        this.tablePrimary = tablePrimary;
        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) {
        const rel = this.findLocalRelation(tableName);
        if (rel) {
            return rel;
        }
        if (this.parent) {
            if (this.parent instanceof QueryRelation) {
                return this.parent.findRelation(tableName);
            }
            else {
                if (this.parent.sinkTableRelation?.tablePrimary.tableName.localeCompare(tableName, undefined, {
                    sensitivity: 'accent'
                }) == 0) {
                    return this.parent.sinkTableRelation;
                }
            }
        }
        return undefined;
    }
    findCTE(tableName) {
        for (const rel of this.ctes) {
            if (tableName.localeCompare(rel[0], undefined, {
                sensitivity: 'accent'
            }) == 0)
                return rel[1];
        }
        return this.parent instanceof QueryRelation ? this.parent.findCTE(tableName) : undefined;
    }
    resolveOrAssumeRelationColumn(columnName, range, tableName) {
        if (tableName !== undefined) {
            const rel = this.findRelation(tableName);
            console.log(tableName, rel, Array.from(this.relations));
            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);
                }
            }
            if (tables.length == 1) {
                return tables[0].addAssumedColumn(columnName, range);
            }
            return undefined;
        }
    }
}
exports.QueryRelation = QueryRelation;
class InsertRelation extends Relation {
    constructor(id, range) {
        super(id, [], undefined, range);
        this.columnIdSeq = 0;
    }
    getNextColumnId() {
        this.columnIdSeq++;
        return `${this.id}_column_${this.columnIdSeq}`;
    }
}
exports.InsertRelation = InsertRelation;
class InsertQueryVisitor extends tree_1.AbstractParseTreeVisitor {
    constructor() {
        super();
        this.relationInsertSeq = 0;
        this.relationSeq = 0;
    }
    getNextInsertRelationId() {
        return `insert_${this.relationInsertSeq++}`;
    }
    getNextRelationId() {
        return `relation_${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)
            }
        };
    }
    reportTableReferences() {
        for (const [alias, relation] of this.currentQueryRelation?.relations || []) {
            if (relation instanceof TableRelation) {
                this.onRelation(relation, alias !== relation.id ? alias : undefined);
            }
        }
    }
    visitInsert(ctx) {
        this.currentInsertRelation = new InsertRelation(this.getNextInsertRelationId());
        const tableName = ctx
            .table()
            .text.split('.')
            .map(t => (0, utils_1.sanitizeText)(t));
        const tablePrimary = (0, query_structure_visitor_1.tablePrimaryFromMultipart)(tableName);
        const columns = [];
        if (ctx.columns()) {
            ctx
                .columns()
                .columnName()
                .forEach(columnNameCtx => {
                columns.push(new Column(this.currentInsertRelation.getNextColumnId(), (0, utils_1.sanitizeText)(columnNameCtx.text), this.rangeFromContext(columnNameCtx), tablePrimary, false));
            });
        }
        const relation = new TableRelation(this.getNextRelationId(), tablePrimary, columns, undefined, this.rangeFromContext(ctx), undefined);
        this.currentInsertRelation.sinkTableRelation = relation;
        const result = this.visitChildren(ctx);
        this.onRelation(relation, tablePrimary.tableName);
        this.currentInsertRelation = undefined;
        return result;
    }
    visitQuery(ctx) {
        this.currentQueryRelation = new QueryRelation(this.getNextRelationId(), this.currentQueryRelation ?? this.currentInsertRelation, this.rangeFromContext(ctx));
        this.lastQueryRelation = this.currentQueryRelation;
        const result = this.visitChildren(ctx);
        this.reportTableReferences();
        if (this.currentQueryRelation.parent instanceof QueryRelation) {
            this.currentQueryRelation = this.currentQueryRelation.parent;
        }
        else {
            console.log(this.currentQueryRelation.parent);
            this.currentQueryRelation = new QueryRelation(this.getNextRelationId());
        }
        return result;
    }
    visitQueryPrimary(ctx) {
        if (this.currentQueryRelation !== undefined) {
            this.currentQueryRelation.columnIdSeq = 0;
            this.reportTableReferences();
            this.currentQueryRelation.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;
    }
    visitTableReference(ctx) {
        if (this.currentQueryRelation !== undefined) {
            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.currentQueryRelation.findCTE(tableName[0]);
                    if (cte !== undefined) {
                        this.currentQueryRelation.relations.set(alias, cte);
                        return this.defaultResult();
                    }
                    const rel = this.currentQueryRelation.findRelation(tableName[0]);
                    if (rel !== undefined) {
                        if (tableName[0] !== alias) {
                            this.currentQueryRelation.relations.set(alias, rel);
                        }
                        return this.defaultResult();
                    }
                }
                const tablePrimary = (0, query_structure_visitor_1.tablePrimaryFromMultipart)(tableName);
                const relation = new TableRelation(this.getNextRelationId(), tablePrimary, [], this.currentQueryRelation, this.rangeFromContext(ctx), tablePrimary);
                this.currentQueryRelation.relations.set(alias, relation);
            }
            else if (ctx.tablePrimary().tableQuery()) {
                const result = this.visitChildren(ctx);
                const relation = this.lastQueryRelation;
                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.currentQueryRelation.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();
    }
    visitColumnReference(ctx) {
        if (this.currentQueryRelation) {
            const columnName = (0, utils_1.sanitizeText)(ctx.columnName().text);
            const range = this.rangeFromContext(ctx);
            const col = this.currentQueryRelation.resolveOrAssumeRelationColumn(columnName, range);
            if (col !== undefined) {
                this.currentQueryRelation.columnReferences.push(col);
                this.onColumnReference(col.tableId, col.columnId);
            }
        }
        return this.defaultResult();
    }
    visitDereference(ctx) {
        if (this.currentQueryRelation) {
            const tableName = (0, utils_1.sanitizeText)(ctx.table().tableName().text);
            const columnName = (0, utils_1.sanitizeText)(ctx.columnName().text);
            const range = this.rangeFromContext(ctx);
            const col = this.currentQueryRelation.resolveOrAssumeRelationColumn(columnName, range, tableName);
            if (col !== undefined) {
                this.currentQueryRelation.columnReferences.push(col);
                this.onColumnReference(col.tableId, col.columnId);
            }
        }
        return this.defaultResult();
    }
}
exports.InsertQueryVisitor = InsertQueryVisitor;
//# sourceMappingURL=insert-query.visitor.js.map