var __rest = (this && this.__rest) || function (s, e) {
    var t = {};
    for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
        t[p] = s[p];
    if (s != null && typeof Object.getOwnPropertySymbols === "function")
        for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
            if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
                t[p[i]] = s[p[i]];
        }
    return t;
};
import { isToken } from 'antlr4ng';
import { SimpleStack } from './simpleStack';
import { ctxToText, isWordRange, tokenToWord, } from './textAndWord';
import { EntityContextType } from './types';
/**
 * TODO: more stmt type should be supported.
 */
export var StmtContextType;
(function (StmtContextType) {
    /** A self-contained and complete statement */
    StmtContextType["COMMON_STMT"] = "commonStmt";
    StmtContextType["CREATE_CATALOG_STMT"] = "createCatalogStmt";
    StmtContextType["CREATE_DATABASE_STMT"] = "createDatabaseStmt";
    StmtContextType["CREATE_TABLE_STMT"] = "createTableStmt";
    StmtContextType["CREATE_VIEW_STMT"] = "createViewStmt";
    StmtContextType["SELECT_STMT"] = "selectStmt";
    StmtContextType["INSERT_STMT"] = "insertStmt";
    StmtContextType["CREATE_FUNCTION_STMT"] = "createFunctionStmt";
    StmtContextType["ALTER_TABLE_STMT"] = "alterTableStmt";
})(StmtContextType || (StmtContextType = {}));
export function toStmtContext(ctx, type, input, rootStmt, parentStmt, isContainCaret) {
    const text = ctxToText(ctx, input);
    if (!text)
        return null;
    const { text: _ } = text, position = __rest(text, ["text"]);
    return {
        stmtContextType: type,
        position,
        rootStmt: rootStmt !== null && rootStmt !== void 0 ? rootStmt : null,
        parentStmt: parentStmt !== null && parentStmt !== void 0 ? parentStmt : null,
        isContainCaret,
    };
}
/**
 * entity's attribute
 * @key comment: entity's comment attribute
 * @key colType: column entity's type attribute
 * @key alias: entity's alias attribute
 * */
