mod common;

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

use anyhow::{Result, bail};
use bytes::Bytes;
use kafkit_client::{
    AcknowledgeType, AdminConfig, ConsumerConfig, KafkaAdmin, KafkaProducer, KafkaShareConsumer,
    NewTopic, ProduceRecord, ProducerConfig, RecordHeader, ShareAcknowledgementCommit,
    ShareAcquireMode, ShareConsumerOptions, ShareRecords,
};

use crate::common::{KafkaHarness, unique_group, unique_topic};

#[tokio::test(flavor = "multi_thread")]
async fn share_consumer_accepts_and_commits_records() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("share-accept-commit");
    let group = unique_group("share-accept-commit");
    create_share_topic(kafka.bootstrap_server(), &topic, 1).await?;

    let completed = Arc::new(Mutex::new(Vec::<ShareAcknowledgementCommit>::new()));
    let completed_for_callback = Arc::clone(&completed);
    let mut consumer = KafkaShareConsumer::connect_with_options(
        ConsumerConfig::new(kafka.bootstrap_server(), group),
        ShareConsumerOptions::default().with_acknowledgement_commit_callback(move |commits| {
            completed_for_callback
                .lock()
                .expect("callback mutex")
                .extend(commits);
        }),
    )
    .await?;
    consumer.subscribe(vec![topic.clone()]).await?;
    wait_for_share_assignment(&mut consumer, Duration::from_secs(60)).await?;

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

    let records = poll_share_until(
        &mut consumer,
        |records| records.len() >= 2,
        Duration::from_secs(60),
    )
    .await?;
    for record in records.iter() {
        consumer.acknowledge(record, AcknowledgeType::Accept);
    }
    consumer.commit_sync().await?;

    let committed_offsets = completed
        .lock()
        .expect("callback mutex")
        .iter()
        .flat_map(|commit| commit.offsets.iter().copied())
        .collect::<Vec<_>>();
    assert!(committed_offsets.contains(&0));
    assert!(committed_offsets.contains(&1));

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

#[tokio::test(flavor = "multi_thread")]
async fn share_consumer_release_makes_record_available_again() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("share-release-redelivery");
    let group = unique_group("share-release-redelivery");
    create_share_topic(kafka.bootstrap_server(), &topic, 1).await?;

    let mut consumer =
        KafkaShareConsumer::connect(ConsumerConfig::new(kafka.bootstrap_server(), group)).await?;
    consumer.subscribe(vec![topic.clone()]).await?;
    wait_for_share_assignment(&mut consumer, Duration::from_secs(60)).await?;

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

    let first = poll_share_until(
        &mut consumer,
        |records| {
            records
                .iter()
                .any(|record| record.record.value.as_deref() == Some(&b"retry-me"[..]))
        },
        Duration::from_secs(60),
    )
    .await?;
    let record = first
        .iter()
        .find(|record| record.record.value.as_deref() == Some(&b"retry-me"[..]))
        .expect("retry-me should be fetched first");
    consumer.acknowledge(record, AcknowledgeType::Release);
    consumer.commit_sync().await?;

    let redelivered = poll_share_until(
        &mut consumer,
        |records| {
            records.iter().any(|record| {
                record.record.value.as_deref() == Some(&b"retry-me"[..])
                    && record.delivery_count >= 2
            })
        },
        Duration::from_secs(60),
    )
    .await?;
    let record = redelivered
        .iter()
        .find(|record| record.record.value.as_deref() == Some(&b"retry-me"[..]))
        .expect("retry-me should be redelivered");
    assert!(record.delivery_count >= 2);
    consumer.acknowledge(record, AcknowledgeType::Accept);
    consumer.commit_sync().await?;

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

