/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! # parsedate
//!
//! See [`HgTime`] and [`HgTime::parse`] for main features.

use std::ops::Add;
use std::ops::Range;
use std::ops::RangeInclusive;
use std::ops::Sub;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;

use chrono::prelude::*;
use chrono::Duration;
use chrono::LocalResult;

/// A simple time structure that matches hg's time representation.
///
/// Internally it's unixtime (in GMT), and offset (GMT -1 = +3600).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct HgTime {
    pub unixtime: i64,
    pub offset: i32,
}

const DEFAULT_FORMATS: [&str; 35] = [
    // mercurial/util.py defaultdateformats
    "%Y-%m-%dT%H:%M:%S", // the 'real' ISO8601
    "%Y-%m-%dT%H:%M",    //   without seconds
    "%Y-%m-%dT%H%M%S",   // another awful but legal variant without :
    "%Y-%m-%dT%H%M",     //   without seconds
    "%Y-%m-%d %H:%M:%S", // our common legal variant
    "%Y-%m-%d %H:%M",    //   without seconds
    "%Y-%m-%d %H%M%S",   // without :
    "%Y-%m-%d %H%M",     //   without seconds
    "%Y-%m-%d %I:%M:%S%p",
    "%Y-%m-%d %H:%M",
    "%Y-%m-%d %I:%M%p",
    "%a %b %d %H:%M:%S %Y",
    "%a %b %d %I:%M:%S%p %Y",
    "%a, %d %b %Y %H:%M:%S", //  GNU coreutils "/bin/date --rfc-2822"
    "%b %d %H:%M:%S %Y",
    "%b %d %I:%M:%S%p %Y",
    "%b %d %H:%M:%S",
    "%b %d %I:%M:%S%p",
    "%b %d %H:%M",
    "%b %d %I:%M%p",
    "%m-%d",
    "%m/%d",
    "%Y-%m-%d",
    "%m/%d/%y",
    "%m/%d/%Y",
    "%b",
    "%b %d",
    "%b %Y",
    "%b %d %Y",
    "%I:%M%p",
    "%H:%M",
    "%H:%M:%S",
    "%I:%M:%S%p",
    "%Y",
    "%Y-%m",
];

const INVALID_OFFSET: i32 = i32::max_value();
static DEFAULT_OFFSET: AtomicI32 = AtomicI32::new(INVALID_OFFSET);
static FORCED_NOW: AtomicU64 = AtomicU64::new(0); // test only

/// Call `TimeZone` methods on either `Local` (system default) or
/// TimeZone specified by `set_default_offset`.
///
/// This cannot be made as a regular function because:
/// - TimeZone<Local> and TimeZone<FixedOffset> are different types.
/// - TimeZone<T> cannot be made into a trait object.
macro_rules! with_local_timezone {
    (|$tz:ident| { $($expr: tt)* }) => {
        {
            let offset = DEFAULT_OFFSET.load(Ordering::Acquire);
            match FixedOffset::west_opt(offset) {
                Some($tz) => {
                    $($expr)*
                },
                None => {
                    let $tz = Local;
                    $($expr)*
                },
            }
        }
    };
}

impl HgTime {
    /// Supported Range. This is to be compatible with Python stdlib.
    ///
    /// The Python `datetime`  library can only express a limited range
    /// of dates (0001-01-01 to 9999-12-31). Its strftime requires
    /// year >= 1900.
    pub const RANGE: RangeInclusive<HgTime> = Self::min_value()..=Self::max_value();

    /// Return the current time with local timezone, or `None` if the timestamp
    /// is outside [`HgTime::RANGE`].
    ///
    /// The local timezone can be affected by `set_default_offset`.
    pub fn now() -> Option<Self> {
        let forced_now = FORCED_NOW.load(Ordering::Acquire);
        if forced_now == 0 {
            Self::try_from(Local::now()).ok().and_then(|mut t: HgTime| {
                let offset = DEFAULT_OFFSET.load(Ordering::Acquire);
                if is_valid_offset(offset) {
                    t.offset = offset;
                }
                t.bounded()
            })
        } else {
            Some(Self::from_compact_u64(forced_now))
        }
    }

    pub fn to_utc(self) -> DateTime<Utc> {
        let naive = NaiveDateTime::from_timestamp_opt(self.unixtime, 0).unwrap();
        DateTime::from_utc(naive, Utc)
    }

    /// Converts to `NaiveDateTime` with local timezone specified by `offset`.
    fn to_naive(self) -> NaiveDateTime {
        NaiveDateTime::from_timestamp_opt(self.unixtime - self.offset as i64, 0).unwrap()
    }

    /// Set as the faked "now". Useful for testing.
    ///
    /// This should only be used for testing.
    pub fn set_as_now_for_testing(self) {
        FORCED_NOW.store(self.to_lossy_compact_u64(), Ordering::SeqCst);
    }

    /// Parse a date string.
    ///
    /// Return `None` if it cannot be parsed.
    ///
    /// This function matches `mercurial.util.parsedate`, and can parse
    /// some additional forms like `2 days ago`.
    pub fn parse(date: &str) -> Option<Self> {
        match date {
            "now" => Self::now(),
            "today" => Self::now().and_then(|now| {
                Self::try_from(now.to_naive().date().and_hms_opt(0, 0, 0).unwrap()).ok()
            }),
            "yesterday" => Self::now().and_then(|now| {
                Self::try_from(
                    (now.to_naive().date() - Duration::days(1))
                        .and_hms_opt(0, 0, 0)
                        .unwrap(),
                )
                .ok()
            }),
            date if date.ends_with(" ago") => {
                let duration_str = &date[..date.len() - 4];
                duration_str
                    .parse::<humantime::Duration>()
                    .ok()
                    .and_then(|duration| Self::now().and_then(|n| n - duration.as_secs()))
            }
            _ => Self::parse_absolute(date, &default_date_lower),
        }
    }

    /// Parse a date string as a range.
    ///
    /// For example, `Apr 2000` covers range `Apr 1, 2000` to `Apr 30, 2000`.
    /// Also support more explicit ranges:
    /// - START to END
    /// - > START
    /// - < END
    pub fn parse_range(date: &str) -> Option<Range<Self>> {
        Self::parse_range_internal(date, true)
    }

    fn parse_range_internal(date: &str, support_to: bool) -> Option<Range<Self>> {
        match date {
            "now" => Self::now().and_then(|n| (n + 1).map(|m| n..m)),
            "today" => Self::now().and_then(|now| {
                let date = now.to_naive().date();
                let start = Self::try_from(date.and_hms_opt(0, 0, 0).unwrap());
                let end = Self::try_from(date.and_hms_opt(23, 59, 59).unwrap()).map(|t| t + 1);
                if let (Ok(start), Ok(Some(end))) = (start, end) {
                    Some(start..end)
                } else {
                    None
                }
            }),
            "yesterday" => Self::now().and_then(|now| {
                let date = now.to_naive().date() - Duration::days(1);
                let start = Self::try_from(date.and_hms_opt(0, 0, 0).unwrap());
                let end = Self::try_from(date.and_hms_opt(23, 59, 59).unwrap()).map(|t| t + 1);
                if let (Ok(start), Ok(Some(end))) = (start, end) {
                    Some(start..end)
                } else {
                    None
                }
            }),
            date if date.starts_with('>') => {
                Self::parse(&date[1..]).map(|start| start..Self::max_value())
            }
            date if date.starts_with("since ") => {
                Self::parse(&date[6..]).map(|start| start..Self::max_value())
            }
            date if date.starts_with('<') => Self::parse(&date[1..])
                .and_then(|end| end + 1)
                .map(|end| Self::min_value()..end),
            date if date.starts_with('-') => {
                // This does not really make much sense. But is supported by hg
                // (see 'hg help dates').
                Self::parse_range(&format!("since {} days ago", &date[1..]))
            }
            date if date.starts_with("before ") => {
                Self::parse(&date[7..]).map(|end| Self::min_value()..end)
            }
            date if support_to && date.contains(" to ") => {
                let phrases: Vec<_> = date.split(" to ").collect();
                if phrases.len() == 2 {
                    if let (Some(start), Some(end)) = (
                        Self::parse_range_internal(&phrases[0], false),
                        Self::parse_range_internal(&phrases[1], false),
                    ) {
                        Some(start.start..end.end)
                    } else {
                        None
                    }
                } else {
                    None
                }
            }
            _ => {
                let start = Self::parse_absolute(date, &default_date_lower);
                let end = Self::parse_absolute(date, &|c| default_date_upper(c, "31"))
                    .or_else(|| Self::parse_absolute(date, &|c| default_date_upper(c, "30")))
                    .or_else(|| Self::parse_absolute(date, &|c| default_date_upper(c, "29")))
                    .or_else(|| Self::parse_absolute(date, &|c| default_date_upper(c, "28")))
                    .and_then(|end| end + 1);
                if let (Some(start), Some(end)) = (start, end) {
                    Some(start..end)
                } else {
                    None
                }
            }
        }
    }

    /// Parse date in an absolute form.
    ///
    /// Return None if it cannot be parsed.
    ///
    /// `default_date` takes a format char, for example, `H`, and returns a
    /// default value of it.
    fn parse_absolute(date: &str, default_date: &dyn Fn(char) -> &'static str) -> Option<Self> {
        let date = date.trim();

        // Hg internal format. "unixtime offset"
        let parts: Vec<_> = date.split(' ').collect();
        if parts.len() == 2 {
            if let Ok(unixtime) = parts[0].parse() {
                if let Ok(offset) = parts[1].parse() {
                    if is_valid_offset(offset) {
                        return Self { unixtime, offset }.bounded();
                    }
                }
            }
        }

        // Normalize UTC timezone name to +0000. The parser does not know
        // timezone names.
        let date = if date.ends_with("GMT") || date.ends_with("UTC") {
            format!("{} +0000", &date[..date.len() - 3])
        } else {
            date.to_string()
        };
        let mut now = None; // cached, lazily calculated "now"

        // Try all formats!
        for naive_format in DEFAULT_FORMATS.iter() {
            // Fill out default fields.  See mercurial.util.strdate.
            // This makes it possible to parse partial dates like "month/day",
            // or "hour:minute", since the missing fields will be filled.
            let mut default_format = String::new();
            let mut date_with_defaults = date.clone();
            let mut use_now = false;
            for part in ["S", "M", "HI", "d", "mb", "Yy"] {
                if part
                    .chars()
                    .any(|ch| naive_format.contains(&format!("%{}", ch)))
                {
                    // For example, if the user specified "d" (day), but
                    // not other things, we should use 0 for "H:M:S", and
                    // "now" for "Y-m" (year, month).
                    use_now = true;
                } else {
                    let format_char = part.chars().next().unwrap();
                    default_format += &format!(" @%{}", format_char);
                    if use_now {
                        // For example, if the user only specified "month/day",
                        // then we should use the current "year", instead of
                        // year 0.
                        now = now.or_else(|| Self::now().map(|n| n.to_naive()));
                        match now {
                            Some(now) => {
                                date_with_defaults +=
                                    &format!(" @{}", now.format(&format!("%{}", format_char)))
                            }
                            None => return None,
                        }
                    } else {
                        // For example, if the user only specified
                        // "hour:minute", then we should use "second 0", instead
                        // of the current second.
                        date_with_defaults += " @";
                        date_with_defaults += default_date(format_char);
                    }
                }
            }

            // Try parse with timezone.
            // See https://docs.rs/chrono/0.4.9/chrono/format/strftime/index.html#specifiers
            let format = format!("{}%#z{}", naive_format, default_format);
            if let Ok(parsed) = DateTime::parse_from_str(&date_with_defaults, &format) {
                if let Ok(parsed) = parsed.try_into() {
                    return Some(parsed);
                }
            }

            // Without timezone.
            let format = format!("{}{}", naive_format, default_format);
            if let Ok(parsed) = NaiveDateTime::parse_from_str(&date_with_defaults, &format) {
                if let Ok(parsed) = parsed.try_into() {
                    return Some(parsed);
                }
            }
        }

        None
    }