export var AttrName;
(function (AttrName) {
    AttrName["comment"] = "_comment";
    AttrName["colType"] = "_colType";
    AttrName["alias"] = "_alias";
})(AttrName || (AttrName = {}));
export function isCommonEntityContext(entity) {
    if (!entity)
        return false;
    return 'relatedEntities' in entity && !('arguments' in entity);
}
export function isFuncEntityContext(entity) {
    if (!entity)
        return false;
    return 'arguments' in entity;
}
export function isColumnEntityContext(entity) {
    if (!entity)
        return false;
    return AttrName.colType in entity;
}
export function toEntityContext(ctx, type, input, belongStmt, attrInfo) {
    var _a, _b;
    const word = ctxToText(ctx, input);
    if (!word)
        return null;
    const { text, startLine, endLine } = word, rest = __rest(word, ["text", "startLine", "endLine"]);
    const position = Object.assign(Object.assign({}, rest), { line: startLine });
    let entityInfo = {
        entityContextType: type,
        text,
        position,
        belongStmt,
        [AttrName.comment]: null,
    };
    switch (entityInfo.entityContextType) {
        case EntityContextType.FUNCTION:
        case EntityContextType.FUNCTION_CREATE: {
            entityInfo.relatedEntities = null;
            entityInfo.arguments = null;
            break;
        }
        case EntityContextType.COLUMN:
        case EntityContextType.COLUMN_CREATE:
            entityInfo[AttrName.colType] = null;
            break;
        default:
            entityInfo.relatedEntities = null;
            break;
    }
    if (attrInfo) {
        for (let k = 0; k < attrInfo.length; k++) {
            const attributeName = (_a = attrInfo[k]) === null || _a === void 0 ? void 0 : _a.attrName;
            const attrToken = findAttribute(ctx, attributeName, (_b = attrInfo[k]) === null || _b === void 0 ? void 0 : _b.endContextList);
            if (attrToken) {
                const attrVal = isToken(attrToken)
                    ? tokenToWord(attrToken, input)
                    : ctxToText(attrToken, input);
                if (attrVal) {
                    if (isWordRange(attrVal)) {
                        entityInfo = Object.assign(entityInfo, {
                            [attributeName]: attrVal,
                        });
                    }
                    else {
                        const { startIndex, endIndex, startColumn, endColumn, text } = attrVal, rest = __rest(attrVal, ["startIndex", "endIndex", "startColumn", "endColumn", "text"]);
                        entityInfo = Object.assign(entityInfo, {
                            [attributeName]: {
                                startIndex,
                                endIndex,
                                startColumn,
                                endColumn,
                                text,
                                line: rest === null || rest === void 0 ? void 0 : rest.startLine,
                            },
                        });
                    }
                }
            }
        }
    }
    return entityInfo;
}
export function findAttribute(ctx, keyName, endContextNameList) {
    var _a, _b;
    const parent = (ctx === null || ctx === void 0 ? void 0 : ctx.parent) || null;
    let attrVal = null;
    if (parent === null || parent === void 0 ? void 0 : parent[keyName]) {
        attrVal = (parent === null || parent === void 0 ? void 0 : parent[keyName]) || null;
        return attrVal;
    }
    else {
        if (((_a = parent === null || parent === void 0 ? void 0 : parent.constructor) === null || _a === void 0 ? void 0 : _a.name) && !endContextNameList.includes((_b = parent === null || parent === void 0 ? void 0 : parent.constructor) === null || _b === void 0 ? void 0 : _b.name)) {
            attrVal = findAttribute(parent, keyName, endContextNameList);
        }
        if (!attrVal) {
            if (parent === null || parent === void 0 ? void 0 : parent.children) {
                attrVal = findAttributeChildren(parent, keyName);
            }
        }
    }
    return attrVal;
}
function findAttributeChildren(ctx, keyName) {
    const visitChildren = (ctx === null || ctx === void 0 ? void 0 : ctx.children) || [];
    let attrVal = null;
    if (visitChildren.length) {
        for (let i = 0; i < visitChildren.length; i++) {
            const child = visitChildren[i] || null;
            if (child === null || child === void 0 ? void 0 : child[keyName]) {
                attrVal = (child === null || child === void 0 ? void 0 : child[keyName]) || null;
                return attrVal;
            }
            else {
                attrVal = findAttributeChildren(child, keyName);
            }
        }
    }
    return attrVal;
}
/**
 * @todo: Handle alias, includes column alias, table alias, query as alias and so on.
 * @todo: [may be need] Combine the entities in each clause.
 */
