#![allow(unused_imports)]

mod common;

use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::{Result, anyhow, bail};
use bytes::Bytes;
use kafkit_client::{
    AdminConfig, AlterConfigOp, AutoOffsetReset, ConfigResource, ConsumerConfig, ConsumerError,
    IsolationLevel, KafkaAdmin, KafkaClient, KafkaConsumer, KafkaMessage, KafkaProducer,
    NewPartitions, NewTopic, ProduceRecord, ProducerCompression, ProducerConfig, RecordHeader,
    SubscriptionPattern, TopicPartition, TopicPartitionTimestamp,
};

use crate::common::{
    KafkaHarness, collect_values, expected_assignment, poll_describe_until, poll_until,
    poll_until_assignment, poll_until_consumer_error, poll_until_with_admin, poll_until_with_api,
    unique_group, unique_topic,
};

#[tokio::test(flavor = "multi_thread")]
async fn producer_partitioning_is_deterministic_for_keys_and_sticky_for_unkeyed_records()
-> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("producer-partitioner");
    let admin = KafkaAdmin::connect(AdminConfig::new(kafka.bootstrap_server())).await?;
    admin
        .create_topics([NewTopic::new(topic.clone(), 3, 1)])
        .await?;

    let keyed_producer =
        KafkaProducer::connect(ProducerConfig::new(kafka.bootstrap_server())).await?;
    let (keyed_one, keyed_two, keyed_three) = tokio::try_join!(
        keyed_producer.send(
            ProduceRecord::without_partition(topic.clone(), Bytes::from_static(b"keyed-one"))
                .with_key(Bytes::from_static(b"stable-key")),
        ),
        keyed_producer.send(
            ProduceRecord::without_partition(topic.clone(), Bytes::from_static(b"keyed-two"))
                .with_key(Bytes::from_static(b"stable-key")),
        ),
        keyed_producer.send(
            ProduceRecord::without_partition(topic.clone(), Bytes::from_static(b"keyed-three"))
                .with_key(Bytes::from_static(b"stable-key")),
        ),
    )?;
    assert_eq!(keyed_one.partition, keyed_two.partition);
    assert_eq!(keyed_two.partition, keyed_three.partition);
    keyed_producer.shutdown().await?;

    let sticky_producer = KafkaProducer::connect(
        ProducerConfig::new(kafka.bootstrap_server())
            .with_batch_size(64 * 1024)
            .with_linger(Duration::from_millis(200)),
    )
    .await?;
    let (sticky_one, sticky_two, sticky_three) = tokio::try_join!(
        sticky_producer.send(ProduceRecord::without_partition(
            topic.clone(),
            Bytes::from_static(b"sticky-one"),
        )),
        sticky_producer.send(ProduceRecord::without_partition(
            topic.clone(),
            Bytes::from_static(b"sticky-two"),
        )),
        sticky_producer.send(ProduceRecord::without_partition(
            topic.clone(),
            Bytes::from_static(b"sticky-three"),
        )),
    )?;
    assert_eq!(sticky_one.partition, sticky_two.partition);
    assert_eq!(sticky_two.partition, sticky_three.partition);
    sticky_producer.shutdown().await?;
    Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn producer_round_trip_with_compression_and_modern_groups() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("roundtrip");
    kafka.create_topic(&topic).await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), unique_group("roundtrip"))
            .with_auto_offset_reset(AutoOffsetReset::Earliest),
    )
    .await?;
    consumer.subscribe(vec![topic.clone()]).await?;

    let producer = KafkaProducer::connect(
        ProducerConfig::new(kafka.bootstrap_server())
            .with_client_id("integration-producer")
            .with_compression(ProducerCompression::Gzip),
    )
    .await?;

    producer
        .send(ProduceRecord::new(
            topic.clone(),
            0,
            Bytes::from_static(b"hello from integration test"),
        ))
        .await?;

    let records = poll_until(
        &consumer,
        |records| {
            records.iter().any(|record| {
                record.topic == topic
                    && record.partition == 0
                    && record.value.as_deref() == Some(&b"hello from integration test"[..])
            })
        },
        Duration::from_secs(20),
    )
    .await?;

    assert!(records.iter().any(|record| {
        record.topic == topic
            && record.partition == 0
            && record.value.as_deref() == Some(&b"hello from integration test"[..])
    }));

    consumer.shutdown().await?;
    producer.shutdown().await?;
    Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn producer_and_consumer_expose_headers_and_timestamps() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("headers-timestamps");
    let produced_timestamp = 1_700_000_123_456i64;
    kafka.create_topic(&topic).await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), unique_group("headers-timestamps"))
            .with_auto_offset_reset(AutoOffsetReset::Earliest),
    )
    .await?;
    consumer.subscribe(vec![topic.clone()]).await?;

    let producer = KafkaProducer::connect(ProducerConfig::new(kafka.bootstrap_server())).await?;
    producer
        .send(
            ProduceRecord::new(topic.clone(), 0, Bytes::from_static(b"payload"))
                .with_key(Bytes::from_static(b"record-key"))
                .with_timestamp(produced_timestamp)
                .with_headers([
                    RecordHeader::new("trace-id", Bytes::from_static(b"trace-123")),
                    RecordHeader::null("nullable"),
                ]),
        )
        .await?;

    let records = poll_until(
        &consumer,
        |records| {
            records.iter().any(|record| {
                record.topic == topic && record.value.as_deref() == Some(&b"payload"[..])
            })
        },
        Duration::from_secs(20),
    )
    .await?;

    let record = records
        .iter()
        .find(|record| record.topic == topic && record.value.as_deref() == Some(&b"payload"[..]))
        .ok_or_else(|| anyhow!("expected produced record in fetched batch"))?;
    assert_eq!(record.key.as_deref(), Some(&b"record-key"[..]));
    assert_eq!(record.timestamp, produced_timestamp);
    assert_eq!(record.headers.len(), 2);
    assert_eq!(record.headers[0].key, "trace-id");
    assert_eq!(record.headers[0].value.as_deref(), Some(&b"trace-123"[..]));
    assert_eq!(record.headers[1].key, "nullable");
    assert!(record.headers[1].value.is_none());

    consumer.shutdown().await?;
    producer.shutdown().await?;
    Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn producer_can_send_tombstones() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("tombstones");
    kafka.create_topic(&topic).await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), unique_group("tombstones"))
            .with_auto_offset_reset(AutoOffsetReset::Earliest),
    )
    .await?;
    consumer.subscribe(vec![topic.clone()]).await?;

    let producer = KafkaProducer::connect(ProducerConfig::new(kafka.bootstrap_server())).await?;
    producer
        .send(
            ProduceRecord::tombstone(topic.clone(), 0)
                .with_key(Bytes::from_static(b"deleted-key"))
                .with_header(RecordHeader::null("nullable")),
        )
        .await?;

    let records = poll_until(
        &consumer,
        |records| {
            records.iter().any(|record| {
                record.topic == topic
                    && record.key.as_deref() == Some(&b"deleted-key"[..])
                    && record.value.is_none()
            })
        },
        Duration::from_secs(20),
    )
    .await?;

    let record = records
        .iter()
        .find(|record| {
            record.topic == topic
                && record.key.as_deref() == Some(&b"deleted-key"[..])
                && record.value.is_none()
        })
        .ok_or_else(|| anyhow!("expected tombstone record in fetched batch"))?;
    assert_eq!(record.headers.len(), 1);
    assert_eq!(record.headers[0].key, "nullable");
    assert!(record.headers[0].value.is_none());

    consumer.shutdown().await?;
    producer.shutdown().await?;
    Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn producer_batches_records_for_one_partition() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("batched-produce");
    kafka.create_topic(&topic).await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), unique_group("batched-produce"))
            .with_auto_offset_reset(AutoOffsetReset::Earliest),
    )
    .await?;
    consumer.subscribe(vec![topic.clone()]).await?;

    let producer = KafkaProducer::connect(
        ProducerConfig::new(kafka.bootstrap_server())
            .with_batch_size(64 * 1024)
            .with_linger(Duration::from_millis(100)),
    )
    .await?;

    let (ack1, ack2, ack3) = tokio::try_join!(
        producer.send(ProduceRecord::new(
            topic.clone(),
            0,
            Bytes::from_static(b"one")
        )),
        producer.send(ProduceRecord::new(
            topic.clone(),
            0,
            Bytes::from_static(b"two")
        )),
        producer.send(ProduceRecord::new(
            topic.clone(),
            0,
            Bytes::from_static(b"three")
        )),
    )?;

    let mut acks = [ack1, ack2, ack3];
    acks.sort_by_key(|ack| ack.offset);
    assert!(
        acks.iter()
            .all(|ack| ack.base_offset == acks[0].base_offset)
    );
    assert_eq!(acks[0].offset, acks[0].base_offset);
    assert_eq!(acks[1].offset, acks[0].base_offset + 1);
    assert_eq!(acks[2].offset, acks[0].base_offset + 2);

    let records = poll_until(
        &consumer,
        |records| records.len() >= 3,
        Duration::from_secs(20),
    )
    .await?;
    assert!(
        records
            .iter()
            .any(|record| record.value.as_deref() == Some(&b"one"[..]))
    );
    assert!(
        records
            .iter()
            .any(|record| record.value.as_deref() == Some(&b"two"[..]))
    );
    assert!(
        records
            .iter()
            .any(|record| record.value.as_deref() == Some(&b"three"[..]))
    );

    consumer.shutdown().await?;
    producer.shutdown().await?;
    Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn producer_flush_waits_for_lingered_batches() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("flush");
    kafka.create_topic(&topic).await?;

    let producer = Arc::new(
        KafkaProducer::connect(
            ProducerConfig::new(kafka.bootstrap_server())
                .with_batch_size(64 * 1024)
                .with_linger(Duration::from_secs(5)),
        )
        .await?,
    );

    let send_one = {
        let producer = producer.clone();
        let topic = topic.clone();
        tokio::spawn(async move {
            producer
                .send(ProduceRecord::new(
                    topic,
                    0,
                    Bytes::from_static(b"flush-one"),
                ))
                .await
        })
    };
    let send_two = {
        let producer = producer.clone();
        let topic = topic.clone();
        tokio::spawn(async move {
            producer
                .send(ProduceRecord::new(
                    topic,
                    0,
                    Bytes::from_static(b"flush-two"),
                ))
                .await
        })
    };

    tokio::time::sleep(Duration::from_millis(100)).await;
    producer.flush().await?;

    let ack_one = tokio::time::timeout(Duration::from_secs(5), send_one).await???;
    let ack_two = tokio::time::timeout(Duration::from_secs(5), send_two).await???;

    let mut acks = [ack_one, ack_two];
    acks.sort_by_key(|ack| ack.offset);
    assert_eq!(acks[0].offset, acks[0].base_offset);
    assert_eq!(acks[1].offset, acks[0].base_offset + 1);

    let values = kafka
        .consume_values(&topic, "read_uncommitted", 2, 10_000)
        .await?;
    assert_eq!(values, vec!["flush-one".to_owned(), "flush-two".to_owned()]);

    Arc::try_unwrap(producer)
        .map_err(|_| anyhow!("producer still had outstanding references"))?
        .shutdown()
        .await?;
    Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn producer_delivery_timeout_fails_lingered_batches() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("delivery-timeout");
    kafka.create_topic(&topic).await?;

    let producer = KafkaProducer::connect(
        ProducerConfig::new(kafka.bootstrap_server())
            .with_linger(Duration::from_secs(5))
            .with_delivery_timeout(Duration::from_millis(250)),
    )
    .await?;

    let result = tokio::time::timeout(
        Duration::from_secs(5),
        producer.send(ProduceRecord::new(
            topic.clone(),
            0,
            Bytes::from_static(b"timed-out"),
        )),
    )
    .await?;
    let error =
        result.expect_err("linger should keep the batch unsent until delivery timeout expires");
    assert!(error.to_string().contains("delivery timed out"));

    let values = kafka
        .consume_values(&topic, "read_uncommitted", 1, 2_000)
        .await?;
    assert!(values.is_empty());

    producer.shutdown().await?;
    Ok(())
}
