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 java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;

import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;

/** A small crash lab: one Kafka record, one local H2 table, one consumer group per topic. */
public final class OffsetFailureLab {
    private static final String BOOTSTRAP = "127.0.0.1:19092";

    private static Properties producerProperties() {
        Properties props = new Properties();
        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, BOOTSTRAP);
        props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        props.put(ProducerConfig.ACKS_CONFIG, "all");
        return props;
    }

    private static Properties consumerProperties(String topic) {
        Properties props = new Properties();
        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, BOOTSTRAP);
        props.put(ConsumerConfig.GROUP_ID_CONFIG, topic + "-group");
        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "1");
        props.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, "classic");
        props.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "6000");
        props.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "1000");
        props.put(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "5000");
        return props;
    }

    private static Connection database(String topic) throws SQLException {
        Connection connection = DriverManager.getConnection(
                "jdbc:h2:file:/tmp/kafka-seo-db-" + topic + ";DB_CLOSE_ON_EXIT=FALSE;WRITE_DELAY=0");
        try (Statement statement = connection.createStatement()) {
            statement.execute("CREATE TABLE IF NOT EXISTS processed "
                    + "(event_id VARCHAR(80) PRIMARY KEY, payload VARCHAR(200))");
        }
        return connection;
    }

    private static void seed(String topic) throws Exception {
        try (AdminClient admin = AdminClient.create(producerProperties())) {
            admin.createTopics(List.of(new NewTopic(topic, 1, (short) 1))).all().get();
        }
        try (KafkaProducer<String, String> producer = new KafkaProducer<>(producerProperties())) {
            var metadata = producer.send(new ProducerRecord<>(topic, "event-1", "order-created")).get();
            System.out.println("seeded topic=" + topic + " offset=" + metadata.offset());
        }
    }

    private static void writeBusinessRow(Connection connection, ConsumerRecord<String, String> record)
            throws SQLException {
        connection.setAutoCommit(false);
        try (PreparedStatement statement = connection.prepareStatement(
                "INSERT INTO processed(event_id, payload) VALUES (?, ?)")) {
            statement.setString(1, record.key());
            statement.setString(2, record.value());
            statement.executeUpdate();
            connection.commit();
            System.out.println("database=committed " + record.key());
        } catch (SQLException error) {
            connection.rollback();
            if (!"23505".equals(error.getSQLState())) throw error;
            System.out.println("database=duplicate " + record.key());
        } finally {
            connection.setAutoCommit(true);
        }
        try (Statement statement = connection.createStatement()) {
            statement.execute("CHECKPOINT SYNC");
        }
    }

    private static void consume(String topic, String mode) throws Exception {
        try (Connection connection = database(topic);
             KafkaConsumer<String, String> consumer = new KafkaConsumer<>(consumerProperties(topic))) {
            consumer.subscribe(List.of(topic));
            ConsumerRecord<String, String> record = null;
            long deadline = System.nanoTime() + Duration.ofSeconds(20).toNanos();
            while (record == null && System.nanoTime() < deadline) {
                var batch = consumer.poll(Duration.ofMillis(500));
                if (!batch.isEmpty()) record = batch.iterator().next();
            }
            if (record == null) throw new IllegalStateException("No record received");

            TopicPartition partition = new TopicPartition(record.topic(), record.partition());
            var committed = consumer.committed(Set.of(partition)).get(partition);
            System.out.printf("received offset=%d position=%d committed=%s mode=%s%n",
                    record.offset(), consumer.position(partition),
                    committed == null ? "none" : committed.offset(), mode);

            if ("crash-before-db".equals(mode)) crash(11);
            if ("commit-before-db".equals(mode)) {
                consumer.commitSync(Map.of(partition,
                        new org.apache.kafka.clients.consumer.OffsetAndMetadata(record.offset() + 1)));
                System.out.println("offset=committed before database write");
                crash(12);
            }

            writeBusinessRow(connection, record);
            if ("crash-after-db".equals(mode)) crash(13);
            if ("slow-after-db".equals(mode)) {
                System.out.println("pausing 12 seconds after database commit");
                Thread.sleep(12000);
            }

            consumer.commitSync(Map.of(partition,
                    new org.apache.kafka.clients.consumer.OffsetAndMetadata(record.offset() + 1)));
            System.out.println("offset=committed after database write");
        }
    }

    private static void crash(int exitCode) {
        System.out.flush();
        Runtime.getRuntime().halt(exitCode);
    }

    private static void status(String topic) throws Exception {
        try (Connection connection = database(topic);
             Statement statement = connection.createStatement();
             ResultSet rows = statement.executeQuery("SELECT COUNT(*) FROM processed")) {
            rows.next();
            System.out.println("database_rows=" + rows.getInt(1));
        }
    }

    public static void main(String[] args) throws Exception {
        if (args.length < 2 || !args[1].matches("[a-z0-9-]+")) {
            throw new IllegalArgumentException("Usage: OffsetFailureLab seed|consume|status <topic> [mode]");
        }
        switch (args[0]) {
            case "seed" -> seed(args[1]);
            case "status" -> status(args[1]);
            case "consume" -> {
                String mode = args.length > 2 ? args[2] : "normal";
                if (!Set.of("normal", "crash-before-db", "crash-after-db", "commit-before-db", "slow-after-db").contains(mode)) {
                    throw new IllegalArgumentException("Unknown mode: " + mode);
                }
                consume(args[1], mode);
            }
            default -> throw new IllegalArgumentException("Unknown command: " + args[0]);
        }
    }
}