use std::{
    io::{self, Write},
    process::ExitCode,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread,
    time::{Duration, Instant},
};

use clap::Parser;
use irtt_cli::{
    format_event, format_human_event_with_options, CliArgs, HumanOutputOptions, OutputMode,
};
use irtt_client::{Client, ClientEvent, OpenOutcome, RecvBudget};

#[cfg(test)]
use irtt_client::ClientTimestamp;

#[cfg(feature = "stats")]
use irtt_stats::{StatsCollector, StatsConfig};

const RECV_BUDGET: RecvBudget = RecvBudget { max_packets: 16 };
const MAX_FINAL_DRAIN: Duration = Duration::from_secs(30);
const IDLE_SLEEP: Duration = Duration::from_millis(5);
const MAX_SLEEP: Duration = Duration::from_millis(20);

fn main() -> ExitCode {
    let args = CliArgs::parse();
    let shutdown_requested = Arc::new(AtomicBool::new(false));
    if let Err(err) = install_signal_handler(Arc::clone(&shutdown_requested)) {
        eprintln!("irtt-rs: failed to install signal handler: {err}");
        return ExitCode::FAILURE;
    }

    match run(args, shutdown_requested.as_ref()) {
        Ok(()) => ExitCode::SUCCESS,
        Err(err) => {
            eprintln!("irtt-rs: {err}");
            ExitCode::FAILURE
        }
    }
}

fn install_signal_handler(shutdown_requested: Arc<AtomicBool>) -> Result<(), ctrlc::Error> {
    ctrlc::set_handler(move || {
        shutdown_requested.store(true, Ordering::Relaxed);
    })
}

fn run(args: CliArgs, shutdown_requested: &AtomicBool) -> Result<(), Box<dyn std::error::Error>> {
    let mode = args.output;
    let continuous = args.is_continuous();
    let mut stdout = io::LineWriter::new(io::stdout().lock());
    #[cfg(feature = "stats")]
    let mut stats = StatsCollector::new(stats_config(continuous));
    let mut output = EventOutput {
        mode,
        human_options: HumanOutputOptions {
            verbose: args.verbose,
        },
        print_finite_summary: !continuous,
        out: &mut stdout,
        #[cfg(feature = "stats")]
        stats: &mut stats,
    };

    if is_shutdown_requested(shutdown_requested) {
        return Ok(());
    }

    let mut client = Client::connect(args.to_client_config())?;

    if is_shutdown_requested(shutdown_requested) {
        return Ok(());
    }

    let open = client.open()?;
    output.print_event(open_event(&open))?;

    let mut interrupted = false;
    while should_continue_run(continuous, client.next_send_deadline(), shutdown_requested) {
        if should_send_probe(client.next_send_deadline(), shutdown_requested) {
            let events = client.send_probe()?;
            output.print_events(&events)?;
        }

        let events = client.recv_available(RECV_BUDGET)?;
        output.print_events(&events)?;

        let events = client.poll_timeouts()?;
        output.print_events(&events)?;

        if is_shutdown_requested(shutdown_requested) {
            interrupted = true;
            break;
        }

        sleep_until_next_send(client.next_send_deadline());
    }
    interrupted |= is_shutdown_requested(shutdown_requested);

    if interrupted {
        eprintln!("interrupted, closing session...");
    }

    if should_drain_final(continuous, interrupted) {
        drain_final_replies(&mut client, &mut output)?;
    }

    let events = client.poll_timeouts()?;
    output.print_events(&events)?;

    let events = client.close()?;
    output.print_events(&events)?;
    output.print_summary()?;
    output.out.flush()?;
    Ok(())
}

fn is_shutdown_requested(shutdown_requested: &AtomicBool) -> bool {
    shutdown_requested.load(Ordering::Relaxed)
}

fn should_continue_run(
    continuous: bool,
    next_send_deadline: Option<Instant>,
    shutdown_requested: &AtomicBool,
) -> bool {
    !is_shutdown_requested(shutdown_requested) && (continuous || next_send_deadline.is_some())
}

fn should_send_probe(next_send_deadline: Option<Instant>, shutdown_requested: &AtomicBool) -> bool {
    !is_shutdown_requested(shutdown_requested)
        && next_send_deadline.is_some_and(|deadline| deadline <= Instant::now())
}

fn should_drain_final(continuous: bool, interrupted: bool) -> bool {
    !continuous || interrupted
}

fn open_event(outcome: &OpenOutcome) -> &ClientEvent {
    match outcome {
        OpenOutcome::Started { event, .. } | OpenOutcome::NoTestCompleted { event, .. } => event,
    }
}

