package com.iformall.plugin;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.StringValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.MultiExpressionList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.create.table.ColDataType;
import net.sf.jsqlparser.statement.create.table.ColumnDefinition;
import net.sf.jsqlparser.statement.create.table.CreateTable;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;

import java.sql.Connection;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;

/**
 *共享数据库的多租户系统实现
 * Created by Stormeye on 2018/10/22.
 */
@Intercepts({
        //@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        //@Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class})
        @Signature( method = "prepare", type = StatementHandler.class,args = {Connection.class, Integer.class}),
        @Signature(method = "query", type = StatementHandler.class, args = {java.sql.Statement.class, ResultHandler.class})
})
public class MultiTenancy implements Interceptor {

    //获取tenantId的接口
    private TenantInfo tenantInfo;
    private String tenantIdColumn = "tenant_id";
    private String dialect = "mysql";
    //属性参数信息
    private Properties properties;

    public MultiTenancy(){
    }

    public Object intercept(Invocation invocation) throws Throwable {
        if (getTenantInfo().getTenantId() != null) {
            // 能从cookie或者session获取到tenantId
            mod(invocation);
        }
        return invocation.proceed();
    }

    /**
     * 更改MappedStatement为新的
     * @param invocation
     * @throws Throwable
     */
    public void mod(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        MetaObject metaStatementHandler = SystemMetaObject.forObject(statementHandler);
        Configuration configuration = (Configuration) metaStatementHandler.getValue("delegate.configuration");

        MappedStatement ms = (MappedStatement) metaStatementHandler.getValue("delegate.mappedStatement");
        if (getTenantInfo().doMappedStatementFIlter(ms))
            return;
        SqlCommandType type = ms.getSqlCommandType();

        BoundSql boundSql = statementHandler.getBoundSql();

        String newSQL = addWhere(boundSql.getSql());
        if (newSQL != null) {
            metaStatementHandler.setValue("delegate.boundSql.sql", newSQL);
        }

    }

    /**
     * 添加租户id条件
     * @param sql
     * @return
     * @throws JSQLParserException
     */
    private String addWhere(String sql) throws Exception {
        Statement stmt = CCJSqlParserUtil.parse(sql);
        if (stmt instanceof Insert) {
            //获得Insert对象
            Insert insert = (Insert) stmt;
            if (getTenantInfo().doTableFilter(insert.getTable().getName()))
                return insert.toString();
            boolean bExist = false;
            for(Column col:insert.getColumns()) {
                if(col.getColumnName().equalsIgnoreCase(getTenantIdColumn()))
                    bExist = true;
            }
            if (!bExist) {
                insert.getColumns().add(new Column(getTenantIdColumn()));

                if (insert.getItemsList() instanceof MultiExpressionList){
                    for (ExpressionList expressionList : ((MultiExpressionList) insert.getItemsList()).getExprList()) {
                        addTenantValue(expressionList);
                    }
                }else {
                    addTenantValue(((ExpressionList) insert.getItemsList()));
                }
            }
            return insert.toString();
        }

        if (stmt instanceof Delete) {
            //获得Delete对象
            Delete deleteStatement = (Delete) stmt;
            if (getTenantInfo().doTableFilter(deleteStatement.getTable().getName()))
                return deleteStatement.toString();
            Expression where = deleteStatement.getWhere();
            if (where instanceof BinaryExpression) {
                EqualsTo equalsTo = new EqualsTo();
                equalsTo.setLeftExpression(new Column(getTenantIdColumn()));
                equalsTo.setRightExpression(new StringValue(getTenantInfo().getTenantId()));
                AndExpression andExpression = new AndExpression(equalsTo, where);
                deleteStatement.setWhere(andExpression);
            }
            return deleteStatement.toString();
        }

        if (stmt instanceof Update) {
            //获得Update对象
            Update updateStatement = (Update) stmt;
            TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
            List<String> tableList = tablesNamesFinder.getTableList(stmt);
            if (tableList.size() == 0) {
                return updateStatement.toString();
            }
            //获得where条件表达式
            Expression where = updateStatement.getWhere();
            for (String table: tableList) {
                if (getTenantInfo().doTableFilter(table))
                    continue;
                if (where != null) {
                    AndExpression andE = addAndExpression(stmt, table, updateStatement.getWhere());
                    if (andE != null) {
                        updateStatement.setWhere(andE);
                    }
                } else {
                    throw new Exception("update语句不能没有where条件:" + sql + Arrays.toString(Thread.currentThread().getStackTrace()));
                }
            }
            return updateStatement.toString();
        }

        if (stmt instanceof Select) {
            //获得Select对象
            Select select = (Select) stmt;
            PlainSelect ps = (PlainSelect) select.getSelectBody();
            TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
            List<String> tableList = tablesNamesFinder.getTableList(select);

            if (tableList.size() == 0) {
                return select.toString();
            }
            for (String table : tableList) {
                if (getTenantInfo().doTableFilter(table))
                    continue;
                if (ps.getWhere() != null) {
                    AndExpression where = addAndExpression(stmt, table, ps.getWhere());
                    // form 和 join 中加载的表
                    if (where != null) {
                        ps.setWhere(where);
                    } else {
                        //子查询中的表
                        findSubSelect(stmt, ps.getWhere());
                    }
                } else {
                    ps.setWhere(addEqualsTo(stmt, table));
                }
            }
            return select.toString();
        }

        if (stmt instanceof CreateTable) {
            CreateTable createTable = (CreateTable) stmt;
            ColumnDefinition columnDefinition = new ColumnDefinition();
            columnDefinition.setColumnName(getTenantIdColumn());
            ColDataType colDataType = new ColDataType();
            colDataType.setDataType("varchar(50)");
            columnDefinition.setColDataType(colDataType);
            createTable.getColumnDefinitions().add(columnDefinition);
            return createTable.toString();
        }
            throw new RuntimeException("非法sql语句,请检查"+sql);
    }

