package example;

import com.zaxxer.hikari.HikariDataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.SQLTransientConnectionException;
import java.sql.Statement;
import java.util.Set;

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

public final class PoolLab {
    private PoolLab() {}

    private static void first() throws SQLException {
        try (HikariDataSource pool = Connections.pool()) {
            FirstQuery.query(pool);
            require(pool.getHikariPoolMXBean().getActiveConnections() == 0, "Connection leaked");
            System.out.println("activeAfterQuery=0");
        }
    }

    private static void reset() throws SQLException {
        try (HikariDataSource pool = Connections.pool()) {
            String before;
            try (Connection connection = pool.getConnection()) {
                before = scalar(connection, "select pg_backend_pid()");
                connection.setReadOnly(true);
                connection.setTransactionIsolation(Connection.TRANSACTION_SERIALIZABLE);
                connection.setAutoCommit(false);
                scalar(connection, "select 1");
            }
            try (Connection connection = pool.getConnection()) {
                require(before.equals(scalar(connection, "select pg_backend_pid()")), "Physical session not reused");
                require(connection.getAutoCommit(), "autoCommit was not reset");
                require(!connection.isReadOnly(), "readOnly was not reset");
                require(connection.getTransactionIsolation() == Connection.TRANSACTION_READ_COMMITTED,
                    "Isolation was not reset");
                System.out.println("sameBackend=true autoCommitReset=true readOnlyReset=true isolationReset=true");
                connection.setAutoCommit(false);
                execute(connection, "insert into pool_marker(id, label) values (9001, 'uncommitted')");
            }
            try (Connection connection = pool.getConnection()) {
                require("0".equals(scalar(connection, "select count(*) from pool_marker where id = 9001")),
                    "Uncommitted insert survived return to pool");
                System.out.println("leftoverTransactionRolledBack=true");
            }
        }
    }

    private static void pollution() throws SQLException {
        try (HikariDataSource pool = Connections.pool()) {
            String baseline;
            String backend;
            try (Connection connection = pool.getConnection()) {
                baseline = scalar(connection, "show application_name");
                backend = scalar(connection, "select pg_backend_pid()");
                execute(connection, "set application_name = 'changed-by-sql'");
            }
            try (Connection connection = pool.getConnection()) {
                require(backend.equals(scalar(connection, "select pg_backend_pid()")), "Physical session not reused");
                require("changed-by-sql".equals(scalar(connection, "show application_name")),
                    "Expected session setting to survive pool return");
                try (PreparedStatement statement = connection.prepareStatement("select set_config('application_name', ?, false)")) {
                    statement.setString(1, baseline);
                    statement.execute();
                }
                require(baseline.equals(scalar(connection, "show application_name")), "Explicit restore failed");
                System.out.println("sqlSessionSettingSurvived=true explicitlyRestored=true");
                connection.setAutoCommit(false);
                execute(connection, "set local application_name = 'transaction-only'");
                require("transaction-only".equals(scalar(connection, "show application_name")), "SET LOCAL did not apply");
                connection.commit();
                require(baseline.equals(scalar(connection, "show application_name")), "SET LOCAL outlived transaction");
                System.out.println("transactionLocalSettingRestored=true");
            }
        }
    }

    private static void exhaustion() throws SQLException {
        try (HikariDataSource pool = Connections.pool()) {
            try (Connection held = pool.getConnection()) {
                require("1".equals(scalar(held, "select 1")), "Held connection unavailable");
                try (Connection unexpected = pool.getConnection()) {
                    throw new AssertionError("Pool gave a second connection: " + unexpected);
                } catch (SQLTransientConnectionException expected) {
                    require(pool.getHikariPoolMXBean().getActiveConnections() == 1, "Unexpected active count");
                    System.out.println("borrowTimedOut=true activeWhileHeld=1");
                }
            }
            try (Connection recovered = pool.getConnection()) {
                require("1".equals(scalar(recovered, "select 1")), "Borrow after release failed");
            }
            require(pool.getHikariPoolMXBean().getActiveConnections() == 0, "Connection still held");
            System.out.println("borrowAfterRelease=true activeAfterClose=0");
        }
    }

