#![allow(unused_imports)]

mod common;

use std::time::Duration;

use anyhow::{Result, anyhow};
use bytes::Bytes;
use kafkit_client::{
    AdminConfig, AutoOffsetReset, ConsumerConfig, KafkaAdmin, KafkaConsumer, KafkaProducer,
    NewTopic, ProduceRecord, ProducerCompression, ProducerConfig, RecordHeader,
};

use crate::common::{
    KafkaHarness, expected_assignment, poll_until, poll_until_assignment, unique_group,
    unique_topic,
};

#[tokio::test(flavor = "multi_thread")]
async fn producer_round_trips_empty_key_empty_value_and_nullable_headers() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("edge-empty-record");
    kafka.create_topic(&topic).await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), unique_group("edge-empty-record"))
            .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::new())
                .with_key(Bytes::new())
                .with_headers([
                    RecordHeader::new("present", Bytes::from_static(b"value")),
                    RecordHeader::new("empty", Bytes::new()),
                    RecordHeader::null("null"),
                ]),
        )
        .await?;

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

    let record = records
        .iter()
        .find(|record| {
            record.topic == topic
                && record.key.as_deref() == Some(&b""[..])
                && record.value.as_deref() == Some(&b""[..])
        })
        .ok_or_else(|| anyhow!("expected empty-key empty-value record"))?;
    assert_eq!(record.headers.len(), 3);
    assert_eq!(record.headers[0].key, "present");
    assert_eq!(record.headers[0].value.as_deref(), Some(&b"value"[..]));
    assert_eq!(record.headers[1].key, "empty");
    assert_eq!(record.headers[1].value.as_deref(), Some(&b""[..]));
    assert_eq!(record.headers[2].key, "null");
    assert!(record.headers[2].value.is_none());

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

#[tokio::test(flavor = "multi_thread")]
async fn producer_round_trips_all_compression_codecs() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("edge-compression");
    kafka.create_topic(&topic).await?;

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

    let codecs = [
        (ProducerCompression::None, "none"),
        (ProducerCompression::Gzip, "gzip"),
        (ProducerCompression::Snappy, "snappy"),
        (ProducerCompression::Lz4, "lz4"),
        (ProducerCompression::Zstd, "zstd"),
    ];

    for (codec, label) in codecs {
        let producer = KafkaProducer::connect(
            ProducerConfig::new(kafka.bootstrap_server())
                .with_client_id(format!("edge-compression-{label}"))
                .with_compression(codec),
        )
        .await?;
        producer
            .send(ProduceRecord::new(
                topic.clone(),
                0,
                Bytes::from(format!("compressed-{label}")),
            ))
            .await?;
        producer.shutdown().await?;
    }

    let records = poll_until(
        &consumer,
        |records| records.len() >= codecs.len(),
        Duration::from_secs(20),
    )
    .await?;
    let mut values = records
        .iter()
        .filter_map(|record| record.value.as_deref())
        .map(|value| String::from_utf8_lossy(value).into_owned())
        .collect::<Vec<_>>();
    values.sort();
    assert_eq!(
        values,
        vec![
            "compressed-gzip".to_owned(),
            "compressed-lz4".to_owned(),
            "compressed-none".to_owned(),
            "compressed-snappy".to_owned(),
            "compressed-zstd".to_owned(),
        ]
    );

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

