use crate::header::{Append, ByteArrayMap, SetCookie, SetCookieBuilder};
use ohkami_lib::map::TupleMap;
use std::borrow::Cow;

#[derive(Clone)]
pub struct Headers {
    standard: ByteArrayMap<N_SERVER_HEADERS, Cow<'static, str>>,
    custom: Option<Box<TupleMap<&'static str, Cow<'static, str>>>>,
    /// using `Option<Box<Vec<...>>>` to reduce the size to 1 word
    /// by the niche optimization.
    #[allow(clippy::box_collection)]
    pub(super) setcookie: Option<Box<Vec<Cow<'static, str>>>>,
    pub(super) size: usize,
}

pub struct SetHeaders<'set>(&'set mut Headers);
impl Headers {
    #[inline]
    pub fn set(&mut self) -> SetHeaders<'_> {
        SetHeaders(self)
    }
}

pub trait HeaderAction<'action> {
    fn perform(self, set: SetHeaders<'action>, key: Header) -> SetHeaders<'action>;
}
const _: () = {
    // append
    impl<'a> HeaderAction<'a> for Append {
        #[inline]
        fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
            set.0.append(key, self.0);
            set
        }
    }

    // insert
    impl<'a> HeaderAction<'a> for &'static str {
        #[inline(always)]
        fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
            set.0.insert(key, Cow::Borrowed(self));
            set
        }
    }
    impl<'a> HeaderAction<'a> for String {
        #[inline(always)]
        fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
            set.0.insert(key, Cow::Owned(self));
            set
        }
    }
    impl<'a> HeaderAction<'a> for std::borrow::Cow<'static, str> {
        fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
            set.0.insert(key, self);
            set
        }
    }

    // remove or insert
    impl<'a> HeaderAction<'a> for Option<Cow<'static, str>> {
        #[inline]
        fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
            match self {
                None => set.0.remove(key),
                Some(v) => set.0.insert(key, v),
            }
            set
        }
    }
};

pub trait CustomHeadersAction<'action> {
    fn perform(self, set: SetHeaders<'action>, key: &'static str) -> SetHeaders<'action>;
}
const _: () = {
    /* append */
    impl<'set> CustomHeadersAction<'set> for Append {
        fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
            set.0.append_custom(key, self.0);
            set
        }
    }

    /* insert */
    impl<'set> CustomHeadersAction<'set> for &'static str {
        #[inline(always)]
        fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
            set.0.insert_custom(key, Cow::Borrowed(self));
            set
        }
    }
    impl<'set> CustomHeadersAction<'set> for String {
        #[inline(always)]
        fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
            set.0.insert_custom(key, Cow::Owned(self));
            set
        }
    }
    impl<'set> CustomHeadersAction<'set> for Cow<'static, str> {
        fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
            set.0.insert_custom(key, self);
            set
        }
    }

    /* remove or insert */
    impl<'set> CustomHeadersAction<'set> for Option<Cow<'static, str>> {
        #[inline]
        fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
            match self {
                None => set.0.remove_custom(key),
                Some(v) => set.0.insert_custom(key, v),
            }
            set
        }
    }
};