    private static void invalid() throws SQLException {
        try (HikariDataSource pool = Connections.pool()) {
            String before;
            try (Connection victim = pool.getConnection()) {
                before = scalar(victim, "select pg_backend_pid()");
                // The normal application role may terminate its own backend, not arbitrary users' sessions.
                try (Connection observer = Connections.independent();
                     PreparedStatement kill = observer.prepareStatement("select pg_terminate_backend(?)")) {
                    kill.setInt(1, Integer.parseInt(before));
                    try (var rows = kill.executeQuery()) {
                        require(rows.next() && rows.getBoolean(1), "Backend termination not accepted");
                    }
                }
                try {
                    scalar(victim, "select 1");
                    throw new AssertionError("Terminated backend was still usable");
                } catch (SQLException expected) {
                    String state = expected.getSQLState();
                    require(Set.of("57P01", "08003", "08006").contains(state), "Unexpected disconnect SQLSTATE: " + state);
                    System.out.println("terminatedBackendRejected=true connectionFailureClassAccepted=true");
                }
            }
            try (Connection recovered = pool.getConnection()) {
                require(!before.equals(scalar(recovered, "select pg_backend_pid()")), "Dead physical session reused");
                require("1".equals(scalar(recovered, "select 1")), "Replacement unavailable");
                System.out.println("replacementBackend=true queryAfterReplacement=true");
            }
        }
    }

    private static void timeout() throws SQLException {
        try (HikariDataSource pool = Connections.pool(); Connection connection = pool.getConnection()) {
            connection.setAutoCommit(false);
            try {
                try (Statement slow = connection.createStatement()) {
                    slow.setQueryTimeout(1);
                    slow.executeQuery("select pg_sleep(3)").close();
                    throw new AssertionError("Slow query unexpectedly completed");
                } catch (SQLException expected) {
                    require("57014".equals(expected.getSQLState()), "Expected query cancellation, got " + expected.getSQLState());
                    System.out.println("queryCanceled=true sqlState=57014");
                }
                try {
                    scalar(connection, "select 1");
                    throw new AssertionError("Failed transaction accepted next query");
                } catch (SQLException expected) {
                    require("25P02".equals(expected.getSQLState()), "Expected failed transaction, got " + expected.getSQLState());
                    System.out.println("transactionNeedsRollback=true sqlState=25P02");
                }
                connection.rollback();
                require("1".equals(scalar(connection, "select 1")), "Query after rollback failed");
                System.out.println("queryAfterRollback=true");
            } finally {
                connection.rollback();
            }
        }
    }

    private static void hold() throws SQLException, InterruptedException {
        try (HikariDataSource pool = Connections.pool(); Connection connection = pool.getConnection()) {
            connection.setAutoCommit(false);
            try {
                String backend = scalar(connection, "select pg_backend_pid()");
                System.out.println("holdingSeconds=45 backendPid=" + backend + " state=idle-in-transaction");
                Thread.sleep(45_000);
            } finally {
                connection.rollback();
            }
        }
        System.out.println("holdFinished=true");
    }

    public static void main(String[] args) throws SQLException, InterruptedException {
        String mode = args.length == 0 ? "all" : args[0];
        switch (mode) {
            case "first" -> first();
            case "reset" -> reset();
            case "pollution" -> pollution();
            case "exhaustion" -> exhaustion();
            case "invalid" -> invalid();
            case "timeout" -> timeout();
            case "hold" -> hold();
            case "all" -> { first(); reset(); pollution(); exhaustion(); invalid(); timeout(); }
            default -> throw new IllegalArgumentException("Mode: first|reset|pollution|exhaustion|invalid|timeout|hold|all");
        }
    }
}