// SPDX-License-Identifier: MIT

use std::{
    io::{self, IsTerminal, Write},
    sync::{Arc, Mutex},
};

use indicatif::{MultiProgress, ProgressBar, ProgressStyle};

/// The progress display currently owning the terminal draw area, if any. Log
/// lines must be written through its `suspend` so they do not corrupt the
/// in-place redraw (stray stderr writes make indicatif reprint the whole
/// group below the old frame).
static ACTIVE_DRAW: Mutex<Option<ActiveDraw>> = Mutex::new(None);

#[derive(Clone)]
enum ActiveDraw {
    Single(ProgressBar),
    Multi(MultiProgress),
}

fn set_active_draw(active: Option<ActiveDraw>) {
    if let Ok(mut guard) = ACTIVE_DRAW.lock() {
        *guard = active;
    }
}

fn current_active_draw() -> Option<ActiveDraw> {
    ACTIVE_DRAW.lock().ok().and_then(|guard| guard.clone())
}

/// Run `f` with the active progress display suspended (bars cleared, draw
/// lock held). The Ctrl-C reporter prints its summary through this: writing
/// to stderr while download workers keep redrawing races with the next frame,
/// which can overwrite the summary and leave a stale stacked frame behind.
pub fn suspend_active_draw<T>(f: impl FnOnce() -> T) -> T {
    match current_active_draw() {
        Some(ActiveDraw::Multi(multi)) => multi.suspend(f),
        Some(ActiveDraw::Single(bar)) => bar.suspend(f),
        None => f(),
    }
}

/// `tracing` writer that routes log lines through the active progress
/// display's `suspend`, so warnings printed mid-transfer do not break the
/// progress redraw.
pub struct LogWriter;

pub struct LogWriterHandle;

impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for LogWriter {
    type Writer = LogWriterHandle;

    fn make_writer(&'a self) -> Self::Writer {
        LogWriterHandle
    }
}

impl Write for LogWriterHandle {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let write_all = || {
            let mut stderr = io::stderr();
            let written = stderr.write(buf)?;
            stderr.flush()?;
            Ok(written)
        };
        match current_active_draw() {
            Some(ActiveDraw::Multi(multi)) => multi.suspend(write_all),
            Some(ActiveDraw::Single(bar)) => bar.suspend(write_all),
            None => write_all(),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        io::stderr().flush()
    }
}

#[derive(Clone)]
pub struct ByteProgress {
    bar: Option<ProgressBar>,
}

impl ByteProgress {
    pub fn new(total: Option<u64>, quiet: bool, label: &str) -> Self {
        if quiet || !std::io::stderr().is_terminal() {
            return Self { bar: None };
        }

        let Some(total) = total else {
            return Self { bar: None };
        };

        let bar = ProgressBar::new(total);
        let style = ProgressStyle::with_template(
            "{msg} [{bar:40.cyan/blue}] {percent:>3}% {bytes}/{total_bytes} {bytes_per_sec} ETA {eta_precise}",
        )
        .unwrap_or_else(|_| ProgressStyle::default_bar())
        .progress_chars("=> ");
        bar.set_style(style);
        bar.set_message(label.to_owned());
        set_active_draw(Some(ActiveDraw::Single(bar.clone())));

        Self { bar: Some(bar) }
    }

    pub fn inc(&self, bytes: u64) {
        if let Some(bar) = &self.bar {
            bar.inc(bytes);
        }
    }

    pub fn set_position(&self, bytes: u64) {
        if let Some(bar) = &self.bar {
            bar.set_position(bytes);
        }
    }