macro_rules! Header {
    ($N:literal; $( $method:ident($konst:ident): $name_bytes:literal, )*) => {
        pub(crate) const N_SERVER_HEADERS: usize = $N;
        const _: [Header; N_SERVER_HEADERS] = [$(Header::$konst),*];

        #[derive(Debug, PartialEq, Clone, Copy)]
        pub enum Header {
            $( $konst, )*
        }

        impl Header {
            #[inline] pub const fn as_bytes(&self) -> &'static [u8] {
                match self {
                    $(
                        Self::$konst => $name_bytes,
                    )*
                }
            }
            pub const fn as_str(&self) -> &'static str {
                unsafe {std::str::from_utf8_unchecked(self.as_bytes())}
            }
            #[inline(always)] const fn len(&self) -> usize {
                match self {
                    $(
                        Self::$konst => $name_bytes.len(),
                    )*
                }
            }

            // Mainly used in tests
            pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
                (0..N_SERVER_HEADERS)
                    .map(|i| unsafe {std::mem::transmute::<_, Header>(i as u8)})
                    .find(|h| h.as_bytes().eq_ignore_ascii_case(bytes))
            }
        }

        impl<T: AsRef<[u8]>> PartialEq<T> for Header {
            fn eq(&self, other: &T) -> bool {
                self.as_bytes().eq_ignore_ascii_case(other.as_ref())
            }
        }

        #[allow(non_snake_case)]
        impl<'set> SetHeaders<'set> {
            $(
                #[inline]
                pub fn $method(self, action: impl HeaderAction<'set>) -> Self {
                    action.perform(self, Header::$konst)
                }

                #[deprecated = "Use snake_case method instead"]
                pub fn $konst(self, action: impl HeaderAction<'set>) -> Self {
                    self.$method(action)
                }
            )*

            #[deprecated = "use `.x` instead"]
            pub fn custom(self, name: &'static str, action: impl CustomHeadersAction<'set>) -> Self {
                self.x(name, action)
            }
            #[inline]
            pub fn x(self, name: &'static str, action: impl CustomHeadersAction<'set>) -> Self {
                action.perform(self, name)
            }
        }

        #[allow(non_snake_case)]
        impl Headers {
            $(
                #[inline]
                pub fn $method(&self) -> Option<&str> {
                    self.get_standard(Header::$konst)
                }

                #[deprecated = "Use snake_case method instead"]
                pub fn $konst(&self) -> Option<&str> {
                    self.$method()
                }
            )*

            #[deprecated = "use `.get` instead"]
            pub fn custom(&self, name: &'static str) -> Option<&str> {
                self.get(name)
            }
            #[inline]
            pub fn get(&self, name: &'static str) -> Option<&str> {
                self.get_custom(name)
                    .or_else(|| self.get_standard(Header::from_bytes(name.as_bytes())?))
            }
        }
    };
}
Header! {48;
    accept_ranges(AcceptRanges): b"Accept-Ranges",
    access_control_allow_credentials(AccessControlAllowCredentials): b"Access-Control-Allow-Credentials",
    access_control_allow_headers(AccessControlAllowHeaders): b"Access-Control-Allow-Headers",
    access_control_allow_methods(AccessControlAllowMethods): b"Access-Control-Allow-Methods",
    access_control_allow_origin(AccessControlAllowOrigin): b"Access-Control-Allow-Origin",
    access_control_expose_headers(AccessControlExposeHeaders): b"Access-Control-Expose-Headers",
    access_control_max_age(AccessControlMaxAge): b"Access-Control-Max-Age",
    age(Age): b"Age",
    allow(Allow): b"Allow",
    alt_svc(AltSvc): b"Alt-Svc",
    cache_control(CacheControl): b"Cache-Control",
    cache_status(CacheStatus): b"Cache-Status",
    cdn_cache_control(CDNCacheControl): b"CDN-Cache-Control",
    connection(Connection): b"Connection",
    content_dispotision(ContentDisposition): b"Content-Disposition",
    content_encoding(ContentEncoding): b"Content-Encoding",
    content_language(ContentLanguage): b"Content-Language",
    content_length(ContentLength): b"Content-Length",
    content_location(ContentLocation): b"Content-Location",
    content_range(ContentRange): b"Content-Range",
    content_security_policy(ContentSecurityPolicy): b"Content-Security-Policy",
    content_security_policy_report_only(ContentSecurityPolicyReportOnly): b"Content-Security-Policy-Report-Only",
    content_type(ContentType): b"Content-Type",
    cross_origin_embedder_policy(CrossOriginEmbedderPolicy): b"Cross-Origin-Embedder-Policy",
    cross_origin_resource_policy(CrossOriginResourcePolicy): b"Cross-Origin-Resource-Policy",
    date(Date): b"Date",
    etag(ETag): b"ETag",
    expires(Expires): b"Expires",
    link(Link): b"Link",
    location(Location): b"Location",
    last_modified(LastModified): b"Last-Modified",
    proxy_authenticate(ProxyAuthenticate): b"Proxy-Authenticate",
    referrer_policy(ReferrerPolicy): b"Referrer-Policy",
    refresh(Refresh): b"Refresh",
    retry_after(RetryAfter): b"Retry-After",
    sec_websocket_accept(SecWebSocketAccept): b"Sec-WebSocket-Accept",
    sec_websocket_protocol(SecWebSocketProtocol): b"Sec-WebSocket-Protocol",
    sec_websocket_version(SecWebSocketVersion): b"Sec-WebSocket-Version",
    server(Server): b"Server",
    strict_transport_security(StrictTransportSecurity): b"Strict-Transport-Security",
    trailer(Trailer): b"Trailer",
    transfer_encoding(TransferEncoding): b"Transfer-Encoding",
    upgrade(Upgrade): b"Upgrade",
    vary(Vary): b"Vary",
    via(Via): b"Via",
    x_content_type_options(XContentTypeOptions): b"X-Content-Type-Options",
    x_frame_options(XFrameOptions): b"X-Frame-Options",
    www_authenticate(WWWAuthenticate): b"WWW-Authenticate",
}

