mod common;

use std::collections::{BTreeSet, HashSet};
use std::env;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use anyhow::{Context, Result, bail};
use kafkit_client::{
    AdminConfig, AutoOffsetReset, ConsumerConfig, KafkaAdmin, KafkaConsumer, KafkaProducer,
    NewTopic, ProduceRecord, ProducerCompression, ProducerConfig, TopicPartition,
};
use tokio::task::JoinSet;

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

const DEFAULT_PARTITIONS: i32 = 8;
const DEFAULT_PRODUCER_TASKS: usize = 4;
const DEFAULT_CONSUMER_TASKS: usize = 4;
const DEFAULT_MESSAGES_PER_TASK: usize = 400;
const DEFAULT_PAYLOAD_BYTES: usize = 512;
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(90);

#[tokio::test(flavor = "multi_thread", worker_threads = 6)]
#[ignore = "stress integration test; run manually"]
async fn stress_parallel_producers_and_single_consumer_report_timing() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("stress-burst");
    let partitions = stress_env_i32("STRESS_PARTITIONS", DEFAULT_PARTITIONS);
    let producer_tasks = stress_env_usize("STRESS_PRODUCER_TASKS", DEFAULT_PRODUCER_TASKS);
    let messages_per_task = stress_env_usize("STRESS_MESSAGES_PER_TASK", DEFAULT_MESSAGES_PER_TASK);
    let payload_bytes = stress_env_usize("STRESS_PAYLOAD_BYTES", DEFAULT_PAYLOAD_BYTES);
    let expected_messages = producer_tasks * messages_per_task;

    let admin = KafkaAdmin::connect(AdminConfig::new(kafka.bootstrap_server())).await?;
    admin
        .create_topics([NewTopic::new(topic.clone(), partitions, 1)])
        .await?;
    wait_for_partition_count(&admin, &topic, partitions, Duration::from_secs(20)).await?;

    let consumer = KafkaConsumer::connect(
        ConsumerConfig::new(kafka.bootstrap_server(), unique_group("stress-burst"))
            .with_auto_offset_reset(AutoOffsetReset::Earliest)
            .with_fetch_max_wait(Duration::from_millis(100))
            .with_fetch_min_bytes(1)
            .with_partition_max_bytes(2 * 1024 * 1024)
            .with_fetch_max_bytes(64 * 1024 * 1024),
    )
    .await?;
    consumer.subscribe(vec![topic.clone()]).await?;
    wait_for_assignment(&consumer, Duration::from_secs(20)).await?;

    let producer_window = Instant::now();
    let bootstrap_server = kafka.bootstrap_server().to_owned();
    let producer_topic = topic.clone();
    let producer_task = tokio::spawn(async move {
        drive_parallel_producers(
            &bootstrap_server,
            &producer_topic,
            partitions,
            producer_tasks,
            messages_per_task,
            payload_bytes,
        )
        .await
    });

    let consumer_metrics =
        drain_consumer_records(&consumer, expected_messages, DEFAULT_TIMEOUT).await?;
    let producer_summary = producer_task.await??;
    let producer_elapsed = producer_window.elapsed();
    consumer.shutdown().await?;

    assert_eq!(producer_summary.samples, expected_messages);
    assert_eq!(consumer_metrics.unique_messages, expected_messages);
    assert_eq!(consumer_metrics.duplicates, 0);
    assert!(
        consumer_metrics.end_to_end.p95 < Duration::from_secs(10),
        "p95 end-to-end latency was {:?}",
        consumer_metrics.end_to_end.p95
    );
    assert!(
        producer_summary.throughput_per_sec > 10.0,
        "producer throughput was too low: {:.2} msg/s",
        producer_summary.throughput_per_sec
    );
    assert!(
        consumer_metrics.throughput_per_sec > 10.0,
        "consumer throughput was too low: {:.2} msg/s",
        consumer_metrics.throughput_per_sec
    );

    eprintln!(
        "stress burst summary: sent={} produced_in={:?} producer={} consumer={} empty_polls={}",
        expected_messages,
        producer_elapsed,
        producer_summary,
        consumer_metrics.end_to_end,
        consumer_metrics.empty_polls
    );
    eprintln!(
        "stress burst polling: polls={} poll_latency={} consume_throughput={:.2} msg/s",
        consumer_metrics.polls, consumer_metrics.poll_latency, consumer_metrics.throughput_per_sec
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 6)]
#[ignore = "stress integration test; run manually"]
async fn stress_manual_assignment_consumers_scale_across_partitions() -> Result<()> {
    let kafka = KafkaHarness::start().await?;
    let topic = unique_topic("stress-manual");
    let partitions = stress_env_i32("STRESS_PARTITIONS", DEFAULT_PARTITIONS).max(2);
    let consumer_tasks = stress_env_usize("STRESS_CONSUMER_TASKS", DEFAULT_CONSUMER_TASKS)
        .min(usize::try_from(partitions).unwrap_or(DEFAULT_CONSUMER_TASKS))
        .max(1);
    let messages_per_partition =
        stress_env_usize("STRESS_MESSAGES_PER_PARTITION", DEFAULT_MESSAGES_PER_TASK);
    let payload_bytes = stress_env_usize("STRESS_PAYLOAD_BYTES", DEFAULT_PAYLOAD_BYTES);
    let expected_messages = usize::try_from(partitions).unwrap_or(0) * messages_per_partition;

    let admin = KafkaAdmin::connect(AdminConfig::new(kafka.bootstrap_server())).await?;
    admin
        .create_topics([NewTopic::new(topic.clone(), partitions, 1)])
        .await?;
    wait_for_partition_count(&admin, &topic, partitions, Duration::from_secs(20)).await?;

    let mut consumers = Vec::new();
    let partition_groups = split_partitions(partitions, consumer_tasks);
    for (idx, partitions_for_consumer) in partition_groups.iter().enumerate() {
        let consumer = KafkaConsumer::connect(
            ConsumerConfig::new(
                kafka.bootstrap_server(),
                unique_group(&format!("stress-manual-{idx}")),
            )
            .with_auto_offset_reset(AutoOffsetReset::Earliest)
            .with_fetch_max_wait(Duration::from_millis(50))
            .with_fetch_min_bytes(1)
            .with_partition_max_bytes(2 * 1024 * 1024),
        )
        .await?;
        let assignment = partitions_for_consumer
            .iter()
            .map(|partition| TopicPartition::new(topic.clone(), *partition))
            .collect::<Vec<_>>();
        consumer.assign(assignment).await?;
        consumers.push((consumer, partitions_for_consumer.clone()));
    }

    let mut consumers_set = JoinSet::new();
    for (consumer, partitions_for_consumer) in consumers {
        let expected_for_consumer = partitions_for_consumer.len() * messages_per_partition;
        consumers_set.spawn(async move {
            let metrics =
                drain_consumer_records(&consumer, expected_for_consumer, DEFAULT_TIMEOUT).await?;
            consumer.shutdown().await?;
            Ok::<_, anyhow::Error>((metrics, partitions_for_consumer))
        });
    }

    let producer = KafkaProducer::connect(
        ProducerConfig::new(kafka.bootstrap_server())
            .with_client_id("stress-manual-producer")
            .with_enable_idempotence(true)
            .with_compression(ProducerCompression::Lz4)
            .with_batch_size(64 * 1024)
            .with_linger(Duration::from_millis(5))
            .with_request_timeout(Duration::from_secs(10))
            .with_delivery_timeout(Duration::from_secs(30)),
    )
    .await?;

    let produce_started = Instant::now();
    let mut produce_latencies = Vec::with_capacity(expected_messages);
    for partition in 0..partitions {
        for seq in 0..messages_per_partition {
            let id = usize::try_from(partition).unwrap_or(0) * messages_per_partition + seq;
            let payload = make_payload(id, payload_bytes);
            let started = Instant::now();
            producer
                .send(ProduceRecord::new(topic.clone(), partition, payload))
                .await?;
            produce_latencies.push(started.elapsed());
        }
    }
    producer.flush().await?;
    producer.shutdown().await?;
    let produce_elapsed = produce_started.elapsed();

    let mut consumed_total = 0usize;
    let mut duplicates = 0usize;
    let mut poll_latencies = Vec::new();
    let mut end_to_end_latencies = Vec::new();
    let mut empty_polls = 0usize;
    let mut seen_partitions = BTreeSet::new();
    while let Some(result) = consumers_set.join_next().await {
        let (metrics, partitions_for_consumer) = result??;
        consumed_total += metrics.unique_messages;
        duplicates += metrics.duplicates;
        empty_polls += metrics.empty_polls;
        poll_latencies.extend(metrics.poll_latency_samples);
        end_to_end_latencies.extend(metrics.end_to_end_samples);
        seen_partitions.extend(partitions_for_consumer);
    }

    let produce_summary = DurationSummary::new(produce_latencies, produce_elapsed);
    let poll_summary = DurationSummary::new(poll_latencies, produce_elapsed);
    let end_to_end_summary = DurationSummary::new(end_to_end_latencies, produce_elapsed);

    assert_eq!(
        seen_partitions.len(),
        usize::try_from(partitions).unwrap_or(0)
    );
    assert_eq!(consumed_total, expected_messages);
    assert_eq!(duplicates, 0);
    assert!(
        end_to_end_summary.p95 < Duration::from_secs(10),
        "p95 end-to-end latency was {:?}",
        end_to_end_summary.p95
    );

    eprintln!(
        "stress manual summary: messages={} producer={} poll_latency={} end_to_end={} empty_polls={}",
        expected_messages, produce_summary, poll_summary, end_to_end_summary, empty_polls
    );

    Ok(())
}