export class EntityCollector {
    constructor(input, caretTokenIndex) {
        this._input = input;
        this._caretTokenIndex = caretTokenIndex !== null && caretTokenIndex !== void 0 ? caretTokenIndex : -1;
        this._entitiesSet = new Set();
        this._stmtStack = new SimpleStack();
        this._entityStack = new SimpleStack();
        this._rootStmt = null;
    }
    visitTerminal() { }
    visitErrorNode() { }
    enterEveryRule() { }
    exitEveryRule() { }
    getRootStmt() {
        return this._rootStmt;
    }
    getEntities() {
        return Array.from(this._entitiesSet);
    }
    enterProgram() {
        this._entitiesSet.clear();
        this._stmtStack.clear();
        this._entityStack.clear();
        this._rootStmt = null;
    }
    pushStmt(ctx, type) {
        var _a;
        let isContainCaret;
        if (this._caretTokenIndex >= 0) {
            isContainCaret =
                !!ctx.start &&
                    !!ctx.stop &&
                    ctx.start.tokenIndex <= this._caretTokenIndex &&
                    ctx.stop.tokenIndex >= this._caretTokenIndex;
        }
        const stmtContext = toStmtContext(ctx, type, this._input, this._rootStmt, this._stmtStack.peek(), isContainCaret);
        if (stmtContext) {
            if (this._stmtStack.isEmpty() ||
                ((_a = this._stmtStack.peek()) === null || _a === void 0 ? void 0 : _a.stmtContextType) === StmtContextType.COMMON_STMT) {
                this._rootStmt = stmtContext;
            }
            this._stmtStack.push(stmtContext);
        }
        return stmtContext;
    }
    popStmt() {
        const stmtContext = this._stmtStack.pop();
        if (stmtContext && this._rootStmt === stmtContext) {
            this._rootStmt = this._stmtStack.peek();
            if (!this._entityStack.isEmpty()) {
                this.combineEntitiesAndAdd(stmtContext);
            }
        }
        return stmtContext;
    }
    pushEntity(ctx, type, attrInfo) {
        const entityContext = toEntityContext(ctx, type, this._input, this._stmtStack.peek(), attrInfo);
        if (entityContext) {
            if (this._stmtStack.isEmpty()) {
                this._entitiesSet.add(entityContext);
            }
            else {
                // If is inside a statement
                this._entityStack.push(entityContext);
            }
        }
        return entityContext;
    }
    /**
     * Combine entities that inside a single statement.
     * e.g. combine tableName and column if they are inside a same createTableStatement.
     * Then add combined entities into result.
     */
    combineEntitiesAndAdd(stmtContext) {
        const entitiesInsideStmt = [];
        while (!this._entityStack.isEmpty() &&
            (this._entityStack.peek().belongStmt === stmtContext ||
                this._entityStack.peek().belongStmt.rootStmt === stmtContext)) {
            entitiesInsideStmt.unshift(this._entityStack.pop());
        }
        const combinedEntities = this.combineRootStmtEntities(stmtContext, entitiesInsideStmt);
        while (combinedEntities.length) {
            const entity = combinedEntities.shift();
            entity && this._entitiesSet.add(entity);
        }
    }
    /**
     * Combined all entities under a rootStmt.
     */
    combineRootStmtEntities(stmtContext, entitiesInsideStmt) {
        if (stmtContext.stmtContextType === StmtContextType.CREATE_VIEW_STMT ||
            stmtContext.stmtContextType === StmtContextType.CREATE_TABLE_STMT) {
            return this.combineCreateTableOrViewStmtEntities(stmtContext, entitiesInsideStmt);
        }
        return entitiesInsideStmt;
    }
    combineCreateTableOrViewStmtEntities(stmtContext, entitiesInsideStmt) {
        const columns = [];
        const relatedEntities = [];
        let mainEntity = null;
        const finalEntities = entitiesInsideStmt.reduce((result, entity) => {
            if (entity.belongStmt !== stmtContext) {
                if (entity.entityContextType !== EntityContextType.COLUMN &&
                    entity.entityContextType !== EntityContextType.COLUMN_CREATE) {
                    relatedEntities.push(entity);
                    result.push(entity);
                }
                return result;
            }
            if (entity.entityContextType === EntityContextType.COLUMN_CREATE) {
                columns.push(entity);
            }
            else if (entity.entityContextType === EntityContextType.TABLE_CREATE ||
                entity.entityContextType === EntityContextType.VIEW_CREATE) {
                mainEntity = entity;
                result.push(entity);
                return result;
            }
            else if (entity.entityContextType !== EntityContextType.COLUMN) {
                relatedEntities.push(entity);
                result.push(entity);
            }
            return result;
        }, []);
        if (mainEntity && columns.length) {
            if (isCommonEntityContext(mainEntity)) {
                mainEntity = Object.assign(mainEntity, {
                    columns,
                });
            }
        }
        if (mainEntity && relatedEntities.length) {
            if (isCommonEntityContext(mainEntity) || isFuncEntityContext(mainEntity)) {
                mainEntity = Object.assign(mainEntity, {
                    relatedEntities,
                });
            }
        }
        return finalEntities;
    }
}