const _: () = {
    #[allow(non_snake_case)]
    impl Headers {
        pub fn set_cookie(&self) -> impl Iterator<Item = SetCookie<'_>> {
            self.setcookie
                .as_ref()
                .map(|setcookies| {
                    setcookies
                        .iter()
                        .filter_map(|raw| match SetCookie::from_raw(raw) {
                            Ok(valid) => Some(valid),
                            Err(_err) => {
                                #[cfg(debug_assertions)]
                                crate::WARNING!("Invalid `Set-Cookie`: {_err}");
                                None
                            }
                        })
                })
                .into_iter()
                .flatten()
        }
    }

    #[allow(non_snake_case)]
    impl<'s> SetHeaders<'s> {
        /// Add new `Set-Cookie` header to the response.
        ///
        /// - When you call this N times, the response has N different
        ///   `Set-Cookie` headers.
        /// - Cookie value (second argument) is percent-encoded when the
        ///   response is sended.
        ///
        /// ---
        /// *example.rs*
        /// ```
        /// use ohkami::Response;
        ///
        /// fn mutate_header(res: &mut Response) {
        ///     res.headers.set()
        ///         .server("ohkami")
        ///         .set_cookie("id", "42", |d|d.path("/").same_site_strict())
        ///         .set_cookie("name", "John", |d|d.path("/where").same_site_lax());
        /// }
        /// ```
        #[inline]
        pub fn set_cookie(
            self,
            name: &'static str,
            value: impl Into<Cow<'static, str>>,
            directives: impl FnOnce(SetCookieBuilder) -> SetCookieBuilder,
        ) -> Self {
            let setcookie: Cow<'static, str> = directives(SetCookieBuilder::new(name, value))
                .build()
                .into();
            self.0.size += "Set-Cookie: ".len() + setcookie.len() + "\r\n".len();
            match self.0.setcookie.as_mut() {
                None => self.0.setcookie = Some(Box::new(vec![setcookie])),
                Some(setcookies) => setcookies.push(setcookie),
            }
            self
        }
    }
};

impl Headers {
    #[inline(always)]
    pub(crate) fn insert(&mut self, name: Header, value: Cow<'static, str>) {
        let (name_len, value_len) = (name.len(), value.len());
        match unsafe { self.standard.get_mut(name as u8) } {
            None => {
                unsafe { self.standard.insert_new(name as u8, value) };
                self.size += name_len + ": ".len() + value_len + "\r\n".len();
            }
            Some(old) => {
                self.size -= old.len();
                self.size += value_len;
                *old = value
            }
        }
    }
    #[inline]
    pub(crate) fn insert_custom(&mut self, name: &'static str, value: Cow<'static, str>) {
        let self_len = value.len();
        match &mut self.custom {
            None => {
                self.custom = Some(Box::new(TupleMap::from_iter([(name, value)])));
                self.size += name.len() + ": ".len() + self_len + "\r\n".len()
            }
            Some(custom) => {
                if let Some(old) = custom.insert(name, value) {
                    self.size -= old.len();
                    self.size += self_len;
                } else {
                    self.size += name.len() + ": ".len() + self_len + "\r\n".len()
                }
            }
        }
    }