async fn drive_parallel_producers(
    bootstrap_server: &str,
    topic: &str,
    partitions: i32,
    producer_tasks: usize,
    messages_per_task: usize,
    payload_bytes: usize,
) -> Result<DurationSummary> {
    let mut producers = JoinSet::new();
    for task_id in 0..producer_tasks {
        let bootstrap_server = bootstrap_server.to_owned();
        let topic = topic.to_owned();
        producers.spawn(async move {
            let producer = KafkaProducer::connect(
                ProducerConfig::new(bootstrap_server)
                    .with_client_id(format!("stress-producer-{task_id}"))
                    .with_enable_idempotence(true)
                    .with_compression(ProducerCompression::Lz4)
                    .with_batch_size(64 * 1024)
                    .with_linger(Duration::from_millis(5))
                    .with_request_timeout(Duration::from_secs(10))
                    .with_delivery_timeout(Duration::from_secs(30)),
            )
            .await?;

            let started = Instant::now();
            let mut latencies = Vec::with_capacity(messages_per_task);
            for seq in 0..messages_per_task {
                let id = task_id * messages_per_task + seq;
                let payload = make_payload(id, payload_bytes);
                let partition =
                    i32::try_from(id % usize::try_from(partitions).unwrap_or(1)).unwrap_or(0);
                let send_started = Instant::now();
                producer
                    .send(ProduceRecord::new(topic.clone(), partition, payload))
                    .await?;
                latencies.push(send_started.elapsed());
            }
            producer.flush().await?;
            producer.shutdown().await?;
            Ok::<_, anyhow::Error>((latencies, started.elapsed()))
        });
    }

    let mut all_latencies = Vec::new();
    let started = Instant::now();
    while let Some(result) = producers.join_next().await {
        let (latencies, _) = result??;
        all_latencies.extend(latencies);
    }

    Ok(DurationSummary::new(all_latencies, started.elapsed()))
}

