// Minimal implementation resembling Python's `tqdm`, Highly inspired by tinygrad's impl[1]
// [1] https://github.com/tinygrad/tinygrad/blob/7d886260687584e1ee6ac0e16207eb5b49c516ff/tinygrad/helpers.py#460

use std::{
    fmt::Write as _,
    io::{self, Write},
    time::Instant,
};

pub struct Tqdm<I> {
    it: I,
    desc: String,
    disable: bool,
    unit: &'static str,
    unit_scale: bool,
    total: usize,
    rate: usize,
    st: Instant,
    last_draw_at: Instant,
    last_draw_n: usize,
    rate_ema: f64,
    i: usize,
    n: usize,
    skip: usize,
    last_rendered_n: Option<usize>,
}

impl<I> Tqdm<I> {
    pub fn new(it: I) -> Self {
        Self {
            it,
            desc: String::new(),
            disable: false,
            unit: "it",
            unit_scale: false,
            total: 0,
            rate: 100,
            st: Instant::now(),
            last_draw_at: Instant::now(),
            last_draw_n: 0,
            rate_ema: 0.0,
            i: usize::MAX,
            n: 0,
            skip: 1,
            last_rendered_n: None,
        }
    }
    pub fn desc(mut self, s: impl AsRef<str>) -> Self {
        let s = s.as_ref();
        self.desc = if s.is_empty() {
            String::new()
        } else {
            format!("{s}: ")
        };
        self
    }
    pub fn disable(mut self, v: bool) -> Self {
        self.disable = v;
        self
    }
    pub fn unit(mut self, u: &'static str) -> Self {
        self.unit = u;
        self
    }
    pub fn unit_scale(mut self, v: bool) -> Self {
        self.unit_scale = v;
        self
    }
    pub fn total(mut self, t: usize) -> Self {
        self.total = t;
        self
    }
    pub fn rate(mut self, r: usize) -> Self {
        self.rate = r.max(1);
        self
    }

    pub fn update(&mut self, dn: usize, close: bool) {
        self.n += dn;
        self.i = self.i.wrapping_add(1);
        if self.disable || (!close && self.i % self.skip != 0) {
            return;
        }

        if close && self.last_rendered_n == Some(self.n) {
            let mut e = io::stderr().lock();
            let _ = e.write_all(b"\n");
            let _ = e.flush();
            return;
        }

        let elapsed = self.st.elapsed().as_secs_f64();
        let cols = term_cols().unwrap_or(80);

        if elapsed > 0.0 && self.i != 0 {
            let ips = (self.i as f64) / elapsed;
            if ips > self.rate as f64 {
                self.skip = ((ips as usize) / self.rate).max(1);
            }
        }

        let prog = if self.total != 0 {
            (self.n as f64) / (self.total as f64)
        } else {
            0.0
        };

        let mut suf = String::new();
        let prog_text = if self.unit_scale {
            let a = si(self.n as f64);
            if self.total != 0 {
                format!("{a}/{}", si(self.total as f64))
            } else {
                format!("{a}{}", self.unit)
            }
        } else {
            if self.total != 0 {
                format!("{}/{}", self.n, self.total)
            } else {
                format!("{}{}", self.n, self.unit)
            }
        };

        let _ = write!(suf, "{prog_text} [{}", hms(elapsed));
        if self.total != 0 {
            if self.n != 0 {
                let eta = elapsed / prog - elapsed;
                let _ = write!(suf, "<{}", hms(eta.max(0.0)));
            } else {
                suf.push_str("<?");
            }
        }
        let it_text = if self.n != 0 && elapsed > 0.0 {
            let draw_elapsed = self.last_draw_at.elapsed().as_secs_f64();
            if draw_elapsed > 0.0 && self.n >= self.last_draw_n {
                let dn = self.n - self.last_draw_n;
                let inst_rate = (dn as f64) / draw_elapsed;
                self.rate_ema = if self.rate_ema > 0.0 {
                    0.3 * inst_rate + 0.7 * self.rate_ema
                } else {
                    inst_rate
                };
            }
            let avg_rate = self.n as f64 / elapsed;
            let shown_rate = if self.rate_ema > 0.0 {
                self.rate_ema.min(avg_rate * 1.25)
            } else {
                avg_rate
            };
            if self.unit_scale {
                si(shown_rate)
            } else {
                format!("{shown_rate:5.2}")
            }
        } else {
            "?".into()
        };
        let _ = write!(
            suf,
            ", {it_text}{}{}/s]",
            if self.unit_scale { "" } else { "" },
            self.unit
        );

        let fixed = self.desc.len() + 3 + 2 + 2 + suf.len();
        let sz = (cols.saturating_sub(fixed)).max(1);

        let mut bar = String::new();
        bar.push_str(&self.desc);

        if self.total != 0 {
            let pct = (100.0 * prog).clamp(0.0, 100.0);
            let _ = write!(bar, "{pct:3.0}%|");
            let num = (sz as f64) * prog.clamp(0.0, 1.0);
            let full = num.floor() as usize;
            bar.extend(std::iter::repeat('█').take(full.min(sz)));
            let frac = ((num * 8.0) as usize) % 8;
            if full < sz && frac != 0 {
                bar.push([' ', '▏', '▎', '▍', '▌', '▋', '▊', '▉'][frac]);
            }
            while bar.len() < 1 + self.desc.len() + 4 + sz {
                bar.push(' ');
            }
            bar.push_str("| ");
        }

        bar.push_str(&suf);
        truncate_to_chars(&mut bar, cols);

        let mut e = io::stderr().lock();
        let _ = write!(e, "\r\x1b[K{bar}");
        let _ = e.write_all(if close { b"\n" } else { b"" });
        let _ = e.flush();
        self.last_draw_at = Instant::now();
        self.last_draw_n = self.n;
        self.last_rendered_n = Some(self.n);
    }