#[tokio::test(flavor = "multi_thread")]
async fn share_consumer_record_limit_mode_fetches_records() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("share-record-limit");
    let group = unique_group("share-record-limit");
    create_share_topic(kafka.bootstrap_server(), &topic, 2).await?;

    let mut consumer = KafkaShareConsumer::connect_with_options(
        ConsumerConfig::new(kafka.bootstrap_server(), group),
        ShareConsumerOptions::default()
            .with_share_acquire_mode(ShareAcquireMode::RecordLimit)
            .with_max_poll_records(1),
    )
    .await?;
    assert_eq!(consumer.share_acquire_mode(), ShareAcquireMode::RecordLimit);
    consumer.subscribe(vec![topic.clone()]).await?;
    wait_for_share_assignment(&mut consumer, Duration::from_secs(60)).await?;

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

    let records = poll_share_until(
        &mut consumer,
        |records| !records.is_empty(),
        Duration::from_secs(60),
    )
    .await?;
    for record in records.iter() {
        consumer.acknowledge(record, AcknowledgeType::Accept);
    }
    consumer.commit_sync().await?;

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

#[tokio::test(flavor = "multi_thread")]
async fn share_consumer_round_trips_empty_payload_key_and_headers() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("share-empty-record");
    let group = unique_group("share-empty-record");
    create_share_topic(kafka.bootstrap_server(), &topic, 1).await?;

    let mut consumer =
        KafkaShareConsumer::connect(ConsumerConfig::new(kafka.bootstrap_server(), group)).await?;
    consumer.subscribe(vec![topic.clone()]).await?;
    wait_for_share_assignment(&mut consumer, Duration::from_secs(60)).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_share_until(
        &mut consumer,
        |records| {
            records.iter().any(|record| {
                record.record.key.as_deref() == Some(&b""[..])
                    && record.record.value.as_deref() == Some(&b""[..])
            })
        },
        Duration::from_secs(60),
    )
    .await?;
    let record = records
        .iter()
        .find(|record| record.record.value.as_deref() == Some(&b""[..]))
        .expect("expected empty share record");
    assert_eq!(record.record.headers.len(), 3);
    assert_eq!(record.record.headers[0].key, "present");
    assert_eq!(
        record.record.headers[0].value.as_deref(),
        Some(&b"value"[..])
    );
    assert_eq!(record.record.headers[1].key, "empty");
    assert_eq!(record.record.headers[1].value.as_deref(), Some(&b""[..]));
    assert_eq!(record.record.headers[2].key, "null");
    assert!(record.record.headers[2].value.is_none());

    consumer.acknowledge(record, AcknowledgeType::Accept);
    consumer.commit_sync().await?;
    consumer.shutdown().await?;
    producer.shutdown().await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn share_consumer_subscribes_to_multiple_topics_and_acknowledges_each() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic_a = unique_topic("share-multi-a");
    let topic_b = unique_topic("share-multi-b");
    let group = unique_group("share-multi");
    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 mut consumer =
        KafkaShareConsumer::connect(ConsumerConfig::new(kafka.bootstrap_server(), group)).await?;
    consumer
        .subscribe(vec![topic_a.clone(), topic_b.clone()])
        .await?;
    wait_for_share_assignment_count(&mut consumer, 2, Duration::from_secs(60)).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_share_until(
        &mut consumer,
        |records| {
            records.iter().any(|record| {
                record.record.topic == topic_a
                    && record.record.value.as_deref() == Some(&b"topic-a-value"[..])
            }) && records.iter().any(|record| {
                record.record.topic == topic_b
                    && record.record.value.as_deref() == Some(&b"topic-b-value"[..])
            })
        },
        Duration::from_secs(60),
    )
    .await?;
    for record in records.iter() {
        consumer.acknowledge(record, AcknowledgeType::Accept);
    }
    consumer.commit_sync().await?;

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