    pub fn finish(&self) {
        if let Some(bar) = &self.bar {
            bar.finish_and_clear();
            set_active_draw(None);
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct SegmentProgressSpec {
    pub len: u64,
    pub initial: u64,
}

#[derive(Clone)]
pub struct SegmentedProgress {
    inner: Arc<SegmentedProgressInner>,
}

struct SegmentedProgressInner {
    bars: Option<SegmentedBars>,
    positions: Mutex<SegmentPositions>,
    total: Option<u64>,
    segment_totals: Vec<u64>,
}

impl Drop for SegmentedProgressInner {
    // Error paths drop the group without calling finish(); without this the
    // last frame stays on screen and every retry stacks a new copy below it.
    fn drop(&mut self) {
        if let Some(bars) = &self.bars {
            set_active_draw(None);
            for bar in &bars.segments {
                bar.finish_and_clear();
            }
            bars.main.finish_and_clear();
        }
    }
}

struct SegmentedBars {
    _multi: MultiProgress,
    main: ProgressBar,
    segments: Vec<ProgressBar>,
}

struct SegmentPositions {
    main: u64,
    segments: Vec<u64>,
}

impl SegmentedProgress {
    pub fn new(
        total: Option<u64>,
        quiet: bool,
        label: &str,
        segments: &[SegmentProgressSpec],
    ) -> Self {
        Self::new_with_options(total, quiet, label, segments, "conn", true)
    }

    pub fn new_with_segment_label(
        total: Option<u64>,
        quiet: bool,
        label: &str,
        segments: &[SegmentProgressSpec],
        segment_label: &str,
    ) -> Self {
        Self::new_with_options(total, quiet, label, segments, segment_label, false)
    }

    fn new_with_options(
        total: Option<u64>,
        quiet: bool,
        label: &str,
        segments: &[SegmentProgressSpec],
        segment_label: &str,
        segment_speed: bool,
    ) -> Self {
        let segment_totals = segments
            .iter()
            .map(|segment| segment.len)
            .collect::<Vec<_>>();
        let segment_positions = segments
            .iter()
            .zip(segment_totals.iter())
            .map(|(segment, total)| segment.initial.min(*total))
            .collect::<Vec<_>>();
        let main_position = segment_positions.iter().sum::<u64>();
        let bars = if quiet || !std::io::stderr().is_terminal() {
            None
        } else {
            total.map(|total| {
                segmented_bars(
                    total,
                    label,
                    &segment_totals,
                    &segment_positions,
                    segment_label,
                    segment_speed,
                )
            })
        };

        Self {
            inner: Arc::new(SegmentedProgressInner {
                bars,
                positions: Mutex::new(SegmentPositions {
                    main: main_position,
                    segments: segment_positions,
                }),
                total,
                segment_totals,
            }),
        }
    }

    pub fn inc(&self, index: usize, bytes: u64) {
        let Some(segment_total) = self.inner.segment_totals.get(index).copied() else {
            debug_assert!(false, "missing segmented progress index {index}");
            return;
        };
        let mut positions = self
            .inner
            .positions
            .lock()
            .expect("segmented progress mutex poisoned");
        let Some(segment_position) = positions.segments.get_mut(index) else {
            debug_assert!(false, "missing segmented progress index {index}");
            return;
        };
        let before = *segment_position;
        *segment_position = segment_position.saturating_add(bytes).min(segment_total);
        let advanced = *segment_position - before;
        positions.main = positions
            .main
            .saturating_add(advanced)
            .min(self.inner.total.unwrap_or(u64::MAX));
        drop(positions);

        if let Some(bars) = &self.inner.bars {
            bars.main.inc(advanced);
            if let Some(bar) = bars.segments.get(index) {
                bar.inc(advanced);
            }
        }
    }

    pub fn set_segment_position(&self, index: usize, bytes: u64) {
        let Some(segment_total) = self.inner.segment_totals.get(index).copied() else {
            debug_assert!(false, "missing segmented progress index {index}");
            return;
        };
        let mut positions = self
            .inner
            .positions
            .lock()
            .expect("segmented progress mutex poisoned");
        let position = bytes.min(segment_total);
        let Some(segment_position) = positions.segments.get_mut(index) else {
            debug_assert!(false, "missing segmented progress index {index}");
            return;
        };
        *segment_position = position;
        positions.main = positions
            .segments
            .iter()
            .sum::<u64>()
            .min(self.inner.total.unwrap_or(u64::MAX));
        let main = positions.main;
        drop(positions);

        if let Some(bars) = &self.inner.bars {
            bars.main.set_position(main);
            if let Some(bar) = bars.segments.get(index) {
                bar.set_position(position);
            }
        }
    }

    pub fn set_segment_message(&self, index: usize, message: String) {
        if let Some(bars) = &self.inner.bars {
            if let Some(bar) = bars.segments.get(index) {
                bar.set_message(message);
            }
        }
    }

    pub fn finish(&self) {
        if let Some(bars) = &self.inner.bars {
            set_active_draw(None);
            for bar in &bars.segments {
                bar.finish_and_clear();
            }
            bars.main.finish_and_clear();
        }
    }

    #[cfg(test)]
    fn snapshot(&self) -> (u64, Vec<u64>) {
        let positions = self
            .inner
            .positions
            .lock()
            .expect("segmented progress mutex poisoned");
        (positions.main, positions.segments.clone())
    }
}

fn segmented_bars(
    total: u64,
    label: &str,
    segment_totals: &[u64],
    segment_positions: &[u64],
    segment_label: &str,
    segment_speed: bool,
) -> SegmentedBars {
    let multi = MultiProgress::new();
    // Preset resume positions via with_position BEFORE the bar joins the draw
    // group: set_position after add() draws one frame whose rate estimator
    // counts the preset as an instantaneous transfer (absurd TiB/s speeds).
    let main = multi.add(
        ProgressBar::new(total)
            .with_style(main_style())
            .with_message(label.to_owned())
            .with_position(segment_positions.iter().sum::<u64>().min(total)),
    );
    main.reset_eta();
    main.tick();

    let mut segments = Vec::with_capacity(segment_totals.len());
    let mut previous = main.clone();
    for (index, total) in segment_totals.iter().copied().enumerate() {
        let prefix = if index + 1 == segment_totals.len() {
            "  └─"
        } else {
            "  ├─"
        };
        let initial = segment_positions
            .get(index)
            .copied()
            .unwrap_or(0)
            .min(total);
        let bar = multi.insert_after(
            &previous,
            ProgressBar::new(total)
                .with_style(segment_style(segment_speed))
                .with_prefix(prefix.to_owned())
                .with_message(format!("{segment_label} {}", index + 1))
                .with_position(initial),
        );
        bar.reset_eta();
        bar.tick();
        previous = bar.clone();
        segments.push(bar);
    }

    set_active_draw(Some(ActiveDraw::Multi(multi.clone())));
    SegmentedBars {
        _multi: multi,
        main,
        segments,
    }
}

fn main_style() -> ProgressStyle {
    ProgressStyle::with_template(
        "{msg} [{bar:40.cyan/blue}] {percent:>3}% {bytes}/{total_bytes} {bytes_per_sec} ETA {eta_precise}",
    )
    .unwrap_or_else(|_| ProgressStyle::default_bar())
    .progress_chars("=> ")
}

fn segment_style(show_speed: bool) -> ProgressStyle {
    let template = if show_speed {
        "{prefix} {msg} [{bar:24.white/blue}] {bytes}/{total_bytes} {bytes_per_sec}"
    } else {
        "{prefix} {msg} [{bar:24.white/blue}] {bytes}/{total_bytes}"
    };
    ProgressStyle::with_template(template)
        .unwrap_or_else(|_| ProgressStyle::default_bar())
        .progress_chars("=> ")
}

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

    #[test]
    fn segmented_progress_initializes_each_segment_position() {
        let progress = SegmentedProgress::new(
            Some(300),
            true,
            "example.bin",
            &[
                SegmentProgressSpec {
                    len: 100,
                    initial: 100,
                },
                SegmentProgressSpec {
                    len: 100,
                    initial: 40,
                },
                SegmentProgressSpec {
                    len: 100,
                    initial: 0,
                },
            ],
        );

        assert_eq!(progress.snapshot(), (140, vec![100, 40, 0]));
    }

    #[test]
    fn segmented_progress_increment_updates_main_and_one_segment() {
        let progress = SegmentedProgress::new(
            Some(20),
            true,
            "example.bin",
            &[
                SegmentProgressSpec {
                    len: 10,
                    initial: 2,
                },
                SegmentProgressSpec {
                    len: 10,
                    initial: 4,
                },
            ],
        );

        progress.inc(0, 3);
        progress.inc(1, 99);

        assert_eq!(progress.snapshot(), (15, vec![5, 10]));
    }

    #[test]
    fn segmented_progress_set_segment_position_recomputes_main() {
        let progress = SegmentedProgress::new(
            Some(30),
            true,
            "example.bin",
            &[
                SegmentProgressSpec {
                    len: 10,
                    initial: 5,
                },
                SegmentProgressSpec {
                    len: 20,
                    initial: 5,
                },
            ],
        );

        progress.set_segment_position(1, 12);

        assert_eq!(progress.snapshot(), (17, vec![5, 12]));
    }
}