#[tokio::test(flavor = "multi_thread")]
async fn consumer_latest_offset_reset_skips_existing_records_and_reads_new_records() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("edge-latest-reset");
    kafka.create_topic(&topic).await?;

    let producer = KafkaProducer::connect(ProducerConfig::new(kafka.bootstrap_server())).await?;
    producer
        .send(ProduceRecord::new(
            topic.clone(),
            0,
            Bytes::from_static(b"before-subscribe"),
        ))
        .await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), unique_group("edge-latest-reset"))
            .with_auto_offset_reset(AutoOffsetReset::Latest),
    )
    .await?;
    consumer.subscribe(vec![topic.clone()]).await?;
    poll_until_assignment(
        &consumer,
        expected_assignment(&topic, 1),
        Duration::from_secs(20),
    )
    .await?;

    let initial = consumer.poll_for(Duration::from_secs(2)).await?;
    assert!(
        initial
            .iter()
            .all(|record| record.value.as_deref() != Some(&b"before-subscribe"[..]))
    );

    producer
        .send(ProduceRecord::new(
            topic.clone(),
            0,
            Bytes::from_static(b"after-subscribe"),
        ))
        .await?;
    let records = poll_until(
        &consumer,
        |records| {
            records
                .iter()
                .any(|record| record.value.as_deref() == Some(&b"after-subscribe"[..]))
        },
        Duration::from_secs(20),
    )
    .await?;
    assert!(
        records
            .iter()
            .all(|record| record.value.as_deref() != Some(&b"before-subscribe"[..]))
    );

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

#[tokio::test(flavor = "multi_thread")]
async fn consumer_subscribes_to_multiple_topics_and_commits_each_partition() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic_a = unique_topic("edge-multi-topic-a");
    let topic_b = unique_topic("edge-multi-topic-b");
    let group = unique_group("edge-multi-topic");
    let admin = KafkaAdmin::connect(AdminConfig::new(kafka.bootstrap_server())).await?;
    admin
        .create_topics([
            NewTopic::new(topic_a.clone(), 1, 1),
            NewTopic::new(topic_b.clone(), 1, 1),
        ])
        .await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), group.clone())
            .with_auto_offset_reset(AutoOffsetReset::Earliest),
    )
    .await?;
    consumer
        .subscribe(vec![topic_a.clone(), topic_b.clone()])
        .await?;

    let producer = KafkaProducer::connect(ProducerConfig::new(kafka.bootstrap_server())).await?;
    producer
        .send(ProduceRecord::new(
            topic_a.clone(),
            0,
            Bytes::from_static(b"topic-a-value"),
        ))
        .await?;
    producer
        .send(ProduceRecord::new(
            topic_b.clone(),
            0,
            Bytes::from_static(b"topic-b-value"),
        ))
        .await?;

    let records = poll_until(
        &consumer,
        |records| {
            records
                .iter()
                .any(|record| record.topic == topic_a && record.offset == 0)
                && records
                    .iter()
                    .any(|record| record.topic == topic_b && record.offset == 0)
        },
        Duration::from_secs(20),
    )
    .await?;
    consumer.commit(&records).await?;
    consumer.shutdown().await?;

    let restarted = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), group)
            .with_auto_offset_reset(AutoOffsetReset::Earliest),
    )
    .await?;
    restarted
        .subscribe(vec![topic_a.clone(), topic_b.clone()])
        .await?;
    let replay = restarted.poll_for(Duration::from_secs(2)).await?;
    assert!(replay.iter().all(|record| {
        record.value.as_deref() != Some(&b"topic-a-value"[..])
            && record.value.as_deref() != Some(&b"topic-b-value"[..])
    }));

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

#[tokio::test(flavor = "multi_thread")]
async fn producer_reports_unknown_explicit_partition_without_hanging() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("edge-unknown-partition");
    kafka.create_topic(&topic).await?;

    let producer = KafkaProducer::connect(
        ProducerConfig::new(kafka.bootstrap_server())
            .with_delivery_timeout(Duration::from_millis(500))
            .with_retry_backoff(Duration::from_millis(50)),
    )
    .await?;

    let result = tokio::time::timeout(
        Duration::from_secs(5),
        producer.send(ProduceRecord::new(
            topic.clone(),
            99,
            Bytes::from_static(b"invalid-partition"),
        )),
    )
    .await?;
    let error = result.expect_err("send to missing partition should fail");
    assert!(
        error.to_string().contains("no leader metadata")
            || error.to_string().contains("delivery timed out"),
        "unexpected error: {error:#}"
    );

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