#[tokio::test(flavor = "multi_thread")]
async fn share_consumer_reject_does_not_redeliver_record() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("share-reject");
    let group = unique_group("share-reject");
    create_share_topic(kafka.bootstrap_server(), &topic, 1).await?;

    let mut consumer =
        KafkaShareConsumer::connect(ConsumerConfig::new(kafka.bootstrap_server(), group)).await?;
    consumer.subscribe(vec![topic.clone()]).await?;
    wait_for_share_assignment(&mut consumer, Duration::from_secs(60)).await?;

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

    let records = poll_share_until(
        &mut consumer,
        |records| {
            records
                .iter()
                .any(|record| record.record.value.as_deref() == Some(&b"reject-me"[..]))
        },
        Duration::from_secs(60),
    )
    .await?;
    let record = records
        .iter()
        .find(|record| record.record.value.as_deref() == Some(&b"reject-me"[..]))
        .expect("reject-me should be fetched");
    consumer.acknowledge(record, AcknowledgeType::Reject);
    consumer.commit_sync().await?;

    let redelivered =
        poll_share_for_value(&mut consumer, b"reject-me", Duration::from_secs(3)).await?;
    assert!(
        !redelivered,
        "rejected share record should not be redelivered"
    );

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

#[tokio::test(flavor = "multi_thread")]
async fn share_consumer_empty_commit_is_noop() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("share-empty-commit");
    let group = unique_group("share-empty-commit");
    create_share_topic(kafka.bootstrap_server(), &topic, 1).await?;

    let completed = Arc::new(Mutex::new(Vec::<ShareAcknowledgementCommit>::new()));
    let completed_for_callback = Arc::clone(&completed);
    let mut consumer = KafkaShareConsumer::connect_with_options(
        ConsumerConfig::new(kafka.bootstrap_server(), group),
        ShareConsumerOptions::default().with_acknowledgement_commit_callback(move |commits| {
            completed_for_callback
                .lock()
                .expect("callback mutex")
                .extend(commits);
        }),
    )
    .await?;
    consumer.subscribe(vec![topic]).await?;
    wait_for_share_assignment(&mut consumer, Duration::from_secs(60)).await?;

    consumer.commit_sync().await?;
    assert!(completed.lock().expect("callback mutex").is_empty());

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

async fn create_share_topic(bootstrap_server: &str, topic: &str, partitions: i32) -> Result<()> {
    let admin = KafkaAdmin::connect(AdminConfig::new(bootstrap_server)).await?;
    admin
        .create_topics([NewTopic::new(topic.to_owned(), partitions, 1)])
        .await?;
    Ok(())
}

async fn poll_share_until<F>(
    consumer: &mut KafkaShareConsumer,
    predicate: F,
    timeout: Duration,
) -> Result<ShareRecords>
where
    F: Fn(&ShareRecords) -> bool,
{
    let deadline = Instant::now() + timeout;
    loop {
        let records = consumer.poll_for(Duration::from_millis(500)).await?;
        if predicate(&records) {
            return Ok(records);
        }

        if Instant::now() >= deadline {
            bail!("timed out waiting for share records");
        }
    }
}

async fn wait_for_share_assignment(
    consumer: &mut KafkaShareConsumer,
    timeout: Duration,
) -> Result<()> {
    let deadline = Instant::now() + timeout;
    loop {
        let _ = consumer.poll_for(Duration::from_millis(500)).await?;
        if !consumer.assignment().is_empty() {
            return Ok(());
        }

        if Instant::now() >= deadline {
            bail!("timed out waiting for share assignment");
        }
    }
}

async fn wait_for_share_assignment_count(
    consumer: &mut KafkaShareConsumer,
    expected_count: usize,
    timeout: Duration,
) -> Result<()> {
    let deadline = Instant::now() + timeout;
    loop {
        let _ = consumer.poll_for(Duration::from_millis(500)).await?;
        if consumer.assignment().len() >= expected_count {
            return Ok(());
        }

        if Instant::now() >= deadline {
            bail!("timed out waiting for {expected_count} share assignments");
        }
    }
}

async fn poll_share_for_value(
    consumer: &mut KafkaShareConsumer,
    value: &[u8],
    timeout: Duration,
) -> Result<bool> {
    let deadline = Instant::now() + timeout;
    loop {
        let records = consumer.poll_for(Duration::from_millis(500)).await?;
        if records
            .iter()
            .any(|record| record.record.value.as_deref() == Some(value))
        {
            return Ok(true);
        }

        if Instant::now() >= deadline {
            return Ok(false);
        }
    }
}