    /**
     * 生成代理对象
     *
     * @param target
     * @return
     */
    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    /**
     * 获取配置信息
     * {dialect=mysql, tenantInfo=org.xue.test.TenantInfoImpl, tenantIdColumn=tenant_id}
     * @param properties
     */
    public void setProperties(Properties properties) {
        this.properties=properties;
        try {
            Class onwClass=Class.forName(this.properties.getProperty("tenantInfo"));
            this.tenantInfo = (TenantInfo)onwClass.newInstance();
            this.tenantIdColumn = this.properties.getProperty("tenantIdColumn");
            this.dialect = this.properties.getProperty("dialect");
        } catch (ClassNotFoundException e) {
            throw new RuntimeException(e);
        } catch (InstantiationException e) {
            throw new RuntimeException(e);
        } catch (IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 插入数据 添加租户id
     * @param expressionList
     * @throws Exception
     */
    private void addTenantValue(ExpressionList expressionList) throws Exception {
        expressionList.getExpressions().add(new StringValue(getTenantInfo().getTenantId()));
    }

    /**
     * 多条件情况下，使用AndExpression给where条件加上tenantid条件
     *
     * @param table
     * @param where
     * @return
     * @throws Exception
     */
    public AndExpression addAndExpression(Statement stmt, String table, Expression where) throws Exception {
        EqualsTo equalsTo = addEqualsTo(stmt, table);
        if (equalsTo != null) {
            return new AndExpression(equalsTo, where);
        } else {
            return null;
        }
    }

    /**
     * 创建一个 EqualsTo相同判断 条件
     *
     * @param stmt  查询对象
     * @param table 表名
     * @return “A=B” 单个where条件表达式
     * @throws Exception
     */
    public EqualsTo addEqualsTo(Statement stmt, String table) throws Exception {
        EqualsTo equalsTo = new EqualsTo();
        String aliasName;
        aliasName = getTableAlias(stmt, table);
        if (aliasName != null) {
            equalsTo.setLeftExpression(new Column(aliasName + '.' + getTenantIdColumn()));
            equalsTo.setRightExpression(new StringValue(getTenantInfo().getTenantId()));
            return equalsTo;
        } else {
            return null;
        }
    }

    /**
     * 获取sql送指定表的别名你，没有别名则返回原表名 如果表名不存在返回null
     * 【仅查询from和join 不含 IN 子查询中的表 】
     *
     * @param stmt
     * @param tableName
     * @return
     */
    public String getTableAlias(Statement stmt, String tableName) {
        String alias = null;

        // 插入不做处理
        if (stmt instanceof Insert) {
            return tableName;
        }

        if (stmt instanceof Delete) {
            //获得Delete对象
            Delete deleteStatement = (Delete) stmt;

            if ((deleteStatement.getTable()).getName().equalsIgnoreCase(tableName)) {
                alias = deleteStatement.getTable().getAlias() != null ? deleteStatement.getTable().getAlias().getName() : tableName;
            }
        }

        if (stmt instanceof Update) {
            //获得Update对象
            Update updateStatement = (Update) stmt;

            if ((updateStatement.getTables().get(0)).getName().equalsIgnoreCase(tableName)) {
                alias = updateStatement.getTables().get(0).getAlias() != null ? updateStatement.getTables().get(0).getAlias().getName() : tableName;
            }
        }

        if (stmt instanceof Select) {
            Select select = (Select) stmt;

            PlainSelect ps = (PlainSelect) select.getSelectBody();
            FromItem item = ps.getFromItem();

            // 判断主表的别名
            if (item instanceof Table) {
                if (((Table) item).getName().equalsIgnoreCase(tableName)) {
                    alias = item.getAlias() != null ? item.getAlias().getName() : tableName;
                }
            } else {

            }
        }
        return alias;
    }

    /**
     * 针对子查询中的表别名查询
     *
     * @param subSelect
     * @param tableName
     * @return
     */
    public String getTableAlias(SubSelect subSelect, String tableName) {
        PlainSelect ps = (PlainSelect) subSelect.getSelectBody();
        // 判断主表的别名
        String alias = null;
        FromItem item = ps.getFromItem();
        if (item instanceof Table) {
            if (((Table) item).getName().equalsIgnoreCase(tableName)){
                if (item.getAlias() != null) {
                    alias = item.getAlias().getName();
                } else {
                    alias = tableName;
                }
            }
        }
        return alias;
    }

    /**
     * 递归处理 子查询中的tenantid-where
     *
     * @param stmt  sql查询对象
     * @param where 当前sql的where条件 where为AndExpression或OrExpression的实例，解析其中的rightExpression，然后检查leftExpression是否为空，
     *              不为空则是AndExpression或OrExpression，再次解析其中的rightExpression
     *              注意tenantid-where是加在子查询上的
     */
    void findSubSelect(Statement stmt, Expression where) throws Exception {

        // and 表达式
        if (where instanceof AndExpression) {
            AndExpression andExpression = (AndExpression) where;
            if (andExpression.getRightExpression() instanceof SubSelect) {
                SubSelect subSelect = (SubSelect) andExpression.getRightExpression();
                doSelect(stmt, subSelect);
            }
            if (andExpression.getLeftExpression() != null) {
                findSubSelect(stmt, andExpression.getLeftExpression());
            }
        } else if (where instanceof OrExpression) {
            //  or表达式
            OrExpression orExpression = (OrExpression) where;
            if (orExpression.getRightExpression() instanceof SubSelect) {
                SubSelect subSelect = (SubSelect) orExpression.getRightExpression();
                doSelect(stmt, subSelect);
            }
            if (orExpression.getLeftExpression() != null) {
                findSubSelect(stmt, orExpression.getLeftExpression());
            }
        }
    }

    /**
     * 处理select 和 subSelect
     *
     * @param stmt   查询对象
     * @param select
     * @return
     * @throws Exception
     */
    Expression doSelect(Statement stmt, Expression select) throws Exception {
        PlainSelect ps = null;
        boolean hasSubSelect = false;

        if (select instanceof SubSelect) {
            ps = (PlainSelect) ((SubSelect) select).getSelectBody();
        }
        if (select instanceof Select) {
            ps = (PlainSelect) ((Select) select).getSelectBody();
        }

        TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
        List<String> tableList = tablesNamesFinder.getTableList(select);
        if (tableList.size() == 0) {
            return select;
        }
        for (String table : tableList) {
            // sql 包含 where 条件的情况 使用 addAndExpression 连接 已有的条件和新条件
            if (ps.getWhere() == null) {
                AndExpression where = addAndExpression(stmt, table, ps.getWhere());
                // form 和 join 中加载的表
                if (where != null) {
                    ps.setWhere(where);
                } else {
                    // 如果在Statement中不存在这个表名，则存在于子查询中
                    hasSubSelect = true;
                }
            } else {
                // sql 不含 where条件 新建一个EqualsTo设置为where条件
                EqualsTo equalsTo = addEqualsTo(stmt, table);
                ps.setWhere(equalsTo);
            }
        }

        if (hasSubSelect) {
            //子查询中的表
            findSubSelect(stmt, ps.getWhere());
        }
        return select;
    }

    private void processSelectBody(SelectBody selectBody) {
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
        } else if (selectBody instanceof WithItem) {
            WithItem withItem = (WithItem) selectBody;
            if (withItem.getSelectBody() != null) {
                processSelectBody(withItem.getSelectBody());
            }
        } else {
            SetOperationList operationList = (SetOperationList) selectBody;
            if (operationList.getOperations() != null && operationList.getSelects().size() > 0) {
                List<SelectBody> plainSelects = operationList.getSelects();
                for (SelectBody sub: plainSelects) {
                    processSelectBody(sub);
                }
            }
            if (!orderByHashParameters(operationList.getOrderByElements())) {
                operationList.setOrderByElements(null);
            }
        }
    }

    private void processPlainSelect(PlainSelect plainSelect) {
        if (!orderByHashParameters(plainSelect.getOrderByElements())) {
            plainSelect.setOrderByElements(null);
        }
        if (plainSelect.getFromItem() != null) {
            processFromItem(plainSelect.getFromItem());
        }
        if (plainSelect.getJoins() != null && plainSelect.getJoins().size() > 0) {
            List<Join> joins = plainSelect.getJoins();
            for (Join join : joins) {
                if (join.getRightItem() != null) {
                    processFromItem(join.getRightItem());
                }
            }
        }
    }

    public void processFromItem(FromItem fromItem) {
        if (fromItem instanceof SubJoin) {
            SubJoin subJoin = (SubJoin) fromItem;
            if (subJoin.getJoinList() != null) {
                for(Join tjoin: subJoin.getJoinList()) {
                    if (tjoin.getRightItem() != null) {
                        processFromItem(tjoin.getRightItem());
                    }
                }
            }
        } else if (fromItem instanceof SubSelect) {
            SubSelect subSelect = (SubSelect) fromItem;
            if (subSelect.getSelectBody() != null) {
                processSelectBody(subSelect.getSelectBody());
            }
        } else if (fromItem instanceof ValuesList) {

        } else if (fromItem instanceof LateralSubSelect) {
            LateralSubSelect lateralSubSelect = (LateralSubSelect) fromItem;
            if (lateralSubSelect.getSubSelect() != null) {
                SubSelect subSelect = (SubSelect) (lateralSubSelect.getSubSelect());
                if (subSelect.getSelectBody() != null) {
                    processSelectBody(subSelect.getSelectBody());
                }
            }
        }
        //Table时不用处理
    }

    public boolean orderByHashParameters(List<OrderByElement> orderByElements) {
        if (orderByElements == null) {
            return false;
        }
        for (OrderByElement orderByElement : orderByElements) {
            if (orderByElement.toString().toUpperCase().contains("?")) {
                return true;
            }
        }
        return false;
    }

    public TenantInfo getTenantInfo() {
        return tenantInfo;
    }

    public void setTenantInfo(TenantInfo tenantInfo) {
        this.tenantInfo = tenantInfo;
    }

    public String getTenantIdColumn() {
        return tenantIdColumn;
    }

    public void setTenantIdColumn(String tenantIdColumn) {
        this.tenantIdColumn = tenantIdColumn;
    }

    public String getDialect() {
        return dialect;
    }

    public void setDialect(String dialect) {
        this.dialect = dialect;
    }


}
