use http::{HeaderMap, HeaderName};
pub use tonic::metadata::MetadataMap;

const TRAILER_PREFIX: &str = "trailer-";
const CONNECT_RESERVED_PREFIX: &str = "connect-";
const TRAILER_CONNECT_PREFIX: &str = "trailer-connect-";

#[derive(Clone, Debug, Default)]
pub struct ConnectHeadersInfo {
    pub reserved: HeaderMap,
    pub trailer_names: HeaderMap<()>,
}

impl ConnectHeadersInfo {
    pub fn was_trailer(&self, key: impl AsRef<str>) -> bool {
        self.trailer_names
            .contains_key(format!("{TRAILER_PREFIX}{}", key.as_ref()))
    }
}

pub fn from_headers(
    headers: HeaderMap,
    mut headers_info: Option<&mut ConnectHeadersInfo>,
) -> MetadataMap {
    let (reserved_headers, headers) = headers.into_iter().partition_by_name(|name| {
        name.as_str().starts_with(CONNECT_RESERVED_PREFIX)
            || name.as_str().starts_with(TRAILER_CONNECT_PREFIX)
    });
    if let Some(headers_info) = &mut headers_info {
        headers_info.reserved.extend(reserved_headers);
    }

    let headers = headers
        .into_iter()
        .filter_map_names(|name| {
            if let Some(prefix) = name.as_str().strip_prefix(TRAILER_PREFIX) {
                let prefix_name: HeaderName = prefix.try_into().unwrap();
                if let Some(headers_info) = &mut headers_info {
                    headers_info.trailer_names.insert(name, ());
                }
                Some(prefix_name)
            } else {
                Some(name)
            }
        })
        .collect_map();

    MetadataMap::from_headers(headers)
}

pub(crate) trait HeaderMapIteratorExt<T> {
    /// Collects this iterator into a new [`HeaderMap`].
    fn collect_map(self) -> HeaderMap<T>;

    /// Filter and map header names.
    ///
    /// The given function is passed each _unique_ header name. If it returns
    /// Some, the returned name replaces the original in the resulting iterator.
    /// If it returns None, the header and its value(s) are dropped.
    fn filter_map_names(
        self,
        f: impl FnMut(HeaderName) -> Option<HeaderName>,
    ) -> impl Iterator<Item = (Option<HeaderName>, T)>;

    fn partition_by_name(self, f: impl FnMut(&HeaderName) -> bool) -> (HeaderMap<T>, HeaderMap<T>);
}

impl<T, U: Iterator<Item = (Option<HeaderName>, T)>> HeaderMapIteratorExt<T> for U {
    fn collect_map(self) -> HeaderMap<T> {
        let mut map = HeaderMap::default();
        map.extend(self);
        map
    }

    fn filter_map_names(
        self,
        mut f: impl FnMut(HeaderName) -> Option<HeaderName>,
    ) -> impl Iterator<Item = (Option<HeaderName>, T)> {
        let mut retain = true;
        self.filter_map(move |(maybe_name, val)| {
            let maybe_name = if let Some(name) = maybe_name {
                let mapped_name = f(name);
                retain = mapped_name.is_some();
                mapped_name
            } else {
                maybe_name
            };
            retain.then_some((maybe_name, val))
        })
    }

    fn partition_by_name(
        self,
        mut f: impl FnMut(&HeaderName) -> bool,
    ) -> (HeaderMap<T>, HeaderMap<T>) {
        let mut left = HeaderMap::default();
        let mut right = HeaderMap::default();
        let mut last_choice = true;
        let mut last_name: Option<HeaderName> = None;
        for (maybe_name, val) in self {
            if let Some(name) = maybe_name {
                last_choice = f(&name);
                last_name = Some(name);
            }
            if last_choice { &mut left } else { &mut right }
                .append(last_name.clone().unwrap(), val);
        }
        (left, right)
    }
}

pub(crate) mod serde {
    use serde::{de::MapAccess, ser::SerializeMap, Deserializer, Serializer};
    use tonic::metadata::{AsciiMetadataKey, BinaryMetadataKey, MetadataMap};

    pub fn serialize<S: Serializer>(
        metadata: &MetadataMap,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        let mut map = serializer.serialize_map(Some(metadata.keys_len()))?;
        let mut last_key = "";
        let mut values = vec![];
        for entry in metadata.iter() {
            use tonic::metadata::KeyAndValueRef::*;
            let (key, value_bytes) = match entry {
                Ascii(key, value) => (key.as_str(), value.as_encoded_bytes()),
                Binary(key, value) => (key.as_str(), value.as_encoded_bytes()),
            };
            if key != last_key {
                map.serialize_entry(last_key, &values)?;
                last_key = key;
                values.clear();
            }
            let value = std::str::from_utf8(value_bytes).map_err(serde::ser::Error::custom)?;
            values.push(value);
        }
        if !last_key.is_empty() {
            map.serialize_entry(last_key, &values)?;
        }
        map.end()
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(
        deserializer: D,
    ) -> Result<MetadataMap, D::Error> {
        struct Visitor;
        impl<'de> serde::de::Visitor<'de> for Visitor {
            type Value = MetadataMap;

            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                write!(formatter, "a map of header names to lists of header values")
            }

            fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
                use serde::de::Error;
                let mut metadata = MetadataMap::with_capacity(map.size_hint().unwrap_or_default());
                while let Some((key, values)) = map.next_entry::<String, Vec<String>>()? {
                    for value in values {
                        if key.ends_with("-bin") {
                            let key: BinaryMetadataKey = key.parse().map_err(Error::custom)?;
                            let value = Vec::from(value).try_into().map_err(Error::custom)?;
                            metadata.append_bin(key, value);
                        } else {
                            let key: AsciiMetadataKey = key.parse().map_err(Error::custom)?;
                            let value = value.try_into().map_err(Error::custom)?;
                            metadata.append(key, value);
                        }
                    }
                }
                Ok(metadata)
            }
        }
        deserializer.deserialize_map(Visitor)
    }
}