package example.postgres.audit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.CallableStatement;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Savepoint;
import java.sql.Statement;
import java.sql.Timestamp;
import java.time.Instant;
import java.util.Objects;
import java.util.UUID;

import static java.lang.String.format;

/**
 * Whenever the DB is manipulated through this connection, Postgres receives the context variables used in the 'trigger_audit' triggers beforehand.
 */
public class AuditingConnection extends DelegatingConnection {

    private static final Logger LOGGER = LoggerFactory.getLogger(AuditingConnection.class);

    // TODO: implement the static method `AuditContextProvider.currentUserId()`,
    // returning whatever value you'd like persisted in the columns `inserted_by`/`updated_by`
    // from a thread-local context
    private static final AuditContextProvider AUDIT_CONTEXT_PROVIDER = new AuditContextProvider();

    /**
     * Remembers the audit.AUDIT_USER set for the ongoing transaction, such that we only send SET LOCAL commands when necessary.
     */
    private UUID userCommunicatedToDB;

    public AuditingConnection(Connection wrappedConnection) {
        super(wrappedConnection);
    }

    private void setContext() {
        UUID user = AUDIT_CONTEXT_PROVIDER.currentUserId();

        // The user reported by AUDIT_CONTEXT_PROVIDER may change within the same ongoing transaction,
        // e.g. because the user session is initialized after the first DB statements.
        // Detect changes and propagate the information to the DB iff necessary:
        if (!Objects.equals(user, userCommunicatedToDB)) {
            Timestamp nowTimestamp = Timestamp.from(Instant.now());

            if (isAutoCommit()) {
                // If SET LOCAL is called outside a transaction, it has no effect besides Postgres issuing a warning.
                // BUT WHY THE HELL ARE WE USING AUTO-COMMIT???
                LOGGER.warn("Did not set AUDIT context variables because connection was in AUTO-COMMIT mode.");
                return;
            }

            try (Statement auditUserStmt = createStatementWithoutAuditing();
                 Statement auditTimestampStmt = createStatementWithoutAuditing()) {
                auditUserStmt.execute(format("SET LOCAL audit.AUDIT_USER = '%s'", user.toString()));
                auditTimestampStmt.execute(format("SET LOCAL audit.AUDIT_TIMESTAMP = '%s'", nowTimestamp.toString()));
            } catch (SQLException e) {
                LOGGER.warn("failed to set AUDIT information", e);
            }

            this.userCommunicatedToDB = user;
        }
    }

    private Statement createStatementWithoutAuditing() throws SQLException {
        return super.createStatement();
    }

    private void onTransactionEnd() {
        // Values set with SET LOCAL are purged automatically by Postgres after a COMMIT/ROLLBACK [TO SAVEPOINT].
        this.userCommunicatedToDB = null;
    }

    private boolean isAutoCommit() {
        try {
            return getAutoCommit();
        } catch (SQLException e) {
            LOGGER.warn("failed to read AUTO-COMMIT status", e);
            return false; // so we attempt to set the audit info anyway...
        }
    }

    /*

    DELEGATION WITH AUDITING:
    These methods either manipulate data in the DB or the transaction state of this Connection.
    All other methods don't and therefore use direct delegation.

     */

    @Override
    public Statement createStatement() throws SQLException {
        setContext();
        return super.createStatement();
    }

    @Override
    public PreparedStatement prepareStatement(String sql) throws SQLException {
        setContext();
        return super.prepareStatement(sql);
    }

    @Override
    public CallableStatement prepareCall(String sql) throws SQLException {
        setContext();
        return super.prepareCall(sql);
    }

    @Override
    public Statement createStatement(int resultSetType, int resultSetConcurrency) throws SQLException {
        setContext();
        return super.createStatement(resultSetType, resultSetConcurrency);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
        setContext();
        return super.prepareStatement(sql, resultSetType, resultSetConcurrency);
    }

    @Override
    public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
        setContext();
        return super.prepareCall(sql, resultSetType, resultSetConcurrency);
    }

    @Override
    public void setAutoCommit(boolean autoCommit) throws SQLException {
        super.setAutoCommit(autoCommit);
        onTransactionEnd();
    }

    @Override
    public void commit() throws SQLException {
        super.commit();
        onTransactionEnd();
    }

    @Override
    public void rollback() throws SQLException {
        super.rollback();
        onTransactionEnd();
    }

    @Override
    public Statement createStatement(int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
        setContext();
        return super.createStatement(resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
        setContext();
        return super.prepareStatement(sql, resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
        setContext();
        return super.prepareCall(sql, resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
        setContext();
        return super.prepareStatement(sql, autoGeneratedKeys);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException {
        setContext();
        return super.prepareStatement(sql, columnIndexes);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException {
        setContext();
        return super.prepareStatement(sql, columnNames);
    }

    @Override
    public void rollback(Savepoint savepoint) throws SQLException {
        super.rollback(savepoint);
        // if the SET LOCAL happened before the 'setSavepoint(savepoint)' call, we might have lost our context values:
        onTransactionEnd();
    }

    @Override
    public void setTransactionIsolation(int level) throws SQLException {
        super.setTransactionIsolation(level);
        onTransactionEnd();
    }

}