package example;

import java.sql.BatchUpdateException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Savepoint;
import java.sql.Statement;
import java.util.Arrays;

import static example.Database.execute;
import static example.Database.require;
import static example.Database.scalar;

public final class TransactionLab {
    private TransactionLab() {}
    @FunctionalInterface private interface SqlAction { void run() throws SQLException; }

    private static SQLException expect(String state, SqlAction action) throws SQLException {
        try { action.run(); } catch (SQLException error) {
            require(state.equals(error.getSQLState()), "Expected " + state + ", got " + error.getSQLState());
            return error;
        }
        throw new AssertionError("Expected SQLSTATE " + state);
    }

    private static void first() throws SQLException {
        Database.reset();
        Transfer.debit(125, "first");
        require(scalar("select balance_cents from account where id = 1") == 875, "Balance mismatch");
        require(scalar("select count(*) from ledger where entry_key = 'first' and delta_cents = -125") == 1,
                "Ledger mismatch");
        System.out.println("balanceCents=875 ledgerRows=1 independentRead=true");
    }

    private static void rollback() throws SQLException {
        Database.reset();
        try (Connection seed = Database.open()) {
            execute(seed, "insert into ledger values ('duplicate', 1, 0)");
        }
        expect("23505", () -> Transfer.debit(125, "duplicate"));
        require(scalar("select balance_cents from account where id = 1") == 1000, "Partial debit committed");
        require(scalar("select count(*) from ledger") == 1, "Unexpected ledger row");
        System.out.println("duplicateRejected=true sqlState=23505 debitRolledBack=true");
    }

    private static void autocommit() throws SQLException {
        Database.reset();
        try (Connection connection = Database.open()) {
            execute(connection, "insert into ledger values ('duplicate', 1, 0)");
            require(connection.getAutoCommit(), "Expected default autoCommit");
            expect("23505", () -> Transfer.debitOn(connection, 125, "duplicate"));
        }
        require(scalar("select balance_cents from account where id = 1") == 875, "Expected partial debit");
        require(scalar("select count(*) from ledger") == 1, "Duplicate row should not be added");
        System.out.println("autoCommitPartialDebit=true balanceCents=875");
        Database.reset();
    }

    private static void savepoint() throws SQLException {
        Database.reset();
        try (Connection connection = Database.open()) {
            connection.setAutoCommit(false);
            try {
                execute(connection, "insert into ledger values ('A', 1, 0)");
                Savepoint beforeOptional = connection.setSavepoint("before_optional");
                expect("23505", () -> execute(connection, "insert into ledger values ('A', 1, 0)"));
                expect("25P02", () -> scalar(connection, "select 1"));
                connection.rollback(beforeOptional);
                execute(connection, "insert into ledger values ('B', 1, 0)");
                connection.releaseSavepoint(beforeOptional);
                connection.commit();
            } catch (SQLException | RuntimeException | Error failure) {
                Database.rollback(connection, failure);
                throw failure;
            }
        }
        require(scalar("select count(*) from ledger where entry_key in ('A', 'B')") == 2, "Savepoint lost valid rows");
        System.out.println("failureState=23505 abortedState=25P02 savedRows=2 savepointRecovered=true");
    }

    private static void batch(boolean rewrite) throws SQLException {
        Database.reset();
        try (Connection connection = Database.open(rewrite)) {
            connection.setAutoCommit(false);
            try (PreparedStatement insert = connection.prepareStatement(
                    "insert into import_item(external_key, amount_cents) values (?, ?)")) {
                for (int i = 1; i <= 3; i++) {
                    insert.setString(1, "batch-" + i);
                    insert.setInt(2, i * 100);
                    insert.addBatch();
                }
                long[] counts = insert.executeLargeBatch();
                require(counts.length == 3, "Wrong batch count length");
                for (long count : counts) require(count == 1 || count == Statement.SUCCESS_NO_INFO,
                        "Unexpected successful update count " + count);
                require(scalar(connection, "select count(*) from import_item") == 3, "Own transaction lost rows");
                require(scalar("select count(*) from import_item") == 0, "Batch committed too early");
                connection.commit();
                require(scalar("select count(*) from import_item") == 3, "Commit not visible");
                System.out.println("rewrite=" + rewrite + " counts=" + Arrays.toString(counts)
                        + " invisibleBeforeCommit=true committedRows=3");
            } catch (SQLException | RuntimeException | Error failure) {
                Database.rollback(connection, failure);
                throw failure;
            }
        }
    }

    private static void batchFailure() throws SQLException {
        Database.reset();
        try (Connection connection = Database.open()) {
            connection.setAutoCommit(false);
            try (PreparedStatement insert = connection.prepareStatement(
                    "insert into import_item(external_key, amount_cents) values (?, 100)")) {
                for (String key : new String[]{"red", "green", "red"}) {
                    insert.setString(1, key);
                    insert.addBatch();
                }
                try {
                    insert.executeLargeBatch();
                    throw new AssertionError("Duplicate batch succeeded");
                } catch (BatchUpdateException failure) {
                    require("23505".equals(failure.getSQLState()), "Unexpected batch SQLSTATE " + failure.getSQLState());
                    long[] counts = failure.getLargeUpdateCounts();
                    require(counts.length > 0 && counts.length <= 3, "Unexpected failure count length");
                    require(Arrays.stream(counts).allMatch(n -> n >= 0 || n == Statement.EXECUTE_FAILED
                            || n == Statement.SUCCESS_NO_INFO), "Illegal batch count");
                    connection.rollback();
                    require(scalar("select count(*) from import_item") == 0, "Batch rollback left rows");
                    System.out.println("batchFailed=true sqlState=23505 counts=" + Arrays.toString(counts)
                            + " rowsAfterRollback=0");
                }
            } catch (SQLException | RuntimeException | Error failure) {
                Database.rollback(connection, failure);
                throw failure;
            }
        }
    }