    /// See [`HgTime::RANGE`] for details.
    pub const fn min_value() -> Self {
        Self {
            unixtime: -2208988800, // 1900-01-01 00:00:00
            offset: 0,
        }
    }

    /// See [`HgTime::RANGE`] for details.
    pub const fn max_value() -> Self {
        Self {
            unixtime: 253402300799, // 9999-12-31 23:59:59
            offset: 0,
        }
    }

    /// Return `None` if timestamp is out of [`HgTime::RANGE`].
    pub fn bounded(self) -> Option<Self> {
        if self < Self::min_value() || self > Self::max_value() {
            None
        } else {
            Some(self)
        }
    }
}

// Convert to compact u64.  Used by FORCED_NOW.
// For testing purpose only (no overflow checking).
impl HgTime {
    fn to_lossy_compact_u64(self) -> u64 {
        ((self.unixtime as u64) << 17) + (self.offset + 50401) as u64
    }

    fn from_compact_u64(value: u64) -> Self {
        let unixtime = (value as i64) >> 17;
        let offset = (((value & 0x1ffff) as i64) - 50401) as i32;
        Self { unixtime, offset }
    }
}

impl From<HgTime> for NaiveDateTime {
    fn from(time: HgTime) -> Self {
        time.to_naive()
    }
}

impl From<HgTime> for DateTime<Utc> {
    fn from(time: HgTime) -> Self {
        time.to_utc()
    }
}

impl Add<u64> for HgTime {
    type Output = Option<Self>;

    fn add(self, seconds: u64) -> Option<Self> {
        seconds.try_into().ok().and_then(|seconds| {
            self.unixtime.checked_add(seconds).and_then(|unixtime| {
                Self {
                    unixtime,
                    offset: self.offset,
                }
                .bounded()
            })
        })
    }
}

impl Sub<u64> for HgTime {
    type Output = Option<Self>;

    fn sub(self, seconds: u64) -> Option<Self> {
        seconds.try_into().ok().and_then(|seconds| {
            self.unixtime.checked_sub(seconds).and_then(|unixtime| {
                Self {
                    unixtime,
                    offset: self.offset,
                }
                .bounded()
            })
        })
    }
}

impl PartialOrd for HgTime {
    fn partial_cmp(&self, other: &HgTime) -> Option<std::cmp::Ordering> {
        self.unixtime.partial_cmp(&other.unixtime)
    }
}

impl<Tz: TimeZone> TryFrom<DateTime<Tz>> for HgTime {
    type Error = ();
    fn try_from(time: DateTime<Tz>) -> Result<Self, ()> {
        if time.timestamp() >= i64::min_value() {
            Self {
                unixtime: time.timestamp(),
                offset: time.offset().fix().utc_minus_local(),
            }
            .bounded()
            .ok_or(())
        } else {
            Err(())
        }
    }
}

impl<Tz: TimeZone> TryFrom<LocalResult<DateTime<Tz>>> for HgTime {
    type Error = ();
    fn try_from(time: LocalResult<DateTime<Tz>>) -> Result<Self, ()> {
        match time {
            LocalResult::Single(datetime) => HgTime::try_from(datetime),
            _ => Err(()),
        }
    }
}

impl TryFrom<NaiveDateTime> for HgTime {
    type Error = ();
    fn try_from(time: NaiveDateTime) -> Result<Self, ()> {
        with_local_timezone!(|tz| { tz.from_local_datetime(&time).try_into() })
    }
}

/// Change default offset (timezone).
pub fn set_default_offset(offset: i32) {
    DEFAULT_OFFSET.store(offset, Ordering::SeqCst);
}

fn is_valid_offset(offset: i32) -> bool {
    // UTC-12 to UTC+14.
    offset >= -50400 && offset <= 43200
}

/// Lower bound for default values in dates.
fn default_date_lower(format_char: char) -> &'static str {
    match format_char {
        'H' | 'M' | 'S' => "00",
        'm' | 'd' => "1",
        _ => unreachable!(),
    }
}

/// Upper bound. Assume a month has `N::to_static_str()` days.
fn default_date_upper(format_char: char, max_day: &'static str) -> &'static str {
    match format_char {
        'H' => "23",
        'M' | 'S' => "59",
        'm' => "12",
        'd' => max_day,
        _ => unreachable!(),
    }
}

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

    #[test]
    fn test_naive_local_roundtrip() {
        let now = Local::now().with_nanosecond(0).unwrap().naive_local();
        let hgtime: HgTime = now.try_into().unwrap();
        let now_again = hgtime.to_naive();
        assert_eq!(now, now_again);
    }

    #[test]
    fn test_parse_date() {
        // Test cases are mostly from test-parse-date.t.
        // Some variants were added.
        set_default_offset(7200);

        // t: parse date
        // d: parse date, compare with now, within expected range
        // The right side of assert_eq! is a string so it's autofix-able.

        assert_eq!(t("2006-02-01 13:00:30"), t("2006-02-01 13:00:30-0200"));
        assert_eq!(t("2006-02-01 13:00:30-0500"), "1138816830 18000");
        assert_eq!(t("2006-02-01 13:00:30 +05:00"), "1138780830 -18000");
        assert_eq!(t("2006-02-01 13:00:30Z"), "1138798830 0");
        assert_eq!(t("2006-02-01 13:00:30 GMT"), "1138798830 0");
        assert_eq!(t("2006-4-5 13:30"), "1144251000 7200");
        assert_eq!(t("1150000000 14400"), "1150000000 14400");
        assert_eq!(t("100000 1400000"), "fail");
        assert_eq!(t("1000000000 -16200"), "1000000000 -16200");
        assert_eq!(t("2006-02-01 1:00:30PM +0000"), "1138798830 0");

        assert_eq!(d("1:00:30PM +0000", Duration::days(1)), "0");
        assert_eq!(d("02/01", Duration::weeks(52)), "0");
        assert_eq!(d("today", Duration::days(1)), "0");
        assert_eq!(d("yesterday", Duration::days(2)), "0");

        // ISO8601
        assert_eq!(t("2016-07-27T12:10:21"), "1469628621 7200");
        assert_eq!(t("2016-07-27T12:10:21Z"), "1469621421 0");
        assert_eq!(t("2016-07-27T12:10:21+00:00"), "1469621421 0");
        assert_eq!(t("2016-07-27T121021Z"), "1469621421 0");
        assert_eq!(t("2016-07-27 12:10:21"), "1469628621 7200");
        assert_eq!(t("2016-07-27 12:10:21Z"), "1469621421 0");
        assert_eq!(t("2016-07-27 12:10:21+00:00"), "1469621421 0");
        assert_eq!(t("2016-07-27 121021Z"), "1469621421 0");

        // Months
        assert_eq!(t("Jan 2018"), "1514772000 7200");
        assert_eq!(t("Feb 2018"), "1517450400 7200");
        assert_eq!(t("Mar 2018"), "1519869600 7200");
        assert_eq!(t("Apr 2018"), "1522548000 7200");
        assert_eq!(t("May 2018"), "1525140000 7200");
        assert_eq!(t("Jun 2018"), "1527818400 7200");
        assert_eq!(t("Jul 2018"), "1530410400 7200");
        assert_eq!(t("Sep 2018"), "1535767200 7200");
        assert_eq!(t("Oct 2018"), "1538359200 7200");
        assert_eq!(t("Nov 2018"), "1541037600 7200");
        assert_eq!(t("Dec 2018"), "1543629600 7200");
        assert_eq!(t("Foo 2018"), "fail");

        // Extra tests not in test-parse-date.t
        assert_eq!(d("Jan", Duration::weeks(52)), "0");
        assert_eq!(d("Jan 1", Duration::weeks(52)), "0"); // 1 is not considered as "year 1"
        assert_eq!(d("4-26", Duration::weeks(52)), "0");
        assert_eq!(d("4/26", Duration::weeks(52)), "0");
        assert_eq!(t("4/26/2000"), "956714400 7200");
        assert_eq!(t("Apr 26 2000"), "956714400 7200");
        assert_eq!(t("2020"), "1577844000 7200"); // 2020 is considered as a "year"
        assert_eq!(t("2020 GMT"), "1577836800 0");
        assert_eq!(t("2020-12"), "1606788000 7200");
        assert_eq!(t("2020-13"), "fail");
        assert_eq!(t("1000"), "fail"); // year 1000 < HgTime::min_value()
        assert_eq!(t("1"), "fail");
        assert_eq!(t("0"), "fail");
        assert_eq!(t("100000000000000000 1400"), "fail");

        assert_eq!(t("Fri, 20 Sep 2019 12:15:13 -0700"), "1569006913 25200"); // date --rfc-2822
        assert_eq!(t("Fri, 20 Sep 2019 12:15:13"), "1568988913 7200");
    }

    #[test]
    fn test_parse_ago() {
        set_default_offset(7200);
        assert_eq!(d("10m ago", Duration::hours(1)), "0");
        assert_eq!(d("10 min ago", Duration::hours(1)), "0");
        assert_eq!(d("10 minutes ago", Duration::hours(1)), "0");
        assert_eq!(d("10 hours ago", Duration::days(1)), "0");
        assert_eq!(d("10 h ago", Duration::days(1)), "0");
        assert_eq!(t("9999999 years ago"), "fail");
    }

    #[test]
    fn test_parse_range() {
        set_default_offset(7200);

        assert_eq!(c("since 1 month ago", "now"), "contains");
        assert_eq!(c("since 1 month ago", "2 months ago"), "does not contain");
        assert_eq!(c("> 1 month ago", "2 months ago"), "does not contain");
        assert_eq!(c("< 1 month ago", "2 months ago"), "contains");
        assert_eq!(c("< 1 month ago", "now"), "does not contain");

        assert_eq!(c("-3", "now"), "contains");
        assert_eq!(c("-3", "2 days ago"), "contains");
        assert_eq!(c("-3", "4 days ago"), "does not contain");

        assert_eq!(c("2018", "2017-12-31 23:59:59"), "does not contain");
        assert_eq!(c("2018", "2018-1-1"), "contains");
        assert_eq!(c("2018", "2018-12-31 23:59:59"), "contains");
        assert_eq!(c("2018", "2019-1-1"), "does not contain");

        assert_eq!(c("2018-5-1 to 2018-6-2", "2018-4-30"), "does not contain");
        assert_eq!(c("2018-5-1 to 2018-6-2", "2018-5-30"), "contains");
        assert_eq!(c("2018-5-1 to 2018-6-2", "2018-6-30"), "does not contain");
        assert_eq!(c("2018-5 to 2018-6", "2018-5-1 0:0:0"), "contains");
        assert_eq!(c("2018-5 to 2018-6", "2018-6-30 23:59:59"), "contains");
        assert_eq!(c("2018-5 to 2018-6 to 2018-7", "2018-6-30"), "fail");

        // 0:0:0 yesterday to 23:59:59 today
        // Usually it's 48 hours. However it might be affected by DST.
        let range = HgTime::parse_range("yesterday to today").unwrap();
        assert!(range.end.unixtime - range.start.unixtime >= (24 + 20) * 3600);
    }

    /// String representation of parse result.
    fn t(date: &str) -> String {
        match HgTime::parse(date) {
            Some(time) => format!("{} {}", time.unixtime, time.offset),
            None => "fail".to_string(),
        }
    }

    /// String representation of (parse result - now) / seconds.
    fn d(date: &str, duration: Duration) -> String {
        match HgTime::parse(date) {
            Some(time) => {
                let value = (time.unixtime as i64 - HgTime::now().unwrap().unixtime as i64).abs()
                    / duration.num_seconds();
                format!("{}", value)
            }
            None => "fail".to_string(),
        }
    }

    /// String "contains" (if range contains date) or "does not contain"
    /// or "fail" (if either range or date fails to parse).
    fn c(range: &str, date: &str) -> &'static str {
        if let (Some(range), Some(date)) = (HgTime::parse_range(range), HgTime::parse(date)) {
            if range.contains(&date) {
                "contains"
            } else {
                "does not contain"
            }
        } else {
            "fail"
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashSet;
use std::path::PathBuf;
use std::pin::Pin;

use curl::easy::Easy2;
use futures::prelude::*;
use url::Url;

use crate::driver::MultiDriver;
use crate::errors::Abort;
use crate::errors::HttpClientError;
use crate::event_listeners::HttpClientEventListeners;
use crate::handler::Buffered;
use crate::handler::HandlerExt;
use crate::handler::Streaming;
use crate::pool::Pool;
use crate::progress::Progress;
use crate::receiver::ChannelReceiver;
use crate::receiver::Receiver;
use crate::request::Method;
use crate::request::Request;
use crate::request::StreamRequest;
use crate::response::AsyncResponse;
use crate::response::Response;
use crate::stats::Stats;

pub type ResponseFuture =
    Pin<Box<dyn Future<Output = Result<AsyncResponse, HttpClientError>> + Send + 'static>>;
pub type StatsFuture =
    Pin<Box<dyn Future<Output = Result<Stats, HttpClientError>> + Send + 'static>>;

/// An async-compatible HTTP client powered by libcurl.
///
/// Essentially a more ergonomic API for working with
/// libcurl's multi interface. See URL for details:
///
/// https://curl.haxx.se/libcurl/c/libcurl-multi.html
///
/// NOTE: If you do not need to perform multiple concurrent
/// requests, you may want to use  `Request::send` instead.
#[derive(Clone)]
pub struct HttpClient {
    pool: Pool,
    event_listeners: HttpClientEventListeners,
    config: Config,
}

#[derive(Clone, Debug)]
pub struct Config {
    pub cert_path: Option<PathBuf>,
    pub key_path: Option<PathBuf>,
    pub ca_path: Option<PathBuf>,
    pub convert_cert: bool,

    pub client_info: Option<String>,
    pub disable_tls_verification: bool,
    pub max_concurrent_requests: Option<usize>,
    pub unix_socket_domains: HashSet<String>,
    pub unix_socket_path: Option<String>,
    pub verbose: bool,
    pub verbose_stats: bool,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            cert_path: None,
            key_path: None,
            ca_path: None,
            convert_cert: cfg!(windows),

            client_info: None,
            disable_tls_verification: false,
            max_concurrent_requests: None, // No limit by default
            unix_socket_domains: HashSet::new(),
            unix_socket_path: None,
            verbose: false,
            verbose_stats: false,
        }
    }
}

