package com.bonus.canteen.core.auth.config;

import cn.hutool.core.text.CharSequenceUtil;
import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SetOperationList;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;

public class LeNiuDataPermissionInterceptor extends JsqlParserSupport implements InnerInterceptor {
   private static final Logger log = LoggerFactory.getLogger(LeNiuDataPermissionInterceptor.class);
   private LeNiuDataPermissionHandler dataPermissionHandler;
   private String databaseType;

   public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
      if (!InterceptorIgnoreHelper.willIgnoreDataPermission(ms.getId())) {
         PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
         mpBs.sql(this.parserSingle(mpBs.sql().replaceAll("(?m)^[ \t]*\r?\n", ""), ms.getId()));
      }
   }

   protected void processSelect(Select select, int index, String sql, Object obj) {
      if (select instanceof PlainSelect) {
         this.setWhere((PlainSelect)select, (String)obj);
      } else if (select instanceof SetOperationList) {
         List<Select> selectBodyList = ((SetOperationList)select).getSelects();
         selectBodyList.forEach((s) -> {
            this.setWhere((PlainSelect)s, (String)obj);
         });
      }

   }

   private void setWhere(PlainSelect plainSelect, String whereSegment) {
      Expression sqlSegment = this.dataPermissionHandler.getSqlSegment(plainSelect, whereSegment);
      if (null != sqlSegment) {
         plainSelect.setWhere(sqlSegment);
      }

   }

   public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
      PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh);
      PluginUtils.MPBoundSql mpBs = mpSh.mPBoundSql();
      String sql = mpBs.sql();
      if (CharSequenceUtil.contains(sql, "`")) {
         sql = "kb".equals(this.databaseType) ? sql.replace("`", "\"") : sql;
         mpBs.sql(sql);
      }

   }

   public LeNiuDataPermissionHandler getDataPermissionHandler() {
      return this.dataPermissionHandler;
   }

   public String getDatabaseType() {
      return this.databaseType;
   }

   public void setDataPermissionHandler(final LeNiuDataPermissionHandler dataPermissionHandler) {
      this.dataPermissionHandler = dataPermissionHandler;
   }

   public void setDatabaseType(final String databaseType) {
      this.databaseType = databaseType;
   }

   public LeNiuDataPermissionInterceptor() {
   }

   public LeNiuDataPermissionInterceptor(final LeNiuDataPermissionHandler dataPermissionHandler, final String databaseType) {
      this.dataPermissionHandler = dataPermissionHandler;
      this.databaseType = databaseType;
   }
}