async fn drain_consumer_records(
    consumer: &KafkaConsumer,
    expected_messages: usize,
    timeout: Duration,
) -> Result<DrainMetrics> {
    let deadline = Instant::now() + timeout;
    let started = Instant::now();
    let mut seen = HashSet::with_capacity(expected_messages);
    let mut duplicates = 0usize;
    let mut polls = 0usize;
    let mut empty_polls = 0usize;
    let mut poll_latency_samples = Vec::new();
    let mut end_to_end_samples = Vec::with_capacity(expected_messages);

    while seen.len() < expected_messages {
        if Instant::now() >= deadline {
            bail!(
                "timed out after {:?}: consumed {} of {} messages",
                timeout,
                seen.len(),
                expected_messages
            );
        }

        let poll_started = Instant::now();
        let records = consumer.poll_for(Duration::from_millis(250)).await?;
        polls += 1;
        poll_latency_samples.push(poll_started.elapsed());

        if records.is_empty() {
            empty_polls += 1;
            continue;
        }

        for record in records.iter() {
            let value = record
                .value
                .as_ref()
                .context("stress record unexpectedly had no value")?;
            let (id, sent_at_micros) = parse_payload(value)?;
            if !seen.insert(id) {
                duplicates += 1;
                continue;
            }
            end_to_end_samples.push(elapsed_since_micros(sent_at_micros));
        }
    }

    let elapsed = started.elapsed();
    let end_to_end = DurationSummary::new(end_to_end_samples.clone(), elapsed);
    let poll_latency = DurationSummary::new(poll_latency_samples.clone(), elapsed);

    Ok(DrainMetrics {
        unique_messages: seen.len(),
        duplicates,
        polls,
        empty_polls,
        throughput_per_sec: throughput_per_sec(seen.len(), elapsed),
        poll_latency,
        end_to_end,
        poll_latency_samples,
        end_to_end_samples,
    })
}