impl HttpClient {
    pub fn new() -> Self {
        Self::from_config(Default::default())
    }

    pub fn from_config(config: Config) -> Self {
        Self {
            config,
            pool: Pool::new(),
            event_listeners: Default::default(),
        }
    }

    pub fn verbose_stats(mut self, verbose: bool) -> Self {
        self.config.verbose_stats = verbose;
        self
    }

    pub fn max_concurrent_requests(mut self, max: Option<usize>) -> Self {
        self.config.max_concurrent_requests = max;
        self
    }

    /// Perform multiple HTTP requests concurrently.
    ///
    /// This function will block until all transfers have completed.
    /// Whenever a transfer completes, the user-provided closure
    /// will be called with the result.
    ///
    /// The closure returns a boolean. If false, this function will
    /// return early and all other pending transfers will be aborted.
    pub fn send<I, F>(&self, requests: I, response_cb: F) -> Result<Stats, HttpClientError>
    where
        I: IntoIterator<Item = Request>,
        F: FnMut(Result<Response, HttpClientError>) -> Result<(), Abort>,
    {
        self.send_with_progress(requests, response_cb, |_| ())
    }

    /// Same as `send()`, but takes an additional closure for
    /// monitoring the collective progress of all of the transfers.
    /// The closure will be called whenever any of the underlying
    /// transfers make progress.
    pub fn send_with_progress<I, F, P>(
        &self,
        requests: I,
        mut response_cb: F,
        progress_cb: P,
    ) -> Result<Stats, HttpClientError>
    where
        I: IntoIterator<Item = Request>,
        F: FnMut(Result<Response, HttpClientError>) -> Result<(), Abort>,
        P: FnMut(Progress),
    {
        let mut multi = self.pool.multi();
        multi
            .get_mut()
            .set_max_total_connections(self.config.max_concurrent_requests.unwrap_or(0))?;
        let driver = MultiDriver::new(multi.get(), progress_cb, self.config.verbose_stats);

        for mut request in requests {
            self.event_listeners.trigger_new_request(request.ctx_mut());
            let handle: Easy2<Buffered> = request.try_into()?;
            driver.add(handle)?;
        }

        let mut tls_error = false;
        let stats = driver.perform(|res| {
            if let Err((_, e)) = &res {
                let e: HttpClientError = e.clone().into();
                if let HttpClientError::Tls(_) = e {
                    tls_error = true;
                }
            }
            let res = res
                .map_err(|(mut easy, e)| {
                    let ctx = easy.get_mut().request_context_mut();
                    let info = ctx.info().clone();
                    ctx.event_listeners().trigger_failure(&info);
                    self.event_listeners.trigger_failed_request(ctx);
                    e.into()
                })
                .and_then(|mut easy| {
                    let ctx = easy.get_mut().request_context_mut();
                    let info = ctx.info().clone();
                    ctx.event_listeners().trigger_success(&info);
                    self.event_listeners.trigger_succeeded_request(ctx);
                    Response::try_from(easy.get_mut())
                });
            response_cb(res)
        })?;

        self.event_listeners.trigger_stats(&stats);

        drop(driver);

        // Don't reuse the connection if we've hit auth issues. We've seen cases where we reuse
        // expired credentials.
        if tls_error {
            multi.discard();
        }

        Ok(stats)
    }

    /// Async version of `send` which runs all of the given request concurrently
    /// in another thread.
    ///
    /// Returns a `Vec` of `ResponseFuture`s corresponding to each of the given
    /// input `Request`s. (They are returned in the same order as they were
    /// passed in.) Each `ResponseFuture` will resolve once all of the headers
    /// for that request have been received. The resulting `AsyncResponse` can
    /// be used to access the headers and body stream.
    pub fn send_async<I: IntoIterator<Item = Request>>(
        &self,
        requests: I,
    ) -> Result<(Vec<ResponseFuture>, StatsFuture), HttpClientError> {
        self.send_async_with_progress(requests, |_| ())
    }

    /// Same as `send_async()`, but takes an additional closure for monitoring
    /// the collective progress of all of the transfers. The closure will be
    /// called whenever any of the underlying transfers make progress.
    pub fn send_async_with_progress<I, P>(
        &self,
        requests: I,
        progress_cb: P,
    ) -> Result<(Vec<ResponseFuture>, StatsFuture), HttpClientError>
    where
        I: IntoIterator<Item = Request>,
        P: FnMut(Progress) + Send + 'static,
    {
        let client = self.clone();

        let mut stream_requests = Vec::new();
        let mut responses = Vec::new();

        for req in requests {
            let request_info = req.ctx().info().clone();
            let (receiver, streams) = ChannelReceiver::new();

            // Create a blocking streaming HTTP request to be dispatched on a
            // separate IO task.
            stream_requests.push(req.into_streaming(receiver));

            // Create response Future to return to the caller. The response is
            // linked to the request via channels, allowing async Rust code to
            // seamlessly receive data from the IO task.
            responses.push(AsyncResponse::new(streams, request_info).boxed());
        }

        let task = tokio::task::spawn_blocking(move || {
            client.stream_with_progress(stream_requests, progress_cb)
        });

        let stats = task
            .err_into::<HttpClientError>()
            .map(|res| Ok(res??))
            .boxed();

        Ok((responses, stats))
    }