    #[inline]
    pub(crate) fn remove(&mut self, name: Header) {
        let name_len = name.len();
        if let Some(v) = unsafe { self.standard.get(name as u8) } {
            self.size -= name_len + ": ".len() + v.len() + "\r\n".len();
            unsafe { self.standard.delete(name as u8) };
        }
    }
    pub(crate) fn remove_custom(&mut self, name: &'static str) {
        if let Some(c) = self.custom.as_mut()
            && let Some(v) = c.remove(name)
        {
            self.size -= name.len() + ": ".len() + v.len() + "\r\n".len();
        }
    }

    #[inline(always)]
    pub(crate) fn get_standard(&self, name: Header) -> Option<&str> {
        unsafe { self.standard.get(name as u8) }.map(Cow::as_ref)
    }
    #[inline]
    pub(crate) fn get_custom(&self, name: &'static str) -> Option<&str> {
        self.custom.as_ref()?.get(&name).map(Cow::as_ref)
    }

    pub(crate) fn append(&mut self, name: Header, value: Cow<'static, str>) {
        let value_len = value.len();
        let target = unsafe { self.standard.get_mut(name as u8) };

        self.size += match target {
            Some(v) => {
                match v {
                    Cow::Borrowed(slice) => {
                        let mut appended = String::with_capacity(slice.len() + 2 + value_len);
                        appended.push_str(slice);
                        appended.push_str(", ");
                        appended.push_str(&value);
                        *v = Cow::Owned(appended);
                    }
                    Cow::Owned(string) => {
                        string.push_str(", ");
                        string.push_str(&value);
                    }
                }
                ", ".len() + value_len
            }
            None => {
                unsafe { self.standard.insert_new(name as u8, value) }
                name.len() + ": ".len() + value_len + "\r\n".len()
            }
        };
    }
    pub(crate) fn append_custom(&mut self, name: &'static str, value: Cow<'static, str>) {
        let value_len = value.len();

        let custom = {
            if self.custom.is_none() {
                self.custom = Some(Box::new(TupleMap::new()));
            }
            unsafe { self.custom.as_mut().unwrap_unchecked() }
        };

        self.size += match custom.get_mut(&name) {
            Some(v) => {
                match v {
                    Cow::Owned(string) => {
                        string.push_str(", ");
                        string.push_str(&value);
                    }
                    Cow::Borrowed(s) => {
                        let mut s = s.to_string();
                        s.push_str(", ");
                        s.push_str(&value);
                        *v = Cow::Owned(s);
                    }
                }
                ", ".len() + value_len
            }
            None => {
                custom.insert(name, value);
                name.len() + ": ".len() + value_len + "\r\n".len()
            }
        };
    }
}

impl Headers {
    #[inline]
    pub(crate) fn new() -> Self {
        let mut this = Self {
            standard: ByteArrayMap::new(),
            custom: None,
            setcookie: None,
            size: "\r\n".len(),
        };

        this.set()
            .date(ohkami_lib::imf_fixdate(crate::util::unix_timestamp()))
            .content_length("0");

        this
    }
    #[cfg(feature = "DEBUG")]
    #[doc(hidden)]
    pub fn _new() -> Self {
        Self::new()
    }

    pub(crate) fn iter_standard(&self) -> impl Iterator<Item = (&str, &str)> {
        self.standard.iter().map(|(i, v)| {
            (
                unsafe { std::mem::transmute::<u8, Header>(*i) }.as_str(),
                &**v,
            )
        })
    }

    pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
        self.iter_standard()
            .chain(
                self.custom
                    .as_ref()
                    .into_iter()
                    .flat_map(|hm| hm.iter().map(|(k, v)| (*k, &**v))),
            )
            .chain(
                self.setcookie
                    .as_ref()
                    .map(|sc| sc.iter().map(Cow::as_ref))
                    .into_iter()
                    .flatten()
                    .map(|sc| ("Set-Cookie", sc)),
            )
    }

    #[cfg(any(feature = "__rt_native__", feature = "DEBUG"))]
    /// SAFETY: `buf` has remaining capacity of at least `self.size`
    pub(crate) unsafe fn write_unchecked_to(&self, buf: &mut Vec<u8>) {
        unsafe {
            for (i, v) in self.standard.iter() {
                let h = std::mem::transmute::<u8, Header>(*i);
                {
                    crate::push_unchecked!(buf <- h.as_bytes());
                    crate::push_unchecked!(buf <- b": ");
                    crate::push_unchecked!(buf <- v.as_bytes());
                    crate::push_unchecked!(buf <- b"\r\n");
                }
            }
            if let Some(custom) = self.custom.as_ref() {
                for (k, v) in custom.iter() {
                    crate::push_unchecked!(buf <- k.as_bytes());
                    crate::push_unchecked!(buf <- b": ");
                    crate::push_unchecked!(buf <- v.as_bytes());
                    crate::push_unchecked!(buf <- b"\r\n");
                }
            }
            if let Some(setcookies) = self.setcookie.as_ref() {
                for setcookie in &**setcookies {
                    crate::push_unchecked!(buf <- b"Set-Cookie: ");
                    crate::push_unchecked!(buf <- setcookie.as_bytes());
                    crate::push_unchecked!(buf <- b"\r\n");
                }
            }
            crate::push_unchecked!(buf <- b"\r\n");
        }
    }

    #[cfg(any(feature = "DEBUG", feature = "__rt_native__"))]
    pub fn _write_to(&self, buf: &mut Vec<u8>) {
        buf.reserve(self.size);
        unsafe { self.write_unchecked_to(buf) }
    }
}

const _: () = {
    impl std::fmt::Debug for Headers {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_map().entries(self.iter()).finish()
        }
    }

    impl PartialEq for Headers {
        fn eq(&self, other: &Self) -> bool {
            for (k, v) in self.iter_standard() {
                if other.get_standard(Header::from_bytes(k.as_bytes()).unwrap()) != Some(v) {
                    return false;
                }
            }

            if self.custom != other.custom {
                return false;
            }

            true
        }
    }

    impl<A: Into<Cow<'static, str>>> FromIterator<(&'static str, A)> for Headers {
        fn from_iter<T: IntoIterator<Item = (&'static str, A)>>(iter: T) -> Self {
            let mut this = Headers::new();
            for (k, v) in iter {
                match Header::from_bytes(k.as_bytes()) {
                    Some(h) => this.insert(h, v.into()),
                    None => {
                        this.set().x(k, v.into());
                    }
                }
            }
            this
        }
    }

    impl IntoIterator for Headers {
        type Item = (&'static str, Cow<'static, str>);
        type IntoIter = HeadersIntoIter;
        fn into_iter(self) -> Self::IntoIter {
            HeadersIntoIter {
                standard: self.standard.into_iter(),
                custom: self.custom.map(|tm| tm.into_iter()),
                setcookie: self.setcookie.map(|sc| sc.into_iter()),
            }
        }
    }
    pub struct HeadersIntoIter {
        standard: std::vec::IntoIter<(u8, Cow<'static, str>)>,
        custom: Option<std::vec::IntoIter<(&'static str, Cow<'static, str>)>>,
        setcookie: Option<std::vec::IntoIter<Cow<'static, str>>>,
    }
    impl Iterator for HeadersIntoIter {
        type Item = (&'static str, Cow<'static, str>);
        fn next(&mut self) -> Option<Self::Item> {
            if let Some((i, v)) = self.standard.next() {
                Some((unsafe { std::mem::transmute::<u8, Header>(i) }.as_str(), v))
            } else if let Some(custom) = self.custom.as_mut().and_then(|c| c.next()) {
                Some(custom)
            } else if let Some(setcookie) = self.setcookie.as_mut().and_then(|c| c.next()) {
                Some(("Set-Cookie", setcookie))
            } else {
                None
            }
        }
    }
};

#[cfg(feature = "rt_worker")]
const _: () = {
    impl From<Headers> for ::worker::Headers {
        #[inline(always)]
        fn from(this: Headers) -> ::worker::Headers {
            let h = ::worker::Headers::new();
            for (k, v) in this.iter() {
                if let Err(_e) = h.append(k, v) {
                    crate::DEBUG!("`worker::Headers::append` failed: {_e:?}");
                }
            }
            h
        }
    }
};
