import org.junit.BeforeClass;
import org.junit.Test;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;

import static org.junit.Assert.assertEquals;

public class MariaDBBatchBugTest {

    // Change these for your test environment.
    private static final String HOST = "localhost";
    private static final String DATABASE = "test";
    private static final String USER = "test";
    private static final String PASSWORD = "test";

    private static final String BASE_URL =
            "jdbc:mariadb://" + HOST + "/" + DATABASE;

    private static final String INSERT_SQL = """
            INSERT INTO test_batch
                (value_a, value_b, flag, value_c, created_at)
            VALUES (?, ?, ?, ?, NOW())
            """;

    private static final String SELECT_SQL =
            "SELECT id FROM test_lookup WHERE id=?";

    @BeforeClass
    public static void loadDriver() throws Exception {
        Class.forName("org.mariadb.jdbc.Driver");
    }

    /**
     * No trigger.
     *
     * Expected: succeeds with bulk statements enabled.
     */
    @Test
    public void testBatchWithNoTrigger() throws Exception {
        try (Connection conn = getConnection(true)) {
            recreateSchema(conn, TriggerType.NONE);

            int[] results = executeBatch(conn);

            System.out.println("No trigger:");
            printResults(results);

            assertEquals(2, results.length);
        }
    }

    /**
     * Empty AFTER INSERT trigger.
     *
     * Expected: succeeds with bulk statements enabled.
     */
    @Test
    public void testBatchWithEmptyTrigger() throws Exception {
        try (Connection conn = getConnection(true)) {
            recreateSchema(conn, TriggerType.EMPTY);

            int[] results = executeBatch(conn);

            System.out.println("Empty AFTER INSERT trigger:");
            printResults(results);

            assertEquals(2, results.length);
        }
    }

    /**
     * AFTER INSERT trigger containing a DELETE.
     *
     * Expected: succeeds with bulk statements disabled.
     */
    @Test
    public void testBatchWithDeleteTriggerBulkDisabled() throws Exception {
        try (Connection conn = getConnection(false)) {
            recreateSchema(conn, TriggerType.DELETE);

            int[] results = executeBatch(conn);

            System.out.println(
                    "DELETE trigger, bulk statements disabled:");
            printResults(results);

            assertEquals(2, results.length);
        }
    }

    /**
     * AFTER INSERT trigger containing a DELETE.
     *
     * Expected: succeeds with bulk statements enabled.
     *
     * Currently fails inside MariaDB Connector/J.
     */
    @Test
    public void testBatchWithDeleteTriggerBulkEnabled() throws Exception {
        try (Connection conn = getConnection(true)) {
            recreateSchema(conn, TriggerType.DELETE);

            int[] results = executeBatch(conn);

            System.out.println(
                    "DELETE trigger, bulk statements enabled:");
            printResults(results);

            assertEquals(2, results.length);
        }
    }

    /**
     * AFTER INSERT trigger using a JOIN in the DELETE.
     *
     * Expected: succeeds with bulk statements enabled.
     */
    @Test
    public void testBatchWithJoinDeleteTrigger() throws Exception {
        try (Connection conn = getConnection(true)) {
            recreateSchema(conn, TriggerType.JOIN_DELETE);

            int[] results = executeBatch(conn);

            System.out.println(
                    "JOIN DELETE trigger, bulk statements enabled:");
            printResults(results);

            assertEquals(2, results.length);
        }
    }

    private static Connection getConnection(boolean useBulkStmts)
            throws SQLException {

        return DriverManager.getConnection(
                BASE_URL + "?useBulkStmts=" + useBulkStmts,
                USER,
                PASSWORD);
    }