    /// Perform the given requests, but stream the responses to the
    /// `Receiver` attached to each respective request rather than
    /// buffering the content of each response.
    ///
    /// Note that this function is not asynchronous; it WILL BLOCK
    /// until all of the transfers are complete, and will return
    /// the total stats across all transfers when complete.
    pub fn stream<I, R>(&self, requests: I) -> Result<Stats, HttpClientError>
    where
        I: IntoIterator<Item = StreamRequest<R>>,
        R: Receiver,
    {
        self.stream_with_progress(requests, |_| ())
    }

    /// Same as `stream()`, but takes an additional closure for
    /// monitoring the collective progress of all of the transfers.
    /// The closure will be called whenever any of the underlying
    /// transfers make progress.
    pub fn stream_with_progress<I, R, P>(
        &self,
        requests: I,
        progress_cb: P,
    ) -> Result<Stats, HttpClientError>
    where
        I: IntoIterator<Item = StreamRequest<R>>,
        R: Receiver,
        P: FnMut(Progress),
    {
        let mut multi = self.pool.multi();
        multi
            .get_mut()
            .set_max_total_connections(self.config.max_concurrent_requests.unwrap_or(0))?;
        let driver = MultiDriver::new(multi.get(), progress_cb, self.config.verbose_stats);
        for mut request in requests {
            self.event_listeners
                .trigger_new_request(request.request.ctx_mut());
            let handle: Easy2<Streaming<R>> = request.try_into()?;
            driver.add(handle)?;
        }

        let mut tls_error = false;
        let result = driver
            .perform(|res| {
                if let Err((_, e)) = &res {
                    let e: HttpClientError = e.clone().into();
                    if let HttpClientError::Tls(_) = e {
                        tls_error = true;
                    }
                }
                self.report_result_and_drop_receiver(res)
            })
            .map(|stats| {
                self.event_listeners.trigger_stats(&stats);
                stats
            });

        drop(driver);

        // Don't reuse the connection if we've hit auth issues. We've seen cases where we reuse
        // expired credentials.
        if tls_error {
            multi.discard();
        }

        result
    }

    /// Obtain the `HttpClientEventListeners` to register callbacks.
    pub fn event_listeners(&mut self) -> &mut HttpClientEventListeners {
        &mut self.event_listeners
    }

    /// Easier way to register event callbacks using the builder pattern.
    pub fn with_event_listeners(mut self, f: impl FnOnce(&mut HttpClientEventListeners)) -> Self {
        f(&mut self.event_listeners);
        self
    }

    /// Callback for `MultiDriver::perform` when working with
    /// a `Streaming` handler. Reports the result of the
    /// completed request to the handler's `Receiver`.
    fn report_result_and_drop_receiver<R: Receiver>(
        &self,
        res: Result<Easy2<Streaming<R>>, (Easy2<Streaming<R>>, curl::Error)>,
    ) -> Result<(), Abort> {
        // We need to get the `Easy2` handle in both the
        // success and error cases since we ultimately
        // need to pass the result to the handler contained
        // therein.
        let (mut easy, res) = match res {
            Ok(mut easy) => {
                let ctx = easy.get_mut().request_context_mut();
                let info = ctx.info().clone();
                ctx.event_listeners().trigger_success(&info);
                self.event_listeners.trigger_succeeded_request(ctx);
                (easy, Ok(()))
            }
            Err((mut easy, e)) => {
                let ctx = easy.get_mut().request_context_mut();
                let info = ctx.info().clone();
                ctx.event_listeners().trigger_failure(&info);
                self.event_listeners.trigger_failed_request(ctx);
                (easy, Err(e.into()))
            }
        };

        // Extract the `Receiver` from the `Streaming` handler
        // inside the Easy2 handle. If it's already gone, just
        // log it and move on. (This shouldn't normally happen.)
        if let Some(receiver) = easy.get_mut().take_receiver() {
            receiver.done(res)
        } else {
            tracing::error!("Cannot report status because receiver is missing");
            Ok(())
        }
    }

    /// Create a request with this client's config applied.
    pub fn new_request(&self, url: Url, method: Method) -> Request {
        self.configure_request(Request::new(url, method))
    }

    /// Create a GET request with this client's config applied.
    pub fn get(&self, url: Url) -> Request {
        self.new_request(url, Method::Get)
    }

    /// Create a HEAD request with this client's config applied.
    pub fn head(&self, url: Url) -> Request {
        self.new_request(url, Method::Head)
    }

    /// Create a POST request with this client's config applied.
    pub fn post(&self, url: Url) -> Request {
        self.new_request(url, Method::Post)
    }

    /// Create a PUT request with this client's config applied.
    pub fn put(&self, url: Url) -> Request {
        self.new_request(url, Method::Put)
    }

