use super::*;
use kafka_protocol::records::RecordBatchDecoder;

#[tokio::test]
async fn fetch_round_trips_headers_from_produced_records() {
    let broker = test_broker();
    let records = vec![Record {
        transactional: false,
        control: false,
        partition_leader_epoch: 0,
        producer_id: -1,
        producer_epoch: -1,
        timestamp_type: TimestampType::Creation,
        offset: 0,
        sequence: -1,
        timestamp: 100,
        key: Some(Bytes::from_static(b"key")),
        value: Some(Bytes::from_static(b"value")),
        headers: IndexMap::from([
            (
                StrBytes::from("trace-id".to_string()),
                Some(Bytes::from_static(b"abc123")),
            ),
            (StrBytes::from("empty".to_string()), None),
        ]),
    }];
    let mut encoded = BytesMut::new();
    RecordBatchEncoder::encode(
        &mut encoded,
        &records,
        &RecordEncodeOptions {
            version: 2,
            compression: Compression::None,
        },
    )
    .unwrap();
    let request = ProduceRequest::default()
        .with_acks(1)
        .with_timeout_ms(5_000)
        .with_topic_data(vec![
            TopicProduceData::default()
                .with_name(TopicName(StrBytes::from("headers.topic".to_string())))
                .with_partition_data(vec![
                    PartitionProduceData::default()
                        .with_index(0)
                        .with_records(Some(encoded.freeze())),
                ]),
        ]);
    let _ = handle_produce(&broker, request).await.unwrap();

    let fetch = handle_fetch(
        &broker,
        FetchRequest::default().with_topics(vec![
            kafka_protocol::messages::fetch_request::FetchTopic::default()
                .with_topic(TopicName(StrBytes::from("headers.topic".to_string())))
                .with_partitions(vec![
                    kafka_protocol::messages::fetch_request::FetchPartition::default()
                        .with_partition(0)
                        .with_fetch_offset(0)
                        .with_partition_max_bytes(1024),
                ]),
        ]),
    )
    .await
    .unwrap();

    let mut payload = fetch.responses[0].partitions[0].records.clone().unwrap();
    let decoded = RecordBatchDecoder::decode_all(&mut payload).unwrap();
    let record = decoded
        .into_iter()
        .flat_map(|batch| batch.records)
        .next()
        .unwrap();

    assert_eq!(
        record
            .headers
            .get(&StrBytes::from("trace-id".to_string()))
            .cloned(),
        Some(Some(Bytes::from_static(b"abc123")))
    );
    assert_eq!(
        record
            .headers
            .get(&StrBytes::from("empty".to_string()))
            .cloned(),
        Some(None)
    );
}

#[tokio::test]
async fn fetch_round_trips_null_values_from_produced_records() {
    let broker = test_broker();
    let records = vec![Record {
        transactional: false,
        control: false,
        partition_leader_epoch: 0,
        producer_id: -1,
        producer_epoch: -1,
        timestamp_type: TimestampType::Creation,
        offset: 0,
        sequence: -1,
        timestamp: 100,
        key: Some(Bytes::from_static(b"tombstone-key")),
        value: None,
        headers: Default::default(),
    }];
    let mut encoded = BytesMut::new();
    RecordBatchEncoder::encode(
        &mut encoded,
        &records,
        &RecordEncodeOptions {
            version: 2,
            compression: Compression::None,
        },
    )
    .unwrap();
    let request = ProduceRequest::default()
        .with_acks(1)
        .with_timeout_ms(5_000)
        .with_topic_data(vec![
            TopicProduceData::default()
                .with_name(TopicName(StrBytes::from("tombstone.topic".to_string())))
                .with_partition_data(vec![
                    PartitionProduceData::default()
                        .with_index(0)
                        .with_records(Some(encoded.freeze())),
                ]),
        ]);
    let _ = handle_produce(&broker, request).await.unwrap();

    let fetch = handle_fetch(
        &broker,
        FetchRequest::default().with_topics(vec![
            kafka_protocol::messages::fetch_request::FetchTopic::default()
                .with_topic(TopicName(StrBytes::from("tombstone.topic".to_string())))
                .with_partitions(vec![
                    kafka_protocol::messages::fetch_request::FetchPartition::default()
                        .with_partition(0)
                        .with_fetch_offset(0)
                        .with_partition_max_bytes(1024),
                ]),
        ]),
    )
    .await
    .unwrap();

    let mut payload = fetch.responses[0].partitions[0].records.clone().unwrap();
    let decoded = RecordBatchDecoder::decode_all(&mut payload).unwrap();
    let record = decoded
        .into_iter()
        .flat_map(|batch| batch.records)
        .next()
        .unwrap();

    assert_eq!(record.key, Some(Bytes::from_static(b"tombstone-key")));
    assert_eq!(record.value, None);
}

#[tokio::test]
async fn fetch_round_trips_tombstone_with_headers() {
    let broker = test_broker();
    let records = vec![Record {
        transactional: false,
        control: false,
        partition_leader_epoch: 0,
        producer_id: -1,
        producer_epoch: -1,
        timestamp_type: TimestampType::Creation,
        offset: 0,
        sequence: -1,
        timestamp: 100,
        key: Some(Bytes::from_static(b"tombstone-key")),
        value: None,
        headers: IndexMap::from([
            (
                StrBytes::from("trace-id".to_string()),
                Some(Bytes::from_static(b"abc123")),
            ),
            (StrBytes::from("empty".to_string()), None),
        ]),
    }];
    let mut encoded = BytesMut::new();
    RecordBatchEncoder::encode(
        &mut encoded,
        &records,
        &RecordEncodeOptions {
            version: 2,
            compression: Compression::None,
        },
    )
    .unwrap();
    let request = ProduceRequest::default()
        .with_acks(1)
        .with_timeout_ms(5_000)
        .with_topic_data(vec![
            TopicProduceData::default()
                .with_name(TopicName(StrBytes::from(
                    "tombstone.headers.topic".to_string(),
                )))
                .with_partition_data(vec![
                    PartitionProduceData::default()
                        .with_index(0)
                        .with_records(Some(encoded.freeze())),
                ]),
        ]);
    let _ = handle_produce(&broker, request).await.unwrap();

    let fetch = handle_fetch(
        &broker,
        FetchRequest::default().with_topics(vec![
            kafka_protocol::messages::fetch_request::FetchTopic::default()
                .with_topic(TopicName(StrBytes::from(
                    "tombstone.headers.topic".to_string(),
                )))
                .with_partitions(vec![
                    kafka_protocol::messages::fetch_request::FetchPartition::default()
                        .with_partition(0)
                        .with_fetch_offset(0)
                        .with_partition_max_bytes(1024),
                ]),
        ]),
    )
    .await
    .unwrap();

    let mut payload = fetch.responses[0].partitions[0].records.clone().unwrap();
    let decoded = RecordBatchDecoder::decode_all(&mut payload).unwrap();
    let record = decoded
        .into_iter()
        .flat_map(|batch| batch.records)
        .next()
        .unwrap();

    assert_eq!(record.key, Some(Bytes::from_static(b"tombstone-key")));
    assert_eq!(record.value, None);
    assert_eq!(
        record
            .headers
            .get(&StrBytes::from("trace-id".to_string()))
            .cloned(),
        Some(Some(Bytes::from_static(b"abc123")))
    );
    assert_eq!(
        record
            .headers
            .get(&StrBytes::from("empty".to_string()))
            .cloned(),
        Some(None)
    );
}