fn drain_final_replies<W: Write>(
    client: &mut Client,
    output: &mut EventOutput<'_, W>,
) -> Result<(), Box<dyn std::error::Error>> {
    let deadline = Instant::now() + final_drain_duration(client.probe_timeout());
    loop {
        let mut printed = false;

        let events = client.recv_available(RECV_BUDGET)?;
        printed |= !events.is_empty();
        output.print_events(&events)?;

        let events = client.poll_timeouts()?;
        printed |= !events.is_empty();
        output.print_events(&events)?;

        if client.is_run_complete() || Instant::now() >= deadline {
            break;
        }

        if !printed {
            thread::sleep(IDLE_SLEEP);
        }
    }
    Ok(())
}

fn sleep_until_next_send(deadline: Option<Instant>) {
    let sleep_for = match deadline {
        Some(deadline) => deadline
            .saturating_duration_since(Instant::now())
            .min(MAX_SLEEP),
        None => Duration::from_millis(1),
    };
    if !sleep_for.is_zero() {
        thread::sleep(sleep_for);
    }
}

struct EventOutput<'a, W: Write> {
    mode: irtt_cli::OutputMode,
    human_options: HumanOutputOptions,
    print_finite_summary: bool,
    out: &'a mut W,
    #[cfg(feature = "stats")]
    stats: &'a mut StatsCollector,
}

impl<W: Write> EventOutput<'_, W> {
    fn print_events(&mut self, events: &[ClientEvent]) -> io::Result<()> {
        for event in events {
            self.print_event(event)?;
        }
        Ok(())
    }

    fn print_event(&mut self, event: &ClientEvent) -> io::Result<()> {
        #[cfg(feature = "stats")]
        let stats_update = self.stats.process(event);

        #[cfg(not(feature = "stats"))]
        let stats_update = irtt_cli::HumanEventStats::default();

        let line =
            if self.mode == OutputMode::Human && !matches!(event, ClientEvent::EchoSent { .. }) {
                Some(format_human_event_with_options(
                    event,
                    Some(stats_update.into()),
                    self.human_options,
                ))
            } else {
                format_event(event, self.mode)
            };

        if let Some(line) = line {
            writeln!(self.out, "{line}")?;
        }
        Ok(())
    }

    fn print_summary(&mut self) -> io::Result<()> {
        if !self.print_finite_summary || !self.mode.prints_summary() {
            return Ok(());
        }

        #[cfg(feature = "stats")]
        {
            write!(
                self.out,
                "{}",
                irtt_cli::summary::format_summary_with_options(
                    &self.stats.snapshot(),
                    irtt_cli::summary::SummaryFormatOptions {
                        verbose: self.human_options.verbose,
                    },
                )
            )?;
        }

        Ok(())
    }
}

#[cfg(feature = "stats")]
fn stats_config(continuous: bool) -> StatsConfig {
    if continuous {
        StatsConfig::continuous()
    } else {
        StatsConfig::finite()
    }
}

fn final_drain_duration(probe_timeout: Duration) -> Duration {
    probe_timeout.min(MAX_FINAL_DRAIN)
}

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

    #[test]
    fn shutdown_flag_stops_run_loop() {
        let shutdown = AtomicBool::new(false);
        assert!(should_continue_run(true, Some(Instant::now()), &shutdown));

        shutdown.store(true, Ordering::Relaxed);
        assert!(!should_continue_run(true, Some(Instant::now()), &shutdown));
    }

    #[test]
    fn shutdown_flag_suppresses_due_probe_send() {
        let shutdown = AtomicBool::new(false);
        let due = Instant::now() - Duration::from_millis(1);
        assert!(should_send_probe(Some(due), &shutdown));

        shutdown.store(true, Ordering::Relaxed);
        assert!(!should_send_probe(Some(due), &shutdown));
    }

    #[test]
    fn interrupted_continuous_run_uses_final_drain_before_close() {
        assert!(should_drain_final(true, true));
        assert!(!should_drain_final(true, false));
        assert!(should_drain_final(false, false));
    }
}

#[cfg(all(test, feature = "stats"))]
mod tests {
    use super::*;
    use irtt_client::{PacketMeta, RttSample, SignedDuration};
    use std::{
        net::{IpAddr, Ipv4Addr, SocketAddr},
        time::UNIX_EPOCH,
    };

    fn test_timestamp(offset: Duration) -> ClientTimestamp {
        ClientTimestamp {
            wall: UNIX_EPOCH + offset,
            mono: Instant::now() + offset,
        }
    }

