package com.ankurm.idgen;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import org.hibernate.engine.jdbc.connections.internal.DriverManagerConnectionProvider;

/**
 * A ConnectionProvider that wraps every Connection and PreparedStatement in a JDK proxy and counts what Hibernate asks of JDBC:
 * statements prepared, rows queued with addBatch, executeBatch calls, and single executes. Hibernate's own Statistics API
 * exposes prepareStatementCount but no batch counter, so this is how "did batching happen" is measured here.
 */
public class JdbcCounter extends DriverManagerConnectionProvider {

    public static final AtomicLong PREPARED = new AtomicLong(), INSERT_PREPARED = new AtomicLong(), SEQUENCE_PREPARED = new AtomicLong(),
            ADD_BATCH = new AtomicLong(), EXECUTE_BATCH = new AtomicLong(), EXECUTE_SINGLE = new AtomicLong();

    /** Every non-INSERT statement text prepared since the last reset (Hibernate's StatementInspector does not see the id generators' own JDBC work). */
    public static final java.util.List<String> OTHER_SQL = new java.util.concurrent.CopyOnWriteArrayList<>();

    public static void reset() {
        OTHER_SQL.clear();
        for (AtomicLong a : new AtomicLong[] {PREPARED, INSERT_PREPARED, SEQUENCE_PREPARED, ADD_BATCH, EXECUTE_BATCH, EXECUTE_SINGLE}) a.set(0);
    }

    /** One-line summary used in transcripts. */
    public static String summary() {
        return "prepared=" + PREPARED.get() + " (insert=" + INSERT_PREPARED.get() + ", sequence=" + SEQUENCE_PREPARED.get() + ")"
                + " addBatch=" + ADD_BATCH.get() + " executeBatch=" + EXECUTE_BATCH.get() + " executeSingle=" + EXECUTE_SINGLE.get();
    }

    /** Statements sent to the database: one per executeBatch plus one per single execute. */
    public static long roundTrips() { return EXECUTE_BATCH.get() + EXECUTE_SINGLE.get(); }

    private final Map<Connection, Connection> realOf = Collections.synchronizedMap(new IdentityHashMap<>());

    @Override
    public void closeConnection(Connection conn) throws SQLException {
        Connection real = realOf.remove(conn);
        super.closeConnection(real != null ? real : conn);
    }

    @Override
    public Connection getConnection() throws SQLException {
        Connection real = super.getConnection();
        Connection proxy = (Connection) Proxy.newProxyInstance(getClass().getClassLoader(), new Class<?>[] {Connection.class}, (p, m, a) -> {
            try {
                Object r = m.invoke(real, a);
                if (m.getName().equals("prepareStatement") && a != null && a[0] instanceof String sql) {
                    PREPARED.incrementAndGet();
                    String s = sql.toLowerCase();
                    if (s.startsWith("insert")) INSERT_PREPARED.incrementAndGet(); else OTHER_SQL.add(sql.replaceAll("\s+", " ").trim());
                    if (s.contains("nextval") || s.contains("next value for")) SEQUENCE_PREPARED.incrementAndGet();
                    return wrap((PreparedStatement) r);
                }
                return r;
            } catch (InvocationTargetException e) { throw e.getCause(); }
        });
        realOf.put(proxy, real);
        return proxy;
    }

    private static PreparedStatement wrap(PreparedStatement real) {
        return (PreparedStatement) Proxy.newProxyInstance(JdbcCounter.class.getClassLoader(), new Class<?>[] {PreparedStatement.class}, (p, m, a) -> {
            switch (m.getName()) {
                case "addBatch" -> ADD_BATCH.incrementAndGet();
                case "executeBatch" -> EXECUTE_BATCH.incrementAndGet();
                case "execute", "executeUpdate", "executeQuery" -> EXECUTE_SINGLE.incrementAndGet();
                default -> { }
            }
            try { return m.invoke(real, a); } catch (InvocationTargetException e) { throw e.getCause(); }
        });
    }
}