//! `fafb-q bench` command — simple publish, subscribe, and latency benchmarks.

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

use bytes::Bytes;
use colored::Colorize;

use crate::cli::{
    BenchAction, CliConfig, CliError, OutputFormat, build_client_config, parse_topic,
    resolve_namespace,
};
use crate::client::FafbClient;
use crate::codec::frame::Frame;
use crate::types::{GroupName, Headers, TopicOffset};

/// Progress update interval for long-running benchmarks.
const PROGRESS_INTERVAL: Duration = Duration::from_secs(2);

pub async fn run(config: CliConfig, action: BenchAction) -> Result<(), CliError> {
    match action {
        BenchAction::Pub {
            topic,
            messages,
            payload_size,
            concurrency,
            no_ack,
        } => run_pub_bench(config, &topic, messages, payload_size, concurrency, no_ack).await,
        BenchAction::Sub {
            topic,
            group,
            from,
            count,
        } => run_sub_bench(config, &topic, group, from, count).await,
        BenchAction::Latency {
            topic,
            messages,
            warmup,
        } => run_latency_bench(config, &topic, messages, warmup).await,
    }
}

// ---------------------------------------------------------------------------
// Publish benchmark
// ---------------------------------------------------------------------------

async fn run_pub_bench(
    config: CliConfig,
    topic_str: &str,
    messages: u64,
    payload_size: usize,
    concurrency: u32,
    no_ack: bool,
) -> Result<(), CliError> {
    let format = OutputFormat::auto_detect(config.json, false);
    let (parsed_ns, topic) = parse_topic(topic_str);
    let ns = resolve_namespace(parsed_ns, config.namespace.as_deref(), &config)
        .map_err(CliError::Error)?;

    let payload = Bytes::from(vec![b'x'; payload_size]);
    let total_bytes = messages * payload_size as u64;

    if matches!(format, OutputFormat::Human) {
        let mode = if no_ack { " (no-ack)" } else { "" };
        println!(
            "bench pub: {messages} messages, {payload_size} bytes each, concurrency={concurrency}{mode}"
        );
    }

    let msgs_per_worker = messages / concurrency as u64;
    let remainder = messages % concurrency as u64;

    let start = Instant::now();
    let mut latencies: Vec<Duration> =
        Vec::with_capacity(if no_ack { 0 } else { messages as usize });

    if concurrency == 1 {
        // Single-connection fast path.
        let client_config = build_client_config(&config, ns)?;
        let mut client = FafbClient::connect(&client_config)
            .await
            .map_err(|e| CliError::Error(format!("connection failed: {e}")))?;

        if no_ack {
            // Fire-and-forget: pipeline publishes without waiting for acks.
            let mut last_progress = Instant::now();
            for i in 0..messages {
                client
                    .send_frame(Frame::Publish {
                        topic: topic.clone(),
                        headers: Headers::new(),
                        ttl_ms: 0,
                        payload: payload.clone(),
                    })
                    .await
                    .map_err(|e| CliError::Error(format!("publish failed: {e}")))?;

                if matches!(format, OutputFormat::Human)
                    && last_progress.elapsed() >= PROGRESS_INTERVAL
                {
                    eprintln!("  sent {}/{messages}...", i + 1);
                    last_progress = Instant::now();
                }
            }
            // Drain incoming PublishAcks so the server doesn't see a broken pipe
            // when writing ack responses. This ensures all publishes are processed.
            let drain_deadline = Instant::now() + Duration::from_secs(5);
            let mut acks_received = 0u64;
            while acks_received < messages && Instant::now() < drain_deadline {
                match tokio::time::timeout(Duration::from_millis(100), client.recv_frame()).await {
                    Ok(Ok(_)) => acks_received += 1,
                    Ok(Err(_)) => break,
                    Err(_) => break, // timeout — no more acks pending
                }
            }
        } else {
            let mut last_progress = Instant::now();
            for i in 0..messages {
                let msg_start = Instant::now();
                client
                    .publish(topic.clone(), payload.clone(), Headers::new(), 0)
                    .await
                    .map_err(|e| CliError::Error(format!("publish failed: {e}")))?;
                latencies.push(msg_start.elapsed());

                if matches!(format, OutputFormat::Human)
                    && last_progress.elapsed() >= PROGRESS_INTERVAL
                {
                    eprintln!("  published {}/{messages}...", i + 1);
                    last_progress = Instant::now();
                }
            }
        }
        client.close().await.ok();
    } else {
        // Concurrent connections via JoinSet.
        let mut join_set = tokio::task::JoinSet::new();

        for i in 0..concurrency {
            let worker_msgs = msgs_per_worker + if (i as u64) < remainder { 1 } else { 0 };
            let cfg = config.clone();
            let ns_clone = ns.clone();
            let topic_clone = topic.clone();
            let payload_clone = payload.clone();

            join_set.spawn(async move {
                let client_config = build_client_config(&cfg, ns_clone)?;
                let mut client = FafbClient::connect(&client_config)
                    .await
                    .map_err(|e| CliError::Error(format!("connection failed: {e}")))?;

                if no_ack {
                    for _ in 0..worker_msgs {
                        client
                            .send_frame(Frame::Publish {
                                topic: topic_clone.clone(),
                                headers: Headers::new(),
                                ttl_ms: 0,
                                payload: payload_clone.clone(),
                            })
                            .await
                            .map_err(|e| CliError::Error(format!("publish failed: {e}")))?;
                    }
                    // Drain acks to ensure server processes all publishes.
                    let drain_deadline = Instant::now() + Duration::from_secs(5);
                    let mut acks = 0u64;
                    while acks < worker_msgs && Instant::now() < drain_deadline {
                        match tokio::time::timeout(Duration::from_millis(100), client.recv_frame())
                            .await
                        {
                            Ok(Ok(_)) => acks += 1,
                            Ok(Err(_)) => break,
                            Err(_) => break,
                        }
                    }
                    client.close().await.ok();
                    Ok::<Vec<Duration>, CliError>(Vec::new())
                } else {
                    let mut worker_latencies = Vec::with_capacity(worker_msgs as usize);
                    for _ in 0..worker_msgs {
                        let msg_start = Instant::now();
                        client
                            .publish(
                                topic_clone.clone(),
                                payload_clone.clone(),
                                Headers::new(),
                                0,
                            )
                            .await
                            .map_err(|e| CliError::Error(format!("publish failed: {e}")))?;
                        worker_latencies.push(msg_start.elapsed());
                    }
                    client.close().await.ok();
                    Ok(worker_latencies)
                }
            });
        }

        while let Some(result) = join_set.join_next().await {
            let worker_lats = result
                .map_err(|e| CliError::Error(format!("worker panicked: {e}")))?
                .map_err(|e: CliError| e)?;
            latencies.extend(worker_lats);
        }
    }

    let elapsed = start.elapsed();
    latencies.sort();

    let throughput_msgs = messages as f64 / elapsed.as_secs_f64();
    let throughput_bytes = total_bytes as f64 / elapsed.as_secs_f64();

    match format {
        OutputFormat::Json => {
            let mut json = serde_json::json!({
                "command": "bench pub",
                "messages": messages,
                "payload_size": payload_size,
                "concurrency": concurrency,
                "no_ack": no_ack,
                "elapsed_ms": elapsed.as_millis() as u64,
                "throughput_msg_per_sec": throughput_msgs as u64,
                "throughput_bytes_per_sec": throughput_bytes as u64,
            });
            if !no_ack {
                let stats = LatencyStats::from_sorted(&latencies);
                json["p50_us"] = (stats.p50.as_micros() as u64).into();
                json["p99_us"] = (stats.p99.as_micros() as u64).into();
                json["p999_us"] = (stats.p999.as_micros() as u64).into();
                json["max_us"] = (stats.max.as_micros() as u64).into();
            }
            println!("{json}");
        }
        _ => {
            println!(
                "{} {} messages in {:.2}s",
                "Published".green(),
                messages,
                elapsed.as_secs_f64()
            );
            println!(
                "  {} {:.0} msg/s ({}/s)",
                "Throughput:".dimmed(),
                throughput_msgs,
                crate::cli::format_bytes(throughput_bytes as u64)
            );
            if !no_ack {
                let stats = LatencyStats::from_sorted(&latencies);
                println!(
                    "  {} {}",
                    "p50 latency: ".dimmed(),
                    format_duration(stats.p50)
                );
                println!(
                    "  {} {}",
                    "p99 latency: ".dimmed(),
                    format_duration(stats.p99)
                );
                println!(
                    "  {} {}",
                    "p999 latency:".dimmed(),
                    format_duration(stats.p999)
                );
                println!(
                    "  {} {}",
                    "max latency: ".dimmed(),
                    format_duration(stats.max)
                );
            }
        }
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Subscribe benchmark
// ---------------------------------------------------------------------------

async fn run_sub_bench(
    config: CliConfig,
    topic_str: &str,
    group: Option<String>,
    from: Option<u64>,
    max_count: Option<u64>,
) -> Result<(), CliError> {
    let format = OutputFormat::auto_detect(config.json, false);
    let (parsed_ns, topic) = parse_topic(topic_str);
    let ns = resolve_namespace(parsed_ns, config.namespace.as_deref(), &config)
        .map_err(CliError::Error)?;

    let client_config = build_client_config(&config, ns)?;
    let mut client = FafbClient::connect(&client_config)
        .await
        .map_err(|e| CliError::Error(format!("connection failed: {e}")))?;

    // Subscribe — from offset if specified, otherwise head.
    let group = group.map(GroupName);
    let sub_frame = if let Some(offset) = from {
        Frame::SubscribeAt {
            topic: topic.clone(),
            offset: TopicOffset(offset),
            group: group.clone(),
        }
    } else {
        Frame::Subscribe {
            topic: topic.clone(),
            group: group.clone(),
        }
    };
    client
        .send_frame(sub_frame)
        .await
        .map_err(|e| CliError::Error(format!("subscribe failed: {e}")))?;

    // Wait for SubscribeAck.
    let ack = client
        .recv_frame()
        .await
        .map_err(|e| CliError::Error(format!("subscribe ack failed: {e}")))?;
    match ack {
        Frame::SubscribeAck { .. } => {}
        Frame::Error { message, .. } => {
            return Err(CliError::Error(format!("subscribe error: {message}")));
        }
        other => {
            return Err(CliError::Error(format!("unexpected frame: {other:?}")));
        }
    }

    if matches!(format, OutputFormat::Human) {
        let stop_hint = match max_count {
            Some(n) => format!("will stop after {n} messages"),
            None => "Ctrl+C to stop".into(),
        };
        println!("bench sub: consuming from {topic_str}, {stop_hint}...");
    }

    let start = Instant::now();
    let mut count: u64 = 0;
    let mut total_bytes: u64 = 0;
    let mut last_progress = Instant::now();

    loop {
        tokio::select! {
            result = client.recv_frame() => {
                match result {
                    Ok(Frame::Deliver { message_id, payload, .. }) => {
                        count += 1;
                        total_bytes += payload.len() as u64;

                        // Auto-ack.
                        client.send_frame(Frame::Ack { message_id }).await
                            .map_err(|e| CliError::Error(format!("ack failed: {e}")))?;

                        // Check count limit.
                        if let Some(limit) = max_count
                            && count >= limit
                        {
                            break;
                        }

                        // Progress output.
                        if matches!(format, OutputFormat::Human)
                            && last_progress.elapsed() >= PROGRESS_INTERVAL
                        {
                            eprintln!("  consumed {count} messages...");
                            last_progress = Instant::now();
                        }
                    }
                    Ok(Frame::GoingAway { .. }) => break,
                    Ok(_) => {} // ignore other frames
                    Err(_) => break, // connection closed or error
                }
            }
            _ = tokio::signal::ctrl_c() => {
                break;
            }
        }
    }

    let elapsed = start.elapsed();
    let throughput_msgs = if elapsed.as_secs_f64() > 0.0 {
        count as f64 / elapsed.as_secs_f64()
    } else {
        0.0
    };
    let throughput_bytes = if elapsed.as_secs_f64() > 0.0 {
        total_bytes as f64 / elapsed.as_secs_f64()
    } else {
        0.0
    };

    match format {
        OutputFormat::Json => {
            let json = serde_json::json!({
                "command": "bench sub",
                "messages": count,
                "bytes": total_bytes,
                "elapsed_ms": elapsed.as_millis() as u64,
                "throughput_msg_per_sec": throughput_msgs as u64,
                "throughput_bytes_per_sec": throughput_bytes as u64,
            });
            println!("{json}");
        }
        _ => {
            println!(
                "{} {count} messages in {:.2}s",
                "Consumed".green(),
                elapsed.as_secs_f64()
            );
            println!(
                "  {} {:.0} msg/s ({}/s)",
                "Throughput:".dimmed(),
                throughput_msgs,
                crate::cli::format_bytes(throughput_bytes as u64)
            );
            println!(
                "  {} {}",
                "Total:".dimmed(),
                crate::cli::format_bytes(total_bytes)
            );
        }
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Latency benchmark
// ---------------------------------------------------------------------------

async fn run_latency_bench(
    config: CliConfig,
    topic_str: &str,
    messages: u64,
    warmup: u64,
) -> Result<(), CliError> {
    let format = OutputFormat::auto_detect(config.json, false);
    let (parsed_ns, topic) = parse_topic(topic_str);
    let ns = resolve_namespace(parsed_ns, config.namespace.as_deref(), &config)
        .map_err(CliError::Error)?;

    // Use two connections: one to subscribe, one to publish.
    let client_config = build_client_config(&config, ns)?;

    let mut sub_client = FafbClient::connect(&client_config)
        .await
        .map_err(|e| CliError::Error(format!("sub connection failed: {e}")))?;

    let mut pub_client = FafbClient::connect(&client_config)
        .await
        .map_err(|e| CliError::Error(format!("pub connection failed: {e}")))?;

    // Subscribe from head.
    sub_client
        .send_frame(Frame::Subscribe {
            topic: topic.clone(),
            group: None,
        })
        .await
        .map_err(|e| CliError::Error(format!("subscribe failed: {e}")))?;

    let ack = sub_client
        .recv_frame()
        .await
        .map_err(|e| CliError::Error(format!("subscribe ack failed: {e}")))?;
    match ack {
        Frame::SubscribeAck { .. } => {}
        Frame::Error { message, .. } => {
            return Err(CliError::Error(format!("subscribe error: {message}")));
        }
        _ => {}
    }

    let total = warmup + messages;
    if matches!(format, OutputFormat::Human) {
        if warmup > 0 {
            println!(
                "bench latency: {warmup} warmup + {messages} measured messages ({total} total)..."
            );
        } else {
            println!("bench latency: measuring round-trip for {messages} messages...");
        }
    }

    let payload = Bytes::from(vec![b'L'; 64]);
    let mut latencies = Vec::with_capacity(messages as usize);

    for i in 0..total {
        let msg_start = Instant::now();

        // Publish.
        pub_client
            .publish(topic.clone(), payload.clone(), Headers::new(), 0)
            .await
            .map_err(|e| CliError::Error(format!("publish failed: {e}")))?;

        // Wait for the Deliver on the subscriber.
        let timeout = Duration::from_secs(config.timeout_secs);
        let frame = tokio::time::timeout(timeout, sub_client.recv_frame())
            .await
            .map_err(|_| CliError::Timeout("timed out waiting for message delivery".into()))?
            .map_err(|e| CliError::Error(format!("recv failed: {e}")))?;

        match frame {
            Frame::Deliver { message_id, .. } => {
                // Only record after warmup.
                if i >= warmup {
                    latencies.push(msg_start.elapsed());
                }
                // Ack.
                sub_client
                    .send_frame(Frame::Ack { message_id })
                    .await
                    .map_err(|e| CliError::Error(format!("ack failed: {e}")))?;
            }
            Frame::Error { message, .. } => {
                return Err(CliError::Error(format!("delivery error: {message}")));
            }
            other => {
                return Err(CliError::Error(format!("unexpected frame: {other:?}")));
            }
        }
    }

    pub_client.close().await.ok();
    sub_client.close().await.ok();

    latencies.sort();
    let stats = LatencyStats::from_sorted(&latencies);

    match format {
        OutputFormat::Json => {
            let mut json = serde_json::json!({
                "command": "bench latency",
                "messages": messages,
                "p50_us": stats.p50.as_micros() as u64,
                "p99_us": stats.p99.as_micros() as u64,
                "p999_us": stats.p999.as_micros() as u64,
                "max_us": stats.max.as_micros() as u64,
            });
            if warmup > 0 {
                json["warmup"] = warmup.into();
            }
            println!("{json}");
        }
        _ => {
            println!("{} {messages} messages", "Latency benchmark:".green());
            println!("  {} {}", "p50: ".dimmed(), format_duration(stats.p50));
            println!("  {} {}", "p99: ".dimmed(), format_duration(stats.p99));
            println!("  {} {}", "p999:".dimmed(), format_duration(stats.p999));
            println!("  {} {}", "max: ".dimmed(), format_duration(stats.max));
        }
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Latency stats helper
// ---------------------------------------------------------------------------

struct LatencyStats {
    p50: Duration,
    p99: Duration,
    p999: Duration,
    max: Duration,
}

impl LatencyStats {
    /// Compute percentile stats from a pre-sorted slice of durations.
    fn from_sorted(sorted: &[Duration]) -> Self {
        if sorted.is_empty() {
            return Self {
                p50: Duration::ZERO,
                p99: Duration::ZERO,
                p999: Duration::ZERO,
                max: Duration::ZERO,
            };
        }
        let n = sorted.len();
        Self {
            p50: sorted[n * 50 / 100],
            p99: sorted[n * 99 / 100],
            p999: sorted[(n * 999 / 1000).min(n - 1)],
            max: sorted[n - 1],
        }
    }
}

/// Format a duration as a human-readable string.
fn format_duration(d: Duration) -> String {
    let us = d.as_micros();
    if us < 1_000 {
        format!("{us}us")
    } else if us < 1_000_000 {
        format!("{:.2}ms", us as f64 / 1_000.0)
    } else {
        format!("{:.2}s", d.as_secs_f64())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn latency_stats_empty() {
        let stats = LatencyStats::from_sorted(&[]);
        assert_eq!(stats.p50, Duration::ZERO);
        assert_eq!(stats.max, Duration::ZERO);
    }

    #[test]
    fn latency_stats_single() {
        let samples = vec![Duration::from_millis(5)];
        let stats = LatencyStats::from_sorted(&samples);
        assert_eq!(stats.p50, Duration::from_millis(5));
        assert_eq!(stats.max, Duration::from_millis(5));
    }

    #[test]
    fn latency_stats_percentiles() {
        let samples: Vec<Duration> = (1..=100).map(|i| Duration::from_micros(i * 10)).collect();
        let stats = LatencyStats::from_sorted(&samples);
        // Index = n * percentile / 100, 0-indexed: p50 → index 50 → value 510
        assert_eq!(stats.p50, Duration::from_micros(510));
        assert_eq!(stats.p99, Duration::from_micros(1000));
        assert_eq!(stats.max, Duration::from_micros(1000));
    }

    #[test]
    fn format_duration_microseconds() {
        assert_eq!(format_duration(Duration::from_micros(42)), "42us");
    }

    #[test]
    fn format_duration_milliseconds() {
        assert_eq!(format_duration(Duration::from_micros(1500)), "1.50ms");
    }

    #[test]
    fn format_duration_seconds() {
        assert_eq!(format_duration(Duration::from_secs(2)), "2.00s");
    }
}