    pub fn write(s: impl AsRef<str>) {
        let mut e = io::stderr().lock();
        let _ = write!(e, "\r\x1b[K{}", s.as_ref());
        let _ = e.flush();
    }
}

impl<I: Iterator> Iterator for Tqdm<I> {
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let x = self.it.next()?;
        self.update(1, false);
        Some(x)
    }
}

impl<I> Drop for Tqdm<I> {
    fn drop(&mut self) {
        self.update(0, true);
    }
}

pub fn trange(n: usize) -> Tqdm<std::ops::Range<usize>> {
    Tqdm::new(0..n).total(n)
}

fn term_cols() -> Option<usize> {
    if let Ok(s) = std::env::var("COLUMNS") {
        if let Ok(n) = s.parse() {
            return Some(n);
        }
    }
    #[cfg(unix)]
    unsafe {
        let mut ws: libc::winsize = std::mem::zeroed();
        if libc::ioctl(2, libc::TIOCGWINSZ, &mut ws) == 0 && ws.ws_col != 0 {
            return Some(ws.ws_col as usize);
        }
    }
    None
}

fn hms(t: f64) -> String {
    if !t.is_finite() || t <= 0.0 {
        return "00:00".into();
    }
    let mut s = t as u64;
    let h = s / 3600;
    s %= 3600;
    let m = s / 60;
    s %= 60;
    if h != 0 {
        format!("{h}:{m:02}:{s:02}")
    } else {
        format!("{m:02}:{s:02}")
    }
}

fn si(x: f64) -> String {
    if !x.is_finite() || x <= 0.0 {
        return "0.00".into();
    }
    const SUF: &[u8] = b" kMGTPEZY";
    let mut g = (x.ln() / 1000f64.ln()).floor() as i32;
    if g < 0 {
        g = 0;
    }
    if g as usize >= SUF.len() {
        g = (SUF.len() - 1) as i32;
    }
    let mut v = x / 1000f64.powi(g);
    let prec =
        (3.0 - 3.0 * ((x.ln() / 1000f64.ln()) - (x.ln() / 1000f64.ln()).floor())).max(0.0) as usize;
    let mut s = format!("{:.*}", prec.min(3), v);
    if s.len() > 4 {
        s.truncate(4);
    }
    while s.ends_with('.') {
        s.pop();
    }
    if s == "1000" && (g as usize) + 1 < SUF.len() {
        v = x / 1000f64.powi(g + 1);
        s = format!("{:.3}", v);
        if s.len() > 4 {
            s.truncate(4);
        }
        while s.ends_with('.') {
            s.pop();
        }
        s.push(' ');
        s.push(SUF[(g as usize) + 1] as char);
        s.trim().into()
    } else {
        let c = SUF[g as usize] as char;
        if c != ' ' {
            s.push(c);
        }
        s
    }
}

fn truncate_to_chars(s: &mut String, max_chars: usize) {
    if s.chars().count() <= max_chars {
        return;
    }
    let byte_idx = s
        .char_indices()
        .nth(max_chars)
        .map(|(idx, _)| idx)
        .unwrap_or(s.len());
    s.truncate(byte_idx);
}