    fn test_remote() -> SocketAddr {
        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 2112)
    }

    fn reply_event(seq: u32, rtt_us: u64) -> ClientEvent {
        let sent_at = test_timestamp(Duration::from_millis(seq as u64));
        let received_at = ClientTimestamp {
            wall: sent_at.wall + Duration::from_micros(rtt_us),
            mono: sent_at.mono + Duration::from_micros(rtt_us),
        };
        ClientEvent::EchoReply {
            seq,
            remote: test_remote(),
            sent_at,
            received_at,
            rtt: RttSample {
                raw: Duration::from_micros(rtt_us),
                adjusted: None,
                effective: Duration::from_micros(rtt_us),
                adjusted_signed: None,
                effective_signed: SignedDuration {
                    ns: i128::from(rtt_us) * 1_000,
                },
            },
            server_timing: None,
            one_way: None,
            received_stats: None,
            bytes: 64,
            packet_meta: PacketMeta::default(),
        }
    }

    #[test]
    fn output_helper_streams_and_collects_events() {
        let sent_at = test_timestamp(Duration::from_secs(1));
        let received_at = test_timestamp(Duration::from_secs(1) + Duration::from_micros(1200));
        let events = [
            ClientEvent::EchoSent {
                seq: 1,
                remote: test_remote(),
                scheduled_at: sent_at.mono,
                sent_at,
                bytes: 64,
                send_call: Duration::from_micros(10),
                timer_error: Duration::ZERO,
            },
            ClientEvent::EchoReply {
                seq: 1,
                remote: test_remote(),
                sent_at,
                received_at,
                rtt: RttSample {
                    raw: Duration::from_micros(1200),
                    adjusted: None,
                    effective: Duration::from_micros(1200),
                    adjusted_signed: None,
                    effective_signed: SignedDuration { ns: 1_200_000 },
                },
                server_timing: None,
                one_way: None,
                received_stats: None,
                bytes: 64,
                packet_meta: PacketMeta::default(),
            },
        ];

        let mut stats = StatsCollector::new(StatsConfig::finite());
        let mut out = Vec::new();
        {
            let mut output = EventOutput {
                mode: irtt_cli::OutputMode::RttUs,
                human_options: HumanOutputOptions::default(),
                print_finite_summary: true,
                out: &mut out,
                stats: &mut stats,
            };
            output.print_events(&events).unwrap();
            output.print_summary().unwrap();
        }

        let rendered = String::from_utf8(out).unwrap();
        let summary = stats.snapshot();
        assert_eq!(rendered, "1200\n");
        assert_eq!(summary.packets.packets_sent, 1);
        assert_eq!(summary.packets.unique_replies, 1);
    }

    #[test]
    fn continuous_output_does_not_print_finite_summary() {
        let mut stats = StatsCollector::new(StatsConfig::continuous());
        let mut out = Vec::new();
        {
            let mut output = EventOutput {
                mode: irtt_cli::OutputMode::Human,
                human_options: HumanOutputOptions::default(),
                print_finite_summary: false,
                out: &mut out,
                stats: &mut stats,
            };
            output.print_summary().unwrap();
        }

        assert!(out.is_empty());
    }

    #[test]
    fn human_output_prints_ipdv_from_stats_update() {
        let events = [reply_event(0, 1200), reply_event(1, 1250)];
        let mut stats = StatsCollector::new(StatsConfig::finite());
        let mut out = Vec::new();
        {
            let mut output = EventOutput {
                mode: irtt_cli::OutputMode::Human,
                human_options: HumanOutputOptions::default(),
                print_finite_summary: false,
                out: &mut out,
                stats: &mut stats,
            };
            output.print_events(&events).unwrap();
        }

        let rendered = String::from_utf8(out).unwrap();
        let mut lines = rendered.lines();
        assert!(lines.next().unwrap().contains("ipdv=n/a"));
        assert!(lines.next().unwrap().contains("ipdv=50.0µs"));
    }

    #[test]
    fn continuous_mode_uses_continuous_stats_config() {
        let mut collector = StatsCollector::new(stats_config(true));
        for seq in 0..5000_u32 {
            let sent_at = test_timestamp(Duration::from_micros(seq as u64));
            let received_at = test_timestamp(Duration::from_micros(seq as u64 + 1));
            collector.process(&ClientEvent::EchoReply {
                seq,
                remote: test_remote(),
                sent_at,
                received_at,
                rtt: RttSample {
                    raw: Duration::from_micros(1),
                    adjusted: None,
                    effective: Duration::from_micros(1),
                    adjusted_signed: None,
                    effective_signed: SignedDuration { ns: 1_000 },
                },
                server_timing: None,
                one_way: None,
                received_stats: None,
                bytes: 64,
                packet_meta: PacketMeta::default(),
            });
        }

        assert_eq!(collector.snapshot().rtt.primary.median_ns, None);
    }

    #[test]
    fn final_drain_uses_capped_probe_timeout() {
        assert_eq!(
            final_drain_duration(Duration::from_secs(4)),
            Duration::from_secs(4)
        );
        assert_eq!(
            final_drain_duration(Duration::from_secs(60)),
            Duration::from_secs(30)
        );
    }
}