package cn.skcks.matrix.v2.casbin.adapter;

import cn.hutool.core.util.IdUtil;
import dev.failsafe.Failsafe;
import org.apache.commons.collections4.CollectionUtils;
import org.casbin.adapter.JDBCAdapter;
import org.casbin.jcasbin.model.Assertion;
import org.casbin.jcasbin.model.Model;

import javax.sql.DataSource;
import java.sql.*;
import java.util.List;
import java.util.Map;

public class CustomJDBCAdapter extends JDBCAdapter {
    static class CasbinRule {
        int id; // Fields reserved for compatibility with other adapters, and the primary key is automatically incremented.
        String ptype;
        String v0;
        String v1;
        String v2;
        String v3;
        String v4;
        String v5;

        public String[] toStringArray() {
            return new String[]{ptype, v0, v1, v2, v3, v4, v5};
        }
    }

    private final int batchSize = 1000;

    public CustomJDBCAdapter(DataSource dataSource, boolean removePolicyFailed, String tableName, boolean autoCreateTable) throws Exception {
        super(dataSource, removePolicyFailed, tableName, autoCreateTable);
    }

    @Override
    protected void migrate() throws SQLException {
        Statement stmt = conn.createStatement();
        String sql = renderActualSql("CREATE TABLE IF NOT EXISTS casbin_rule(id bigint NOT NULL PRIMARY KEY auto_increment, ptype VARCHAR(100) NOT NULL, v0 VARCHAR(100), v1 VARCHAR(100), v2 VARCHAR(100), v3 VARCHAR(100), v4 VARCHAR(100), v5 VARCHAR(100))");
        String productName = conn.getMetaData().getDatabaseProductName();

        switch (productName) {
            case "MySQL":
                String hasTableSql = renderActualSql("SHOW TABLES LIKE 'casbin_rule';");
                ResultSet rs = stmt.executeQuery(hasTableSql);
                if (rs.next()) {
                    return;
                }
                break;
        }

        stmt.executeUpdate(sql);
    }

    private CasbinRule savePolicyLine(String ptype, List<String> rule) {
        CasbinRule line = new CasbinRule();

        line.ptype = ptype;
        if (rule.size() > 0) {
            line.v0 = rule.get(0);
        }
        if (rule.size() > 1) {
            line.v1 = rule.get(1);
        }
        if (rule.size() > 2) {
            line.v2 = rule.get(2);
        }
        if (rule.size() > 3) {
            line.v3 = rule.get(3);
        }
        if (rule.size() > 4) {
            line.v4 = rule.get(4);
        }
        if (rule.size() > 5) {
            line.v5 = rule.get(5);
        }

        return line;
    }

    @Override
    public void addPolicies(String sec, String ptype, List<List<String>> rules) {
        if (CollectionUtils.isEmpty(rules)) {
            return;
        }

        try (Connection conn = dataSource.getConnection()) {
            String sql = renderActualSql("INSERT INTO casbin_rule (ptype,v0,v1,v2,v3,v4,v5,id) VALUES(?,?,?,?,?,?,?,?)");

            Failsafe.with(retryPolicy).run(ctx -> {
                if (ctx.isRetry()) {
                    retry(ctx);
                }
                conn.setAutoCommit(false);
                int count = 0;
                try (PreparedStatement ps = conn.prepareStatement(sql)) {
                    for (List<String> rule : rules) {
                        CasbinRule line = savePolicyLine(ptype, rule);

                        ps.setString(1, line.ptype);
                        ps.setString(2, line.v0);
                        ps.setString(3, line.v1);
                        ps.setString(4, line.v2);
                        ps.setString(5, line.v3);
                        ps.setString(6, line.v4);
                        ps.setString(7, line.v5);
                        ps.setString(8, IdUtil.getSnowflakeNextIdStr());
                        ps.addBatch();
                        if (++count == batchSize) {
                            count = 0;
                            ps.executeBatch();
                            ps.clearBatch();
                        }
                    }
                    if (count != 0) {
                        ps.executeBatch();
                    }
                    conn.commit();
                } catch (SQLException e) {
                    conn.rollback();

                    e.printStackTrace();
                    throw e;
                } finally {
                    conn.setAutoCommit(true);
                }
            });
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }

    }

    @Override
    public void savePolicy(Model model) {
        String cleanSql = renderActualSql("delete from casbin_rule");
        String addSql = renderActualSql("INSERT INTO casbin_rule (ptype,v0,v1,v2,v3,v4,v5,id) VALUES(?,?,?,?,?,?,?,?)");

        try (Connection conn = dataSource.getConnection()) {
            Failsafe.with(retryPolicy).run(ctx -> {
                if (ctx.isRetry()) {
                    retry(ctx);
                }
                conn.setAutoCommit(false);

                int count = 0;

                try (Statement statement = conn.createStatement(); PreparedStatement ps = conn.prepareStatement(addSql)) {
                    statement.execute(cleanSql);
                    for (Map.Entry<String, Assertion> entry : model.model.get("p").entrySet()) {
                        String ptype = entry.getKey();
                        Assertion ast = entry.getValue();

                        for (List<String> rule : ast.policy) {
                            CasbinRule line = savePolicyLine(ptype, rule);

                            ps.setString(1, line.ptype);
                            ps.setString(2, line.v0);
                            ps.setString(3, line.v1);
                            ps.setString(4, line.v2);
                            ps.setString(5, line.v3);
                            ps.setString(6, line.v4);
                            ps.setString(7, line.v5);
                            ps.setString(8, IdUtil.getSnowflakeNextIdStr());
                            ps.addBatch();
                            if (++count == batchSize) {
                                count = 0;
                                ps.executeBatch();
                                ps.clearBatch();
                            }
                        }
                    }

                    for (Map.Entry<String, Assertion> entry : model.model.get("g").entrySet()) {
                        String ptype = entry.getKey();
                        Assertion ast = entry.getValue();

                        for (List<String> rule : ast.policy) {
                            CasbinRule line = savePolicyLine(ptype, rule);

                            ps.setString(1, line.ptype);
                            ps.setString(2, line.v0);
                            ps.setString(3, line.v1);
                            ps.setString(4, line.v2);
                            ps.setString(5, line.v3);
                            ps.setString(6, line.v4);
                            ps.setString(7, line.v5);
                            ps.setString(8, IdUtil.getSnowflakeNextIdStr());

                            ps.addBatch();
                            if (++count == batchSize) {
                                count = 0;
                                ps.executeBatch();
                                ps.clearBatch();
                            }
                        }
                    }

                    if (count != 0) {
                        ps.executeBatch();
                    }

                    conn.commit();
                } catch (SQLException e) {
                    conn.rollback();

                    e.printStackTrace();
                    throw e;
                } finally {
                    conn.setAutoCommit(true);
                }
            });
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }
}