    private static void recreateSchema(
            Connection conn,
            TriggerType triggerType) throws SQLException {

        try (Statement stmt = conn.createStatement()) {

            // Remove objects from previous test runs.
            stmt.executeUpdate(
                    "DROP TRIGGER IF EXISTS test_batch_after_insert");

            stmt.executeUpdate(
                    "DROP TABLE IF EXISTS test_batch");

            stmt.executeUpdate(
                    "DROP TABLE IF EXISTS test_delete");

            stmt.executeUpdate(
                    "DROP TABLE IF EXISTS test_lookup");

            // Table receiving the batch INSERTs.
            stmt.executeUpdate("""
                    CREATE TABLE test_batch (
                        id INT NOT NULL AUTO_INCREMENT,
                        value_a INT NOT NULL,
                        value_b INT NOT NULL,
                        flag TINYINT(1) NOT NULL,
                        value_c INT NOT NULL,
                        created_at DATETIME NOT NULL,
                        PRIMARY KEY (id)
                    ) ENGINE=InnoDB
                    """);

            // Table used by the DELETE trigger.
            stmt.executeUpdate("""
                    CREATE TABLE test_delete (
                        value_a INT NOT NULL,
                        value_b INT NOT NULL
                    ) ENGINE=InnoDB
                    """);

            // Table used by the SELECT performed between batch additions.
            stmt.executeUpdate("""
                    CREATE TABLE test_lookup (
                        id INT NOT NULL PRIMARY KEY
                    ) ENGINE=InnoDB
                    """);

            stmt.executeUpdate(
                    "INSERT INTO test_lookup (id) VALUES (1001), (1002)");

            // Rows which the DELETE trigger can remove.
            stmt.executeUpdate("""
                    INSERT INTO test_delete (value_a, value_b)
                    VALUES
                        (1738, 1001),
                        (1738, 1002)
                    """);

            createTrigger(stmt, triggerType);
        }
    }

    private static void createTrigger(
            Statement stmt,
            TriggerType triggerType) throws SQLException {

        switch (triggerType) {

            case NONE:
                break;

            case EMPTY:
                stmt.executeUpdate("""
                        CREATE TRIGGER test_batch_after_insert
                        AFTER INSERT ON test_batch
                        FOR EACH ROW
                        BEGIN
                        END
                        """);
                break;

            case DELETE:
                /*
                 * This trigger reproduces the Connector/J problem
                 * when bulk statements are enabled.
                 */
                stmt.executeUpdate("""
                        CREATE TRIGGER test_batch_after_insert
                        AFTER INSERT ON test_batch
                        FOR EACH ROW
                        DELETE FROM test_delete
                        WHERE value_a = NEW.value_a
                          AND value_b = NEW.value_b
                        """);
                break;

            case JOIN_DELETE:
                /*
                 * This performs the same logical DELETE using a JOIN.
                 * This version does not reproduce the problem.
                 */
                stmt.executeUpdate("""
                        CREATE TRIGGER test_batch_after_insert
                        AFTER INSERT ON test_batch
                        FOR EACH ROW
                        DELETE td
                          FROM test_delete AS td
                          JOIN test_lookup AS tl
                            ON tl.id = NEW.value_b
                         WHERE td.value_a = NEW.value_a
                           AND td.value_b = tl.id
                        """);
                break;

            default:
                throw new IllegalArgumentException(
                        "Unknown trigger type: " + triggerType);
        }
    }

    private static int[] executeBatch(Connection conn)
            throws SQLException {

        try (PreparedStatement insert =
                     conn.prepareStatement(INSERT_SQL);
             PreparedStatement lookup =
                     conn.prepareStatement(SELECT_SQL)) {

            // These parameters remain constant for both batch entries.
            insert.setInt(1, 1738);
            insert.setInt(4, 2);

            for (int id : new int[] {1001, 1002}) {

                /*
                 * Execute another PreparedStatement between calls to
                 * addBatch(), matching the pattern in the application.
                 */
                lookup.setInt(1, id);

                try (ResultSet rs = lookup.executeQuery()) {
                    if (rs.next()) {
                        insert.setInt(2, id);
                        insert.setBoolean(3, false);

                        System.out.println("Adding: " + insert);
                        insert.addBatch();
                    }
                }
            }

            return insert.executeBatch();
        }
    }

    private static void printResults(int[] results) {
        System.out.println(
                "Number of results: " + results.length);

        for (int i = 0; i < results.length; i++) {
            System.out.println(
                    "results[" + i + "] = " + results[i]);
        }
    }

    private enum TriggerType {
        NONE,
        EMPTY,
        DELETE,
        JOIN_DELETE
    }
}