async fn wait_for_partition_count(
    admin: &KafkaAdmin,
    topic: &str,
    expected_partitions: i32,
    timeout: Duration,
) -> Result<()> {
    let deadline = Instant::now() + timeout;
    loop {
        let described = admin.describe_topics([topic.to_owned()]).await?;
        if described.first().is_some_and(|topic| {
            i32::try_from(topic.partitions.len()).ok() == Some(expected_partitions)
        }) {
            return Ok(());
        }
        if Instant::now() >= deadline {
            bail!(
                "timed out waiting for topic '{}' to expose {} partitions",
                topic,
                expected_partitions
            );
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
}

async fn wait_for_assignment(consumer: &KafkaConsumer, timeout: Duration) -> Result<()> {
    let deadline = Instant::now() + timeout;
    loop {
        let _ = consumer.poll_for(Duration::from_millis(200)).await?;
        if !consumer.assignment().await?.is_empty() {
            return Ok(());
        }
        if Instant::now() >= deadline {
            bail!("timed out waiting for consumer assignment");
        }
    }
}

fn split_partitions(partitions: i32, buckets: usize) -> Vec<Vec<i32>> {
    let mut grouped = vec![Vec::new(); buckets];
    for partition in 0..partitions {
        grouped[usize::try_from(partition).unwrap_or(0) % buckets].push(partition);
    }
    grouped
        .into_iter()
        .filter(|group| !group.is_empty())
        .collect()
}

fn make_payload(id: usize, payload_bytes: usize) -> String {
    let sent_at_micros = unix_time_micros();
    let mut payload = format!("{id}:{sent_at_micros}:");
    let target_padding = payload_bytes.saturating_sub(payload.len());
    if target_padding > 0 {
        payload.push_str(&"x".repeat(target_padding));
    }
    payload
}

fn parse_payload(value: &[u8]) -> Result<(usize, u128)> {
    let payload = std::str::from_utf8(value).context("stress payload was not utf-8")?;
    let mut fields = payload.splitn(3, ':');
    let id = fields
        .next()
        .context("stress payload did not contain an id")?
        .parse::<usize>()
        .context("stress payload id was not numeric")?;
    let sent_at_micros = fields
        .next()
        .context("stress payload did not contain a timestamp")?
        .parse::<u128>()
        .context("stress payload timestamp was not numeric")?;
    Ok((id, sent_at_micros))
}

fn unix_time_micros() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_micros()
}

fn elapsed_since_micros(sent_at_micros: u128) -> Duration {
    let now = unix_time_micros();
    let delta = now.saturating_sub(sent_at_micros);
    Duration::from_micros(u64::try_from(delta).unwrap_or(u64::MAX))
}

fn stress_env_usize(name: &str, default: usize) -> usize {
    env::var(name)
        .ok()
        .and_then(|value| value.trim().parse().ok())
        .filter(|value| *value > 0)
        .unwrap_or(default)
}

fn stress_env_i32(name: &str, default: i32) -> i32 {
    env::var(name)
        .ok()
        .and_then(|value| value.trim().parse().ok())
        .filter(|value| *value > 0)
        .unwrap_or(default)
}

fn throughput_per_sec(samples: usize, elapsed: Duration) -> f64 {
    let seconds = elapsed.as_secs_f64();
    if seconds <= f64::EPSILON {
        samples as f64
    } else {
        samples as f64 / seconds
    }
}

struct DrainMetrics {
    unique_messages: usize,
    duplicates: usize,
    polls: usize,
    empty_polls: usize,
    throughput_per_sec: f64,
    poll_latency: DurationSummary,
    end_to_end: DurationSummary,
    poll_latency_samples: Vec<Duration>,
    end_to_end_samples: Vec<Duration>,
}

#[derive(Clone)]
struct DurationSummary {
    samples: usize,
    min: Duration,
    p50: Duration,
    p95: Duration,
    max: Duration,
    throughput_per_sec: f64,
}

impl DurationSummary {
    fn new(mut samples: Vec<Duration>, elapsed: Duration) -> Self {
        samples.sort_unstable();
        let len = samples.len();
        let min = samples.first().copied().unwrap_or_default();
        let max = samples.last().copied().unwrap_or_default();
        let p50 = percentile(&samples, 0.50);
        let p95 = percentile(&samples, 0.95);
        Self {
            samples: len,
            min,
            p50,
            p95,
            max,
            throughput_per_sec: throughput_per_sec(len, elapsed),
        }
    }
}

impl std::fmt::Display for DurationSummary {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "samples={} throughput={:.2} msg/s min={:?} p50={:?} p95={:?} max={:?}",
            self.samples, self.throughput_per_sec, self.min, self.p50, self.p95, self.max
        )
    }
}

fn percentile(samples: &[Duration], percentile: f64) -> Duration {
    if samples.is_empty() {
        return Duration::ZERO;
    }
    let position = ((samples.len() - 1) as f64 * percentile).round() as usize;
    samples[position.min(samples.len() - 1)]
}