    private static void chunks() throws SQLException {
        Database.reset();
        try (Connection connection = Database.open()) {
            connection.setAutoCommit(false);
            try (PreparedStatement insert = connection.prepareStatement(
                    "insert into import_item(external_key, amount_cents) values (?, 100)")) {
                insert.setString(1, "chunk-first");
                insert.addBatch();
                insert.executeBatch();
                require(scalar(connection, "select count(*) from import_item") == 1, "First chunk not executed");
                require(scalar("select count(*) from import_item") == 0, "First chunk already committed");
                insert.setString(1, "chunk-second");
                insert.addBatch();
                insert.executeBatch();
                require(scalar(connection, "select count(*) from import_item") == 2, "Second chunk not executed");
                require(scalar("select count(*) from import_item") == 0, "Second chunk already committed");
                connection.rollback();
            } catch (SQLException | RuntimeException | Error failure) {
                Database.rollback(connection, failure);
                throw failure;
            }
        }
        require(scalar("select count(*) from import_item") == 0, "Chunks survived rollback");
        System.out.println("twoExecuteBatchCalls=true noEarlyCommit=true rollbackRemovedBoth=true");
    }

    private static void keys() throws SQLException {
        Database.reset();
        try (Connection connection = Database.open()) {
            connection.setAutoCommit(false);
            try (PreparedStatement insert = connection.prepareStatement(
                    "insert into import_item(external_key, amount_cents) values (?, 100)", new String[]{"id"})) {
                insert.setString(1, "generated-key");
                require(insert.executeUpdate() == 1, "Insert count mismatch");
                long id;
                try (ResultSet keys = insert.getGeneratedKeys()) {
                    require(keys.next(), "Missing generated key");
                    id = keys.getLong(1);
                    require(id > 0 && !keys.next(), "Invalid generated key result");
                }
                require(scalar("select count(*) from import_item") == 0, "Generated key implied commit");
                connection.rollback();
                require(scalar("select count(*) from import_item") == 0, "Rolled back generated row exists");
                System.out.println("generatedKeyReturned=true invisibleBeforeCommit=true rolledBackRowAbsent=true");
            } catch (SQLException | RuntimeException | Error failure) {
                Database.rollback(connection, failure);
                throw failure;
            }
        }
    }

    private static void cursor() throws SQLException {
        try (Connection connection = Database.open()) {
            connection.setAutoCommit(false);
            try (PreparedStatement statement = connection.prepareStatement(
                    "select i from generate_series(1, 13) as s(i)",
                    ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) {
                statement.setFetchSize(5);
                try (ResultSet rows = statement.executeQuery()) {
                    long portals = scalar(connection, "select count(*) from pg_cursors where name <> ''"
                            + " and statement like '%generate_series(1, 13)%'");
                    require(portals > 0, "Expected named portal with positive fetch size in transaction");
                    int count = 0;
                    while (rows.next()) { count++; require(rows.getInt(1) == count, "Cursor row order mismatch"); }
                    require(count == 13, "Incomplete cursor read");
                    System.out.println("namedPortalObserved=true fetchSize=5 readRows=13");
                }
                connection.rollback();
            } catch (SQLException | RuntimeException | Error failure) {
                Database.rollback(connection, failure);
                throw failure;
            }
        }
        try (Connection connection = Database.open(); PreparedStatement statement = connection.prepareStatement(
                "select i from generate_series(1, 13) as s(i)")) {
            statement.setFetchSize(5);
            try (ResultSet rows = statement.executeQuery()) {
                require(scalar(connection, "select count(*) from pg_cursors where name <> ''"
                        + " and statement like '%generate_series(1, 13)%'") == 0, "Unexpected cursor in autoCommit");
                int count = 0;
                while (rows.next()) count++;
                require(count == 13, "AutoCommit query incomplete");
                System.out.println("autoCommitCursorFallback=true readRows=13");
            }
        }
    }

    public static void main(String[] args) throws SQLException {
        String mode = args.length == 0 ? "all" : args[0];
        switch (mode) {
            case "first" -> first();
            case "rollback" -> rollback();
            case "autocommit" -> autocommit();
            case "savepoint" -> savepoint();
            case "batch" -> { batch(false); batch(true); }
            case "batch-failure" -> batchFailure();
            case "chunks" -> chunks();
            case "keys" -> keys();
            case "cursor" -> cursor();
            case "all" -> { first(); rollback(); autocommit(); savepoint(); batch(false); batch(true);
                batchFailure(); chunks(); keys(); cursor(); }
            default -> throw new IllegalArgumentException("Unknown mode: " + mode);
        }
    }
}