    fn configure_request(&self, mut req: Request) -> Request {
        req.set_client_info(&self.config.client_info);
        req.set_convert_cert(self.config.convert_cert);
        req.set_verbose(self.config.verbose);

        if let Some(domain) = req.ctx().url().domain() {
            if self.config.unix_socket_domains.contains(domain) {
                req.set_auth_proxy_socket_path(self.config.unix_socket_path.clone());
            }
        }

        if let Some(cert_path) = &self.config.cert_path {
            req.set_cert(cert_path);
        }

        if let Some(key_path) = &self.config.key_path {
            req.set_key(key_path);
        }

        if let Some(ca_path) = &self.config.ca_path {
            req.set_cainfo(ca_path);
        }

        req.set_verify_tls_cert(!self.config.disable_tls_verification);
        req.set_verify_tls_host(!self.config.disable_tls_verification);

        req
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use anyhow::Result;
    use http::StatusCode;
    use mockito::mock;
    use url::Url;

    use super::*;
    use crate::receiver::testutil::TestReceiver;
    use crate::Method;
    use crate::RequestContext;

    #[test]
    fn test_client() -> Result<()> {
        let body1 = b"body1";
        let body2 = b"body2";
        let body3 = b"body3";

        let mock1 = mock("GET", "/test1")
            .with_status(201)
            .with_body(&body1)
            .create();

        let mock2 = mock("GET", "/test2")
            .with_status(201)
            .with_body(&body2)
            .create();

        let mock3 = mock("GET", "/test3")
            .with_status(201)
            .with_body(&body3)
            .create();

        let server_url = Url::parse(&mockito::server_url())?;

        let url1 = server_url.join("test1")?;
        let req1 = Request::get(url1);

        let url2 = server_url.join("test2")?;
        let req2 = Request::get(url2);

        let url3 = server_url.join("test3")?;
        let req3 = Request::get(url3);

        let mut not_received = HashSet::new();
        not_received.insert(body1.to_vec());
        not_received.insert(body2.to_vec());
        not_received.insert(body3.to_vec());

        let client = HttpClient::new();
        let stats = client.send(vec![req1, req2, req3], |res| {
            let res = res.unwrap();
            assert_eq!(res.head.status, StatusCode::CREATED);
            assert!(not_received.remove(&*res.body));
            Ok(())
        })?;

        mock1.assert();
        mock2.assert();
        mock3.assert();

        assert!(not_received.is_empty());
        assert_eq!(stats.requests, 3);

        Ok(())
    }

    #[test]
    fn test_stream() -> Result<()> {
        let body1 = b"body1";
        let body2 = b"body2";
        let body3 = b"body3";

        let mock1 = mock("GET", "/test1")
            .with_status(201)
            .with_body(&body1)
            .create();

        let mock2 = mock("GET", "/test2")
            .with_status(201)
            .with_body(&body2)
            .create();

        let mock3 = mock("GET", "/test3")
            .with_status(201)
            .with_body(&body3)
            .create();

        let server_url = Url::parse(&mockito::server_url())?;

        let url1 = server_url.join("test1")?;
        let rcv1 = TestReceiver::new();
        let req1 = Request::get(url1).into_streaming(rcv1.clone());

        let url2 = server_url.join("test2")?;
        let rcv2 = TestReceiver::new();
        let req2 = Request::get(url2).into_streaming(rcv2.clone());

        let url3 = server_url.join("test3")?;
        let rcv3 = TestReceiver::new();
        let req3 = Request::get(url3).into_streaming(rcv3.clone());

        let client = HttpClient::new();
        let stats = client.stream(vec![req1, req2, req3])?;

        mock1.assert();
        mock2.assert();
        mock3.assert();
        assert_eq!(stats.requests, 3);

        assert_eq!(rcv1.status().unwrap(), StatusCode::CREATED);
        let body = rcv1.chunks().into_iter().flatten().collect::<Vec<_>>();
        assert_eq!(&*body, body1);

        assert_eq!(rcv2.status().unwrap(), StatusCode::CREATED);
        let body = rcv2.chunks().into_iter().flatten().collect::<Vec<_>>();
        assert_eq!(&*body, body2);

        assert_eq!(rcv3.status().unwrap(), StatusCode::CREATED);
        let body = rcv3.chunks().into_iter().flatten().collect::<Vec<_>>();
        assert_eq!(&*body, body3);

        Ok(())
    }

    #[tokio::test]
    async fn test_async() -> Result<()> {
        let body1 = b"body1";
        let body2 = b"body2";
        let body3 = b"body3";

        let mock1 = mock("GET", "/test1")
            .with_status(201)
            .with_body(&body1)
            .create();

        let mock2 = mock("GET", "/test2")
            .with_status(201)
            .with_body(&body2)
            .create();

        let mock3 = mock("GET", "/test3")
            .with_status(201)
            .with_body(&body3)
            .create();

        let server_url = Url::parse(&mockito::server_url())?;

        let url1 = server_url.join("test1")?;
        let req1 = Request::get(url1);

        let url2 = server_url.join("test2")?;
        let req2 = Request::get(url2);

        let url3 = server_url.join("test3")?;
        let req3 = Request::get(url3);

        let client = HttpClient::new();
        let (futures, stats) = client.send_async(vec![req1, req2, req3])?;

        let mut responses = Vec::new();
        for fut in futures {
            responses.push(fut.await?);
        }

        mock1.assert();
        mock2.assert();
        mock3.assert();

        let mut not_received = HashSet::new();
        not_received.insert(body1.to_vec());
        not_received.insert(body2.to_vec());
        not_received.insert(body3.to_vec());

        for res in responses {
            assert_eq!(res.head.status, StatusCode::CREATED);
            let body = res.into_body().raw().try_concat().await?;
            assert!(not_received.remove(&*body));
        }

        assert!(not_received.is_empty());

        let stats = stats.await?;
        assert_eq!(stats.requests, 3);

        Ok(())
    }

    #[tokio::test]
    async fn test_event_listeners() -> Result<()> {
        let server_url = Url::parse(&mockito::server_url())?;

        // this is actually used, it changes how mockito behaves
        const BODY: &[u8] = b"body";
        let _mock1 = mock("GET", "/test1")
            .with_status(201)
            .with_body(BODY)
            .create();

        let url = server_url.join("test1")?;
        let request = Request::get(url);

        let (tx, rx) = crossbeam::channel::unbounded();
        let (msg_tx, msg_rx) = crossbeam::channel::unbounded();

        let client = HttpClient::new().with_event_listeners(|l| {
            l.on_stats(move |stats| tx.send(stats.clone()).expect("send stats over channel"));

            let check_request = &|r: &RequestContext| assert!(matches!(r.method(), Method::Get));

            l.on_new_request({
                let msg_tx = msg_tx.clone();
                move |r| {
                    check_request(r);
                    msg_tx.send("on_new_request").unwrap();

                    let l = r.event_listeners();
                    l.on_download_bytes({
                        let msg_tx = msg_tx.clone();
                        move |req, n| {
                            assert_eq!(n, BODY.len());
                            check_request(req);
                            msg_tx.send("on_download_bytes").unwrap();
                        }
                    });
                    l.on_upload_bytes({
                        let msg_tx = msg_tx.clone();
                        move |req, n| {
                            assert_eq!(n, 0);
                            check_request(req);
                            msg_tx.send("on_upload_bytes").unwrap();
                        }
                    });
                    l.on_content_length({
                        let msg_tx = msg_tx.clone();
                        move |req, n| {
                            assert_eq!(n, BODY.len());
                            check_request(req);
                            msg_tx.send("on_content_length").unwrap();
                        }
                    });
                }
            });

            l.on_failed_request({
                let msg_tx = msg_tx.clone();
                move |r| {
                    check_request(r);
                    msg_tx.send("on_failed_request").unwrap();
                }
            });

            l.on_succeeded_request({
                let msg_tx = msg_tx.clone();
                move |r| {
                    check_request(r);
                    msg_tx.send("on_succeeded_request").unwrap();
                }
            });
        });

        let check_events = |has_content_length: bool| {
            assert_eq!(msg_rx.recv().unwrap(), "on_new_request");
            if has_content_length {
                assert_eq!(msg_rx.recv().unwrap(), "on_content_length");
            }
            assert_eq!(msg_rx.recv().unwrap(), "on_download_bytes");
            assert_eq!(msg_rx.recv().unwrap(), "on_succeeded_request");
        };

        let stats = client.send(vec![request.clone()], |_| Ok(()))?;
        assert_eq!(stats, rx.recv()?);
        check_events(true);

        let stats = client.send_with_progress(vec![request.clone()], |_| Ok(()), |_| ())?;
        assert_eq!(stats, rx.recv()?);
        check_events(true);

        let (_stream, stats) = client.send_async(vec![request.clone()])?;
        let stats = stats.await?;
        assert_eq!(stats, rx.recv()?);
        check_events(false);

        let (_stream, stats) = client.send_async_with_progress(vec![request.clone()], |_| ())?;
        let stats = stats.await?;
        assert_eq!(stats, rx.recv()?);
        check_events(false);

        let my_stream_req = || request.clone().into_streaming(TestReceiver::new());

        let stats = client.stream(vec![my_stream_req()])?;
        assert_eq!(stats, rx.recv()?);
        check_events(false);

        let stats = client.stream_with_progress(vec![my_stream_req()], |_| ())?;
        assert_eq!(stats, rx.recv()?);
        check_events(false);

        drop((client, msg_tx));

        // All msg_tx should be dropped. recv() should not be blocking.
        msg_rx.recv().unwrap_err();

        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cell::RefCell;
use std::time::Duration;
use std::time::Instant;

use anyhow::Context;
use curl::easy::Easy2;
use curl::multi::Easy2Handle;
use curl::multi::Message;
use curl::multi::Multi;

use crate::errors::Abort;
use crate::errors::HttpClientError;
use crate::handler::HandlerExt;
use crate::progress::Progress;
use crate::progress::ProgressReporter;
use crate::stats::Stats;

/// Maximum time that libcurl should wait for socket activity during a call to
/// `Multi::wait`. The Multi session maintains its own timeout internally based
/// on the state of the underlying transfers; this default value will only be
/// used if there is no internal timer value set at the time `wait` is called.
const MULTI_WAIT_TIMEOUT: Duration = Duration::from_secs(10);

/// A complete transfer, along with the associated error
/// if the transfer did not complete successfully.
struct Complete<H> {
    token: usize,
    handle: Easy2<H>,
    result: Result<(), curl::Error>,
}

impl<H> Complete<H> {
    /// If we encountered an error, we should still return the
    /// handle, as the callback may want to access the Handler
    /// inside.
    fn into_result(self) -> Result<Easy2<H>, (Easy2<H>, curl::Error)> {
        let Self { handle, result, .. } = self;
        match result {
            Ok(()) => Ok(handle),
            Err(e) => Err((handle, e)),
        }
    }
}

/// Struct that manages a curl::Multi session, synchronously driving
/// all of the transfers therein to completion.
pub(crate) struct MultiDriver<'a, H, P>
where
    H: HandlerExt,
    P: FnMut(Progress),
{
    multi: &'a Multi,
    handles: RefCell<Vec<Option<Easy2Handle<H>>>>,
    progress: ProgressReporter<P>,
    verbose: bool,
}

impl<'a, H, P> MultiDriver<'a, H, P>
where
    H: HandlerExt,
    P: FnMut(Progress),
{
    pub(crate) fn new(multi: &'a Multi, progress_cb: P, verbose: bool) -> Self {
        Self {
            multi,
            handles: RefCell::new(Vec::new()),
            progress: ProgressReporter::with_callback(progress_cb),
            verbose,
        }
    }

    pub(crate) fn num_transfers(&self) -> usize {
        (&*self.handles.borrow()).len()
    }

    /// Add an Easy2 handle to the Multi stack.
    pub(crate) fn add(&self, mut easy: Easy2<H>) -> Result<(), HttpClientError> {
        // Register this Easy2 handle's Handler with our ProgressReporter
        // so we can aggregate progress across all transfers in the stack.
        easy.get_mut()
            .request_context_mut()
            .event_listeners()
            .on_progress({
                let updater = self.progress.updater();
                move |_req, progress| updater.update(progress)
            });

        // Assign a token to this Easy2 handle so we can correlate messages
        // for this handle with the corresponding Easy2Handle while the
        // Easy2 is owned by the Multi handle.
        let mut handles = self.handles.borrow_mut();
        let token = handles.len();
        let mut handle = self.multi.add2(easy)?;
        handle.set_token(token)?;

        handles.push(Some(handle));

        Ok(())
    }

    /// Drive all of the Easy2 handles in the Multi stack to completion.
    ///
    /// The user-supplied callback will be called whenever a transfer
    /// completes, successfully or otherwise. The callback may cause this
    /// method to return early (aborting all other active transfers).
    pub(crate) fn perform<F>(&self, mut callback: F) -> Result<Stats, HttpClientError>
    where
        F: FnMut(Result<Easy2<H>, (Easy2<H>, curl::Error)>) -> Result<(), Abort>,
    {
        let total = self.num_transfers();
        let mut in_progress = total;

        tracing::debug!("Performing {} transfer(s)", total);

        let start = Instant::now();

        loop {
            let active_transfers = self.multi.perform()? as usize;

            // Check how many transfers have completed so far.
            if active_transfers != in_progress {
                tracing::debug!(
                    "{}/{} transfer(s) complete",
                    total - active_transfers,
                    total
                );
                in_progress = active_transfers;
            }

            // Check for messages. A message indicates a transfer completed or failed.
            let mut completed = Vec::new();
            self.multi.messages(|msg| match self.handle_msg(&msg) {
                Ok(complete) => {
                    tracing::trace!("Transfer {} complete", complete.token);
                    completed.push(complete);
                }
                Err(e) => {
                    tracing::error!("Failed to handle message: {}", e);
                }
            });

            // Run the user-provided callback on each completed transfer. If it returns an
            // error (signalling that we should return early) abort all remaining transfers.
            for c in completed {
                let token = c.token;
                callback(c.into_result())?;
                tracing::trace!("Successfully handled transfer: {}", token);
            }

            // If any tranfers reported progress, notify the user.
            self.progress.report_if_updated();

            if in_progress == 0 {
                tracing::trace!("All transfers finished successfully");
                break;
            }

            tracing::trace!("Waiting for socket activity");
            let active_sockets = self.multi.wait(&mut [], MULTI_WAIT_TIMEOUT)?;
            if active_sockets == 0 {
                tracing::trace!("Timed out waiting for activity");
            }
        }

        let elapsed = start.elapsed();

        let progress = self.progress.aggregate();
        let latency = self
            .progress
            .first_byte_received()
            .unwrap_or(start)
            .duration_since(start);

        let stats = Stats {
            downloaded: progress.downloaded,
            uploaded: progress.uploaded,
            requests: self.num_transfers(),
            time: elapsed,
            latency,
        };

        tracing::debug!("{}", &stats);
        if self.verbose {
            eprintln!("{}", &stats);
        }

        Ok(stats)
    }

    /// Handle a message emitted by the Multi session for any of the
    /// underlying transfers. Based on the current implementation details
    /// of libcurl, a message should only be emitted when a transfer has
    /// completed (successfully or otherwise).
    fn handle_msg(&self, msg: &Message<'_>) -> Result<Complete<H>, HttpClientError> {
        let (token, result) = {
            let token = msg.token()?;
            let handles = self.handles.borrow();
            let handle = handles[token].as_ref().context("Handle already removed")?;
            let result = msg
                .result_for2(handle)
                .context("Failed to get result for handle")?;
            (token, result)
        };

        // If we've gotten this far, we can conclude the transfer has completed
        // (successfully or otherwise), so it can be removed from the stack.
        let handle = self.remove(token)?.context("Handle already removed")?;

        Ok(Complete {
            token,
            handle,
            result,
        })
    }

    /// Remove and return an Easy2 handle from the Multi stack.
    fn remove(&self, index: usize) -> Result<Option<Easy2<H>>, HttpClientError> {
        if let Some(handle) = self.handles.borrow_mut()[index].take() {
            let easy = self.multi.remove2(handle)?;
            Ok(Some(easy))
        } else {
            Ok(None)
        }
    }

    /// Drop all of the outstanding Easy2 handles in the Multi stack.
    fn drop_all(&mut self) {
        let mut dropped = 0;
        for handle in self.handles.borrow_mut().drain(..) {
            if let Some(handle) = handle {
                let _ = self.multi.remove2(handle);
                dropped += 1;
            }
        }

        if dropped > 0 {
            tracing::debug!("Dropped {} outstanding transfers", dropped);
        }
    }
}

impl<'a, H, P> Drop for MultiDriver<'a, H, P>
where
    H: HandlerExt,
    P: FnMut(Progress),
{
    fn drop(&mut self) {
        self.drop_all();
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use futures::channel::oneshot;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum HttpClientError {
    #[error(transparent)]
    Curl(curl::Error),
    #[error(transparent)]
    Tls(#[from] TlsError),
    #[error(transparent)]
    CurlMulti(#[from] curl::MultiError),
    #[error(transparent)]
    CallbackAborted(#[from] Abort),
    #[error("Received invalid or malformed HTTP response")]
    BadResponse(anyhow::Error),
    #[error("The request was dropped before it could complete")]
    RequestDropped(#[from] oneshot::Canceled),
    #[error("The I/O task terminated unexpectedly: {}", .0)]
    IoTaskFailed(#[from] tokio::task::JoinError),
    #[error(transparent)]
    CborError(#[from] serde_cbor::Error),
    #[error(transparent)]
    CborStreamError(#[from] crate::stream::CborStreamError),
    #[error("Could not decode response: {}", .0)]
    DecompressionFailed(futures::io::Error),
    #[error(transparent)]
    Other(#[from] anyhow::Error),
}

impl From<curl::Error> for HttpClientError {
    fn from(e: curl::Error) -> Self {
        TlsError::try_from(e).map_or_else(HttpClientError::Curl, HttpClientError::Tls)
    }
}

/// Error type for user-provided callbacks. Indicates that the client should
/// abort the operation and return early. The user may optionally provide a
/// reason for aborting.
#[derive(Error, Debug)]
pub enum Abort {
    #[error("Operation aborted by user callback: {0}")]
    WithReason(#[source] anyhow::Error),
    #[error("Operation aborted by user callback")]
    Unspecified,
}

impl Abort {
    pub fn abort<E: Into<anyhow::Error>>(reason: E) -> Self {
        Abort::WithReason(reason.into())
    }
}

/// A strongly-typed representation of all possible TLS-related error codes
/// from libcurl. It is useful to separate these from other kinds of libcurl
/// errors because the problem is often related to the client's configuration
/// (e.g., expired client certificate or wrong CA bundle).
///
/// Full list of error codes: https://curl.se/libcurl/c/libcurl-errors.html
#[derive(Error, Debug)]
#[error("Try renewing your certificates. Run `eden doctor`. TlsError: {source}")]
pub struct TlsError {
    pub source: curl::Error,
    pub kind: TlsErrorKind,
}

impl TryFrom<curl::Error> for TlsError {
    type Error = curl::Error;

    fn try_from(source: curl::Error) -> Result<Self, Self::Error> {
        match TlsErrorKind::from_curl_error(&source) {
            Some(kind) => Ok(Self { source, kind }),
            None => Err(source),
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
pub enum TlsErrorKind {
    RecvError = curl_sys::CURLE_RECV_ERROR as isize,
    CaCert = curl_sys::CURLE_SSL_CACERT as isize,
    CaCertBadFile = curl_sys::CURLE_SSL_CACERT_BADFILE as isize,
    CertProblem = curl_sys::CURLE_SSL_CERTPROBLEM as isize,
    Cipher = curl_sys::CURLE_SSL_CIPHER as isize,
    ConnectError = curl_sys::CURLE_SSL_CONNECT_ERROR as isize,
    CrlBadFile = curl_sys::CURLE_SSL_CRL_BADFILE as isize,
    EngineInitFailed = curl_sys::CURLE_SSL_ENGINE_INITFAILED as isize,
    EngineNotFound = curl_sys::CURLE_SSL_ENGINE_NOTFOUND as isize,
    EngineSetFailed = curl_sys::CURLE_SSL_ENGINE_SETFAILED as isize,
    InvalidCertStatus = curl_sys::CURLE_SSL_INVALIDCERTSTATUS as isize,
    IssuerError = curl_sys::CURLE_SSL_ISSUER_ERROR as isize,
    PinnedPubKeyNotMatch = curl_sys::CURLE_SSL_PINNEDPUBKEYNOTMATCH as isize,
    ShutdownFailed = curl_sys::CURLE_SSL_SHUTDOWN_FAILED as isize,
}

impl TlsErrorKind {
    fn from_curl_error(source: &curl::Error) -> Option<Self> {
        use TlsErrorKind::*;

        Some(match source.code() {
            curl_sys::CURLE_RECV_ERROR => ssl_categorize_recv_error(source),
            curl_sys::CURLE_SSL_CACERT => CaCert,
            curl_sys::CURLE_SSL_CACERT_BADFILE => CaCertBadFile,
            curl_sys::CURLE_SSL_CERTPROBLEM => CertProblem,
            curl_sys::CURLE_SSL_CIPHER => Cipher,
            curl_sys::CURLE_SSL_CONNECT_ERROR => ConnectError,
            curl_sys::CURLE_SSL_CRL_BADFILE => CrlBadFile,
            curl_sys::CURLE_SSL_ENGINE_INITFAILED => EngineInitFailed,
            curl_sys::CURLE_SSL_ENGINE_NOTFOUND => EngineNotFound,
            curl_sys::CURLE_SSL_ENGINE_SETFAILED => EngineSetFailed,
            curl_sys::CURLE_SSL_INVALIDCERTSTATUS => InvalidCertStatus,
            curl_sys::CURLE_SSL_ISSUER_ERROR => IssuerError,
            curl_sys::CURLE_SSL_PINNEDPUBKEYNOTMATCH => PinnedPubKeyNotMatch,
            curl_sys::CURLE_SSL_SHUTDOWN_FAILED => ShutdownFailed,
            _ => return None,
        })
    }
}

/// XXX: When libcurl's underlying TLS engine (e.g., OpenSSL) encounters an
/// error during the transfer, libcurl might report a generic error (such as
/// CURLE_RECV_ERROR) instead of a TLS-specific error. In these cases, the
/// "extra description" field of the error usually contains the actual error
/// message from the TLS engine. This function tries to detect this situation
/// using a crude string search heuristic.
fn ssl_categorize_recv_error(error: &curl::Error) -> TlsErrorKind {
    let extra = error.extra_description();
    if let Some(extra) = extra {
        let extra = extra.to_lowercase();
        // Covers: "alert certificate required" and "alert bad certificate"
        if extra.contains("certificate") {
            return TlsErrorKind::CertProblem;
        }
    }

    TlsErrorKind::RecvError
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use crate::progress::Progress;
use crate::request::Request;
use crate::request::RequestContext;
use crate::request::RequestInfo;
use crate::stats::Stats;

/// Generate a struct for holding event listeners (callbacks).
///
/// For each event `event_name(*args)`, two methods are generated:
/// - `on_event_name(f: impl Fn(*args))`: register a listener (callback).
/// - `trigger_event_name(*args)`: trigger the event, call registered callbacks.
macro_rules! gen_event_listeners {
    (
        #[doc=$struct_doc:literal]
        $struct_name:ident {
            $(
                #[doc=$event_doc:literal]
                $event_name:ident ( $($arg_name:ident : $arg_ty:ty),* ),
            )*
        }
    ) => {
        paste::paste! {
            #[doc=$struct_doc]
            #[derive(Default, Clone)]
            pub struct $struct_name {
                $(
                    #[doc=$event_doc]
                    [< $event_name _listeners >] : Vec::<Arc<dyn Fn( $($arg_ty,)* ) + Send + Sync>>,
                )*
            }

            impl ::std::fmt::Debug for $struct_name {
                fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
                    f.write_str(stringify!($struct_name))
                }
            }

            impl $struct_name {
                $(
                    /// Register a callback for the given event type.
                    #[doc=$event_doc]
                    pub fn [< on_ $event_name >](&mut self, f: impl Fn( $($arg_ty,)* ) + Send + Sync + 'static) {
                        self. [< $event_name _listeners >].push(Arc::new(f));
                    }

                    /// Returns true if there are callbacks registered on the event.
                    pub(crate) fn [< should_trigger_ $event_name >](&self) -> bool {
                        !self. [< $event_name _listeners >] .is_empty()
                    }

                    /// Call all registered callbacks for the given event type.
                    pub(crate) fn [< trigger_ $event_name >](&self, $($arg_name: $arg_ty,)*) {
                        for cb in &self. [< $event_name _listeners >] {
                            cb( $($arg_name,)* );
                        }
                    }
                )*
            }
        }
    }
}

gen_event_listeners! {
    /// Events for a `HttpClient`.
    HttpClientEventListeners {
        /// A request is created.
        new_request(req: &mut RequestContext),

        /// A request is completed successfully.
        succeeded_request(req: &RequestContext),

        /// A request is failed.
        failed_request(req: &RequestContext),

        /// One or more requests have completed with statistics.
        stats(stats: &Stats),
    }
}

gen_event_listeners! {
    /// Events for a single Request.
    RequestEventListeners {
        /// Downloaded `n` bytes. Does not include HTTP headers.
        download_bytes(req: &RequestContext, n: usize),

        /// Uploaded `n` bytes.
        upload_bytes(req: &RequestContext, n: usize),

        /// Received Content-Length as `n`.
        content_length(req: &RequestContext, n: usize),

        /// On first byte of network activity.
        first_activity(req: &RequestContext),

        /// On progress update. Note: this is called periodically even if there is no progress.
        progress(req: &RequestContext, progress: Progress),

        /// The request completed successfully.
        success(req: &RequestInfo),

        /// The request completed unsuccessfully.
        failure(req: &RequestInfo),
    }
}

gen_event_listeners! {
    /// Events for request creation (both independent requests and requests via `HttpClient`)
    RequestCreationEventListeners {
        /// A request is created.
        new_request(req: &mut Request),
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering::SeqCst;
    use std::sync::Arc;

    gen_event_listeners! {
        /// Testing event listeners.
        HelloEventListeners {
            /// A "hello" event.
            hello(msg: &str),
        }
    }

    #[test]
    fn test_hello_event_listener() {
        let mut listeners = HelloEventListeners::default();
        let len = Arc::new(AtomicUsize::new(0));
        listeners.on_hello({
            let len = len.clone();
            move |s| {
                len.fetch_add(s.len(), SeqCst);
            }
        });
        listeners.on_hello({
            let len = len.clone();
            move |s| {
                len.fetch_add(s.len() * 100, SeqCst);
            }
        });
        listeners.trigger_hello("abcde");
        assert_eq!(len.load(SeqCst), 505);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::str;

use http::header::HeaderName;
use http::header::HeaderValue;
use http::status::StatusCode;
use http::version::Version;
use once_cell::sync::Lazy;
use regex::Regex;
use thiserror::Error;

static STATUS_REGEX: Lazy<Regex> =
    Lazy::new(|| Regex::new(r"(?i)HTTP/([0-9.]+) ([0-9]+)").unwrap());

/// A parsed header line.
///
/// This enum represents a line from the header section of a
/// response. Note that in addition to the headers, libcurl
/// also passes the initial status line and trailing blank
/// line to the user-specified header callback, so we must
/// be able to handle those cases.
#[derive(Eq, PartialEq, Debug)]
pub enum Header {
    Status(Version, StatusCode),
    Header(HeaderName, HeaderValue),
    EndOfHeaders,
}

#[derive(Debug, Error)]
#[error("Malformed header: {:?}", String::from_utf8_lossy(.0))]
pub struct BadHeader<'a>(&'a [u8]);

impl Header {
    /// Parse a header line. The input is expected to be a CRLF-terminated
    /// line which can be decoded as UTF-8. Note that per RFC 7230, header
    /// values can sometimes contain arbitrary binary data, but in practice
    /// they are limited to ASCII characters, so for simplicity we reject
    /// non-UTF-8 header values. Aside from the values, the specification
    /// restricts all other parts of a header line to be limited to ASCII.
    pub fn parse(line: &[u8]) -> Result<Self, BadHeader> {
        let header = str::from_utf8(line)
            .map_err(|_| BadHeader(line))?
            .trim_end(); // Strip off trailing CRLF.

        if header.is_empty() {
            return Ok(Self::EndOfHeaders);
        }
        if let Some((name, value)) = parse_header(header) {
            return Ok(Self::Header(name, value));
        }
        if let Some((version, code)) = parse_status(header) {
            return Ok(Self::Status(version, code));
        }

        Err(BadHeader(line))
    }
}

/// Parse a status line, e.g. "HTTP/1.1 200 OK".
fn parse_status(line: &str) -> Option<(Version, StatusCode)> {
    let captures = STATUS_REGEX.captures(&line)?;

    let version_str = captures.get(1).map(|m| m.as_str())?;
    let version = parse_version(version_str)?;

    let code_str = captures.get(2).map(|m| m.as_str())?;
    let code = StatusCode::from_u16(code_str.parse().ok()?).ok()?;

    Some((version, code))
}

/// Parse an HTTP version number.
fn parse_version(version: &str) -> Option<Version> {
    Some(match version {
        "0.9" => Version::HTTP_09,
        "1.0" => Version::HTTP_10,
        "1.1" => Version::HTTP_11,
        "2" | "2.0" => Version::HTTP_2,
        "3" | "3.0 " => Version::HTTP_3,
        _ => return None,
    })
}

/// Parse a header name-value pair, e.g. "Content-Length: 42\r\n".
fn parse_header(header: &str) -> Option<(HeaderName, HeaderValue)> {
    let parts = header.splitn(2, ':').collect::<Vec<_>>();
    let (name, value) = if parts.len() > 1 {
        (parts[0], parts[1].trim_start())
    } else {
        (parts[0], "")
    };

    let name = HeaderName::from_bytes(name.as_bytes()).ok()?;
    let value = HeaderValue::from_bytes(value.as_bytes()).ok()?;

    Some((name, value))
}

#[cfg(test)]
mod test {
    use anyhow::Result;
    use http::header;

    use super::*;

    #[test]
    fn test_parse_header() -> Result<()> {
        let header = Header::parse(b"Content-Length: 42\r\n")?;
        let expected = Header::Header(header::CONTENT_LENGTH, HeaderValue::from_static("42"));
        assert_eq!(header, expected);

        let header = Header::parse(b"X-Non-Standard: test\r\n")?;
        let expected = Header::Header(
            HeaderName::from_static("x-non-standard"),
            HeaderValue::from_static("test"),
        );
        assert_eq!(header, expected);

        let header = Header::parse(b"X-No-Value\r\n")?;
        let expected = Header::Header(
            HeaderName::from_static("x-no-value"),
            HeaderValue::from_static(""),
        );
        assert_eq!(header, expected);

        let header = Header::parse(b"X-Whitespace:  hello  world  \r\n")?;
        let expected = Header::Header(
            HeaderName::from_static("x-whitespace"),
            HeaderValue::from_static("hello  world"),
        );
        assert_eq!(header, expected);

        let header = Header::parse("X-Non-ASCII-Value: \u{1F980}\r\n".as_ref())?;
        let expected = Header::Header(
            HeaderName::from_static("x-non-ascii-value"),
            HeaderValue::from_bytes("\u{1F980}".as_ref())?,
        );
        assert_eq!(header, expected);

        assert!(Header::parse("\u{1F980}: Non-ASCII name\r\n".as_ref()).is_err());
        Ok(())
    }

    #[test]
    fn test_parse_status() -> Result<()> {
        let status = Header::parse(b"HTTP/2 201 CREATED\r\n")?;
        let expected = Header::Status(Version::HTTP_2, StatusCode::CREATED);
        assert_eq!(status, expected);
        Ok(())
    }

    #[test]
    fn test_parse_crlf() -> Result<()> {
        assert_eq!(Header::parse(b"\r\n")?, Header::EndOfHeaders);
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! An async-compatible HTTP client built on top of libcurl.

#![allow(dead_code)]

mod client;
mod driver;
mod errors;
mod event_listeners;
mod handler;
mod header;
mod pool;
mod progress;
mod receiver;
mod request;
mod response;
mod stats;
mod stream;

pub use client::Config;
pub use client::HttpClient;
pub use client::ResponseFuture;
pub use client::StatsFuture;
pub use curl::easy::HttpVersion;
pub use errors::Abort;
pub use errors::HttpClientError;
pub use errors::TlsError;
pub use errors::TlsErrorKind;
pub use header::Header;
pub use progress::Progress;
pub use receiver::Receiver;
pub use request::Encoding;
pub use request::Method;
pub use request::MinTransferSpeed;
pub use request::Request;
pub use request::RequestContext;
pub use request::RequestInfo;
pub use request::StreamRequest;
pub use response::AsyncBody;
pub use response::AsyncResponse;
pub use response::Response;
pub use stats::Stats;
pub use stream::BufferedStream;
pub use stream::CborStream;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cmp::Ordering;
use std::collections::BinaryHeap;
use std::mem;
use std::sync::Arc;

use curl::multi::Multi;
use parking_lot::Mutex;

// Maximum number of handles in the pool. Number chosen arbitrarily.
// Note that this number does not limit the number of handles the pool
// will supply to the user, but instead limits the number of handles
// that are retained once they are returned to the pool.
const MAX_POOL_SIZE: usize = 1024;

/// A pool of libcurl `Multi` handles. Each multi session maintains caches for
/// things like TCP connections, TLS sessions, DNS information, etc. As such,
/// it makes sense to reuse multi sessions when possible to benefit from this
/// caching. Since a multi session can only be used by a single thread at a
/// time, this pool provides a mechanism by which threads can use and return
/// `Multi` handles, allowing them to be used repeatedly.
///
/// The `Pool` maintains a priority queue of `Multi` handles, each ranked by
/// how many times it has been used. The assumption is that a handle that has
/// been used more will have a warmer cache and should therefore be preferred.
/// When a borrowed handle is dropped by the thread using it, it will be
/// automatically returned to the pool with its priority incremented.
///
/// Those familiar with libcurl may be aware of libcurl's "share" interface,
/// which allows multiple curl handles (potentially on different threads) to
/// share caches. Unfortunately, the `curl` crate does not provide safe
/// bindings to the share interface, and implementing them would require
/// a lot of subtle unsafe code. While it would be ideal to use the share
/// interface, for now, juggling multi sessions should suffice.
#[derive(Clone)]
pub(crate) struct Pool {
    inner: Arc<PoolInner>,
}

impl Pool {
    pub(crate) fn new() -> Self {
        Self {
            inner: Arc::new(PoolInner::new()),
        }
    }

    pub(crate) fn multi(&self) -> PoolMulti {
        self.inner.clone().pop_or_init()
    }
}

/// A borrowed `Multi` handle that is associated with a particular `Pool`.
pub(crate) struct PoolMulti {
    pool: Arc<PoolInner>,
    entry: Option<PoolEntry>,
    valid: bool,
}

impl PoolMulti {
    pub(crate) fn get(&self) -> &Multi {
        &self.entry.as_ref().unwrap().multi
    }

    pub(crate) fn get_mut(&mut self) -> &mut Multi {
        &mut self.entry.as_mut().unwrap().multi
    }

    /// Release the `PoolMulti` without returning it to the pool.
    pub(crate) fn discard(mut self) {
        self.valid = false;
    }
}

impl Drop for PoolMulti {
    fn drop(&mut self) {
        if self.valid {
            self.pool.push(self.entry.take().unwrap());
        }
    }
}

/// Shared state between a `Pool` and its associated `PoolMulti`s.
struct PoolInner {
    heap: Mutex<BinaryHeap<PoolEntry>>,
}

impl PoolInner {
    fn new() -> Self {
        Self {
            heap: Mutex::new(BinaryHeap::new()),
        }
    }

    /// Pop an existing `Multi` or create a new one if none are available.
    fn pop_or_init(self: Arc<PoolInner>) -> PoolMulti {
        let mut heap = self.heap.lock();
        let entry = heap.pop().unwrap_or_else(PoolEntry::new);
        PoolMulti {
            pool: self.clone(),
            entry: Some(entry),
            valid: true,
        }
    }

    /// Return a `Multi` to the pool.
    fn push(&self, mut entry: PoolEntry) {
        entry.priority += 1;
        let mut heap = self.heap.lock();
        heap.push(entry);
        if heap.len() >= MAX_POOL_SIZE {
            truncate(&mut *heap, MAX_POOL_SIZE);
        }
    }
}

/// Truncate the heap to length n, dropping the items with lowest priority.
///
/// XXX: This is inefficient as it has to allocate a new vector and copy over
/// all the items (and then re-heapify them); it is mostly here just to prevent
/// unbounded memory usage in the case where the user performs a very large
/// number of concurrent requests. (It's hard to imagine this number exceeding
/// the pool size limit, but better to be safe than sorry...)
fn truncate<T: Ord>(heap: &mut BinaryHeap<T>, n: usize) {
    let mut sorted = mem::take(heap).into_sorted_vec();
    if let Some(i) = sorted.len().checked_sub(n) {
        *heap = sorted.drain(i..).collect();
    }
}

/// A `Multi` along with its priority in the pool.
struct PoolEntry {
    multi: Multi,
    priority: usize,
}

impl PoolEntry {
    fn new() -> Self {
        Self {
            multi: Multi::new(),
            priority: 0,
        }
    }
}

impl Ord for PoolEntry {
    fn cmp(&self, other: &Self) -> Ordering {
        self.priority.cmp(&other.priority)
    }
}

impl PartialOrd for PoolEntry {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for PoolEntry {
    fn eq(&self, other: &Self) -> bool {
        self.priority == other.priority
    }
}

impl Eq for PoolEntry {}

/// From [libcurl's documentation][1]:
///
/// > You must never share the same handle in multiple threads. You can pass the
/// > handles around among threads, but you must never use a single handle from
/// > more than one thread at any given time.
///
/// `Multi` does not implement `Send` or `Sync` because of the above
/// constraints. In particular, it is not `Sync` because libcurl has no
/// internal synchronization, and it is not `Send` because a single Multi
/// session can only be used to drive transfers from a single thread at
/// any one time.
///
/// In this case, when a `PoolEntry` is taken from the pool, it will only
/// be used by a single thread, and returned only when it is dropped. As
/// such, there is no risk of it being used by multiple theads at the same
/// time, so it is safe to mark it as `Send`.
///
/// [1]: https://curl.haxx.se/libcurl/c/threadsafe.html
unsafe impl Send for PoolEntry {}

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

    #[test]
    fn test_pool() {
        let pool = Pool::new();

        {
            // Get a handle for the first time.
            let multi = pool.multi();
            assert_eq!(multi.entry.as_ref().unwrap().priority, 0);
        }

        {
            // Check that we reused the existing handle.
            let multi = pool.multi();
            assert_eq!(multi.entry.as_ref().unwrap().priority, 1);
        }

        {
            // Get 2 handles. The first should be the existing one,
            // the second should be a new one.
            let multi1 = pool.multi();
            let multi2 = pool.multi();
            assert_eq!(multi1.entry.as_ref().unwrap().priority, 2);
            assert_eq!(multi2.entry.as_ref().unwrap().priority, 0);
        }

        {
            // Get two handles again. This time we should get the
            // two existing handles in order of number of uses.
            let multi1 = pool.multi();
            let multi2 = pool.multi();
            assert_eq!(multi1.entry.as_ref().unwrap().priority, 3);
            assert_eq!(multi2.entry.as_ref().unwrap().priority, 1);
        }
    }

    #[test]
    fn test_truncate() {
        let mut heap = BinaryHeap::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
        truncate(&mut heap, 5);
        assert_eq!(heap.into_sorted_vec(), vec![5, 6, 7, 8, 9]);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cell::RefCell;
use std::fmt;
use std::iter::Sum;
use std::ops::Add;
use std::ops::AddAssign;
use std::ops::Sub;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::AcqRel;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::atomic::Ordering::Release;
use std::sync::Arc;
use std::time::Instant;

use once_cell::sync::OnceCell;

#[derive(Default, Debug, Copy, Clone, Eq, PartialEq)]
pub struct Progress {
    pub downloaded: usize,
    pub total_downloaded: usize,
    pub uploaded: usize,
    pub total_uploaded: usize,
}

#[derive(Default, Debug)]
struct MutableProgress {
    downloaded: AtomicUsize,
    total_downloaded: AtomicUsize,
    uploaded: AtomicUsize,
    total_uploaded: AtomicUsize,
}

impl Progress {
    pub fn new(
        downloaded: usize,
        total_downloaded: usize,
        uploaded: usize,
        total_uploaded: usize,
    ) -> Self {
        Self {
            downloaded,
            total_downloaded,
            uploaded,
            total_uploaded,
        }
    }

    /// Create a `Progress` struct from progress values in the format provided
    /// by libcurl.
    ///
    /// For historical reasons, libcurl's CURLOPT_PROGRESSFUNCTION provides the
    /// callback with floating-point values. More recently, libcurl has added
    /// a new API called CURLOPT_XFERINFOFUNCTION that uses integers instead.
    /// Unfortunately, the Rust bindings to libcurl do not expose
    /// CURLOPT_XFERINFOFUNCTION, so we need to manually cast to `usize`.
    pub fn from_curl(dltotal: f64, dlnow: f64, ultotal: f64, ulnow: f64) -> Self {
        Self::new(
            dlnow as usize,
            dltotal as usize,
            ulnow as usize,
            ultotal as usize,
        )
    }

    pub fn as_tuple(&self) -> (usize, usize, usize, usize) {
        (
            self.downloaded,
            self.total_downloaded,
            self.uploaded,
            self.total_uploaded,
        )
    }
}

impl MutableProgress {
    fn to_progress(&self) -> Progress {
        let downloaded = self.downloaded.load(Relaxed);
        let uploaded = self.uploaded.load(Relaxed);
        let total_downloaded = self.total_downloaded.load(Acquire);
        let total_uploaded = self.total_uploaded.load(Acquire);
        Progress {
            downloaded,
            total_downloaded,
            uploaded,
            total_uploaded,
        }
    }

    fn set(&self, progress: Progress) {
        self.total_downloaded
            .store(progress.total_downloaded, Release);
        self.downloaded.store(progress.downloaded, Relaxed);
        self.total_uploaded.store(progress.total_uploaded, Release);
        self.uploaded.store(progress.uploaded, Relaxed);
    }

    fn add_assign(&self, progress: Progress) {
        self.total_downloaded
            .fetch_add(progress.total_downloaded, AcqRel);
        self.downloaded.fetch_add(progress.downloaded, Relaxed);
        self.total_uploaded
            .fetch_add(progress.total_uploaded, AcqRel);
        self.uploaded.fetch_add(progress.uploaded, Relaxed);
    }
}

impl fmt::Display for Progress {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Downloaded: {}/{} bytes; Uploaded {}/{} bytes",
            self.downloaded, self.total_downloaded, self.uploaded, self.total_uploaded
        )
    }
}

impl Add for Progress {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        Self {
            downloaded: self.downloaded + other.downloaded,
            total_downloaded: self.total_downloaded + other.total_downloaded,
            uploaded: self.uploaded + other.uploaded,
            total_uploaded: self.total_uploaded + other.total_uploaded,
        }
    }
}

impl AddAssign for Progress {
    fn add_assign(&mut self, other: Progress) {
        *self = *self + other
    }
}

impl Sub for Progress {
    type Output = Self;

    fn sub(self, other: Self) -> Self {
        Self {
            downloaded: self.downloaded.saturating_sub(other.downloaded),
            total_downloaded: self.total_downloaded.saturating_sub(other.total_downloaded),
            uploaded: self.uploaded.saturating_sub(other.uploaded),
            total_uploaded: self.total_uploaded.saturating_sub(other.total_uploaded),
        }
    }
}

impl Sum for Progress {
    fn sum<I: Iterator<Item = Progress>>(iter: I) -> Progress {
        iter.fold(Default::default(), Add::add)
    }
}

/// Struct that keeps track of the collective progress of
/// a collection of active transfers. Its main purpose is
/// to report the aggregate progress of these transfers
/// as if they were a single transfer.
pub(crate) struct ProgressReporter<P> {
    inner: Arc<ProgressInner>,
    callback: RefCell<P>,
    last_progress: MutableProgress,
}

impl<P: FnMut(Progress)> ProgressReporter<P> {
    /// Create a new progress reporter that will call the provided
    /// callback whenever one of its underlying transfers reports
    /// progress.
    pub(crate) fn with_callback(callback: P) -> Self {
        Self {
            inner: Arc::new(ProgressInner::default()),
            callback: RefCell::new(callback),
            last_progress: Default::default(),
        }
    }

    /// Allocate a slot for a new transfer in the reporter,
    /// and return an updater so that the transfer handler
    /// can update the values as the transfer makes progress.
    pub(crate) fn updater(&self) -> ProgressUpdater {
        let inner = Arc::clone(&self.inner);
        ProgressUpdater {
            inner,
            last_progress: Default::default(),
        }
    }

    /// Sum all of the progress values across all slots.
    pub(crate) fn aggregate(&self) -> Progress {
        self.inner.aggregate()
    }

    /// Report the instant at which the first byte
    /// of any of the transfers was received.
    pub(crate) fn first_byte_received(&self) -> Option<Instant> {
        self.inner.first_byte_received.get().cloned()
    }

    /// Call the user-provided progress callback if any of
    /// the transfers have made progress since the last time
    /// this method was called.
    pub(crate) fn report_if_updated(&self) {
        let inner = &self.inner;
        let progress = inner.aggregate();
        if progress != self.last_progress.to_progress() {
            (&mut *self.callback.borrow_mut())(inner.aggregate());
            self.last_progress.set(progress);
        }
    }
}

/// Handle representing a slot in the progress reporter.
/// The handle may be used to update the status of an
/// individual transfer with the reporter.
pub(crate) struct ProgressUpdater {
    inner: Arc<ProgressInner>,
    last_progress: MutableProgress,
}

impl ProgressUpdater {
    pub fn update(&self, progress: Progress) {
        self.inner
            .update(self.last_progress.to_progress(), progress);
        self.last_progress.set(progress);
    }
}

/// Shared state between the updater and reporter.
/// Those structs are expected to manager interior
/// mutability, so the methods on the shared state
/// are free to take exclusive references to self.
#[derive(Default)]
struct ProgressInner {
    total_progress: MutableProgress,
    first_byte_received: OnceCell<Instant>,
}

impl ProgressInner {
    fn update(&self, last_progress: Progress, progress: Progress) {
        if progress.downloaded > 0 {
            self.first_byte_received.get_or_init(Instant::now);
        }
        self.total_progress.add_assign(progress - last_progress);
    }

    fn aggregate(&self) -> Progress {
        self.total_progress.to_progress()
    }
}

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

    fn p(dl: usize, dlt: usize, ul: usize, ult: usize) -> Progress {
        Progress::new(dl, dlt, ul, ult)
    }

    #[test]
    fn test_progress() {
        let mut reported = Vec::new();
        let callback = |progress| {
            reported.push(progress);
        };

        let reporter = ProgressReporter::with_callback(callback);
        let updater1 = reporter.updater();
        let updater2 = reporter.updater();

        reporter.report_if_updated(); // No-op.
        assert_eq!(reporter.first_byte_received(), None);

        updater1.update(p(0, 0, 0, 0));
        assert!(reporter.first_byte_received().is_none());

        updater1.update(p(1, 2, 3, 4));
        updater2.update(p(1, 2, 3, 4));

        assert_eq!(reporter.aggregate(), p(2, 4, 6, 8));
        assert!(reporter.first_byte_received().is_some());

        reporter.report_if_updated();
        reporter.report_if_updated(); // No-op.

        updater2.update(p(4, 3, 102, 101));

        reporter.report_if_updated();

        let expected = vec![p(2, 4, 6, 8), p(5, 5, 105, 105)];
        assert_eq!(&expected, &reported);
    }
}