use miette::{IntoDiagnostic, Result, WrapErr};
use smol_str::{SmolStr, ToSmolStr};
use std::collections::BTreeSet;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::str::FromStr;
use url::Url;

use super::Hydrant;
use crate::db::keys;
use crate::pds_meta::HostStatus;

/// unix sockets are owner-only unless told otherwise, since the api has no auth of its own.
pub const DEFAULT_UNIX_SOCKET_MODE: u32 = 0o600;

/// one address the api listens on.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ApiBind {
    Tcp(SocketAddr),
    Unix(PathBuf),
}

impl ApiBind {
    pub fn tcp(&self) -> Option<SocketAddr> {
        match self {
            Self::Tcp(addr) => Some(*addr),
            Self::Unix(_) => None,
        }
    }
}

impl From<SocketAddr> for ApiBind {
    fn from(addr: SocketAddr) -> Self {
        Self::Tcp(addr)
    }
}

impl FromStr for ApiBind {
    type Err = miette::Report;

    /// parses `host:port`, or `unix:` followed by a socket path.
    fn from_str(s: &str) -> Result<Self> {
        match s.strip_prefix("unix:") {
            Some("") => miette::bail!("`unix:` bind is missing a socket path"),
            Some(path) => Ok(Self::Unix(path.into())),
            None => s
                .parse()
                .map(Self::Tcp)
                .map_err(|e| miette::miette!("invalid bind address `{s}`: {e}")),
        }
    }
}

impl std::fmt::Display for ApiBind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Tcp(addr) => write!(f, "{addr}"),
            Self::Unix(path) => write!(f, "unix:{}", path.display()),
        }
    }
}

#[derive(Debug, Clone)]
pub struct ApiBinds {
    pub(crate) head: ApiBind,
    pub(crate) tail: Vec<ApiBind>,
    pub(crate) unix_mode: u32,
    pub(crate) tcp_management: bool,
}

impl ApiBinds {
    pub fn new(bind: impl Into<ApiBind>) -> Self {
        Self {
            head: bind.into(),
            tail: Vec::new(),
            unix_mode: DEFAULT_UNIX_SOCKET_MODE,
            tcp_management: false,
        }
    }

    pub fn try_from_iter<I>(iter: I) -> Option<Self>
    where
        I: IntoIterator,
        I::Item: Into<ApiBind>,
    {
        let mut iter = iter.into_iter().map(Into::into);
        let head = iter.next()?;
        Some(Self {
            head,
            tail: iter.collect(),
            unix_mode: DEFAULT_UNIX_SOCKET_MODE,
            tcp_management: false,
        })
    }

    /// sets the permissions unix sockets are created with. only the `0o777` bits are used.
    pub fn with_unix_mode(mut self, mode: u32) -> Self {
        self.unix_mode = mode & 0o777;
        self
    }

    /// serves the management endpoints on tcp binds too. they're unix-only by default, since
    /// they have no auth and a unix socket's permissions are the only thing guarding them.
    pub fn with_tcp_management(mut self, enabled: bool) -> Self {
        self.tcp_management = enabled;
        self
    }

    /// whether any bind serves the management endpoints.
    pub fn serves_management(&self) -> bool {
        self.tcp_management || self.iter().any(|bind| matches!(bind, ApiBind::Unix(_)))
    }

    pub fn iter(&self) -> impl Iterator<Item = &ApiBind> + '_ {
        std::iter::once(&self.head).chain(self.tail.iter())
    }

    /// reads `HYDRANT_API_BIND`, `HYDRANT_API_SOCKET_MODE` and `HYDRANT_API_TCP_MANAGEMENT`
    /// through `lookup`. an unset bind gives `default`, and `none` gives `None`, meaning no api
    /// at all.
    pub fn from_lookup(
        lookup: impl Fn(&str) -> Option<String>,
        default: Option<ApiBinds>,
    ) -> Result<Option<Self>> {
        let binds = match lookup("HYDRANT_API_BIND") {
            Some(raw) if raw.trim().eq_ignore_ascii_case("none") => None,
            Some(raw) => Some(
                raw.parse::<Self>()
                    .map_err(|e| miette::miette!("invalid HYDRANT_API_BIND: {e}"))?,
            ),
            None => default,
        };
        let Some(mut binds) = binds else {
            return Ok(None);
        };
        if let Some(mode) = lookup("HYDRANT_API_SOCKET_MODE") {
            binds = binds.with_unix_mode(parse_socket_mode(&mode)?);
        }
        let tcp_management =
            crate::config::env::read(&lookup, "HYDRANT_API_TCP_MANAGEMENT", binds.tcp_management)?;
        Ok(Some(binds.with_tcp_management(tcp_management)))
    }
}

/// octal permission bits like `0660`, with or without the leading `0` or `0o`.
fn parse_socket_mode(s: &str) -> Result<u32> {
    let digits = s.trim();
    let digits = digits.strip_prefix("0o").unwrap_or(digits);
    u32::from_str_radix(digits, 8)
        .ok()
        .filter(|mode| !digits.starts_with('+') && *mode <= 0o777)
        .ok_or_else(|| {
            miette::miette!(
                "invalid HYDRANT_API_SOCKET_MODE `{s}`: want octal permissions like 0660"
            )
        })
}

impl FromStr for ApiBinds {
    type Err = miette::Report;

    /// parses a comma-separated list of [`ApiBind`]s, the format of `HYDRANT_API_BIND`.
    fn from_str(s: &str) -> Result<Self> {
        let binds = s
            .split(',')
            .map(str::trim)
            .filter(|s| !s.is_empty())
            .map(str::parse::<ApiBind>)
            .collect::<Result<Vec<_>>>()?;
        Self::try_from_iter(binds).ok_or_else(|| miette::miette!("no addresses to bind to"))
    }
}

#[derive(Debug, Clone)]
/// information about a host hydrant is consuming from.
pub struct Host {
    /// hostname of the host.
    pub name: SmolStr,
    /// latest seq hydrant has processed from this host.
    pub seq: i64,
    /// the amount of accounts hydrant has seen from this host.
    pub account_count: u64,
    /// the status of this host in hydrant.
    pub status: crate::pds_meta::HostStatus,
}

impl Hydrant {
    /// get the status of a (firehose) host we are consuming from.
    ///
    /// returns the seq we are on for this host.
    pub async fn get_host_status(&self, hostname: &str) -> Result<Option<Host>> {
        let state = self.state.clone();
        let hostname = hostname.to_smolstr();

        let state_closure = state.clone();
        state
            .db
            .run(move |db| {
                let key = keys::firehose_cursor_key(&hostname);

                let mut seq = 0;
                if let Some(cursor_bytes) = db.cursors.get(&key).into_diagnostic()? {
                    seq = i64::from_be_bytes(cursor_bytes.as_ref().try_into().into_diagnostic()?);
                } else {
                    // if it has no cursor, check if it's explicitly tracked in hosts map
                    // or firehose tasks (recently added via API but no messages yet)
                    let meta = state_closure.pds_meta.load();
                    if !meta.hosts.contains_key(hostname.as_str()) {
                        // we should also allow it if it's an active firehose ingestor
                        let mut found_in_cursors = false;
                        state_closure.firehose_cursors.iter_sync(|u, _| {
                            if u.host_str() == Some(hostname.as_str()) {
                                found_in_cursors = true;
                            }
                            !found_in_cursors // continue if not found
                        });

                        if !found_in_cursors {
                            return Ok(None);
                        }
                    }
                }

                let account_count = db.get_count_sync(&keys::pds_account_count_key(&hostname));
                let status = state_closure.pds_meta.load().status(&hostname);
                Ok(Some(Host {
                    name: hostname,
                    seq,
                    account_count,
                    status,
                }))
            })
            .await
    }

    /// enumerates all hosts hydrant is consuming from.
    ///
    /// returns hosts enumerated in this pagination and the cursor to paginate from.
    pub async fn list_hosts(
        &self,
        cursor: Option<&str>,
        limit: usize,
    ) -> Result<(Vec<Host>, Option<SmolStr>)> {
        let state = self.state.clone();
        let cursor = cursor.map(str::to_string);
        let state_closure = state.clone();

        state
            .db
            .run(move |db| {
                let start_bound = match &cursor {
                    Some(after) => std::ops::Bound::Included(keys::firehose_cursor_key(after)),
                    None => std::ops::Bound::Included(keys::FIREHOSE_CURSOR_PREFIX.to_vec()),
                };

                let prefix_end = {
                    let mut end = keys::FIREHOSE_CURSOR_PREFIX.to_vec();
                    *end.last_mut().unwrap() += 1;
                    end
                };
                let end_bound = std::ops::Bound::Excluded(prefix_end);

                let mut db_hosts = Vec::new();
                for item in db.cursors.range((start_bound, end_bound)) {
                    let (k, _) = item.into_inner().into_diagnostic()?;
                    let hostname = std::str::from_utf8(&k[keys::FIREHOSE_CURSOR_PREFIX.len()..])
                        .into_diagnostic()
                        .wrap_err("firehose cursor key contains non-utf8 hostname")?;

                    if cursor
                        .as_ref()
                        .is_some_and(|after| hostname <= after.as_str())
                    {
                        continue;
                    }

                    db_hosts.push(SmolStr::new(hostname));
                    if db_hosts.len() > limit {
                        break;
                    }
                }

                let mut meta_hosts = Vec::new();
                {
                    let meta = state_closure.pds_meta.load();
                    for hostname in meta.hosts.keys() {
                        if cursor
                            .as_ref()
                            .is_some_and(|after| hostname.as_str() <= after.as_str())
                        {
                            continue;
                        }
                        meta_hosts.push(SmolStr::new(hostname));
                    }
                }
                meta_hosts.sort();
                meta_hosts.truncate(limit + 1);

                let mut merged = BTreeSet::new();
                for host in db_hosts {
                    merged.insert(host);
                }
                for host in meta_hosts {
                    merged.insert(host);
                }

                let selected: Vec<SmolStr> = merged.into_iter().take(limit + 1).collect();

                let mut hosts: Vec<Host> = Vec::with_capacity(selected.len().min(limit));
                for hostname in selected.iter().take(limit) {
                    let seq = db
                        .cursors
                        .get(keys::firehose_cursor_key(hostname))
                        .into_diagnostic()?
                        .map(|v| {
                            v.as_ref()
                                .try_into()
                                .into_diagnostic()
                                .wrap_err("cursor value is not 8 bytes")
                                .map(i64::from_be_bytes)
                        })
                        .transpose()?
                        .unwrap_or(0);
                    let account_count = db.get_count_sync(&keys::pds_account_count_key(hostname));
                    let status = state_closure.pds_meta.load().status(hostname);
                    hosts.push(Host {
                        name: hostname.clone(),
                        seq,
                        account_count,
                        status,
                    });
                }

                let next_cursor = if selected.len() > limit {
                    hosts.last().map(|h| h.name.clone())
                } else {
                    None
                };

                Ok((hosts, next_cursor))
            })
            .await
    }
}

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

    #[test]
    fn parses_tcp_and_unix_binds() {
        assert_eq!(
            "127.0.0.1:3000".parse::<ApiBind>().unwrap(),
            ApiBind::Tcp("127.0.0.1:3000".parse().unwrap())
        );
        assert_eq!(
            "[::1]:3000".parse::<ApiBind>().unwrap(),
            ApiBind::Tcp("[::1]:3000".parse().unwrap())
        );
        assert_eq!(
            "unix:/run/hydrant/api.sock".parse::<ApiBind>().unwrap(),
            ApiBind::Unix("/run/hydrant/api.sock".into())
        );
        assert!("unix:".parse::<ApiBind>().is_err());
        assert!("localhost:3000".parse::<ApiBind>().is_err());
        assert!("/run/hydrant/api.sock".parse::<ApiBind>().is_err());
    }

    #[test]
    fn display_round_trips() {
        for s in ["127.0.0.1:3000", "[::1]:3000", "unix:/run/hydrant/api.sock"] {
            assert_eq!(s.parse::<ApiBind>().unwrap().to_string(), s);
        }
    }

    #[test]
    fn parses_bind_lists() {
        let binds: ApiBinds = " 127.0.0.1:3000, unix:/tmp/api.sock ,".parse().unwrap();
        assert_eq!(
            binds.iter().cloned().collect::<Vec<_>>(),
            [
                ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()),
                ApiBind::Unix("/tmp/api.sock".into()),
            ]
        );
        assert!("".parse::<ApiBinds>().is_err());
        assert!(" , ".parse::<ApiBinds>().is_err());
        assert!("127.0.0.1:3000,nope".parse::<ApiBinds>().is_err());
    }

    fn lookup<'a>(pairs: &'a [(&str, &str)]) -> impl Fn(&str) -> Option<String> + 'a {
        |key| {
            pairs
                .iter()
                .find(|(k, _)| *k == key)
                .map(|(_, v)| v.to_string())
        }
    }

    #[test]
    fn from_lookup_reads_bind_and_socket_mode() {
        let default = ApiBinds::new(ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()));
        let unset = ApiBinds::from_lookup(lookup(&[]), Some(default.clone())).unwrap();
        assert_eq!(unset.unwrap().iter().count(), 1);
        assert!(ApiBinds::from_lookup(lookup(&[]), None).unwrap().is_none());

        for none in ["none", " NONE "] {
            let binds =
                ApiBinds::from_lookup(lookup(&[("HYDRANT_API_BIND", none)]), Some(default.clone()));
            assert!(binds.unwrap().is_none(), "{none:?} should turn the api off");
        }

        let unix =
            ApiBinds::from_lookup(lookup(&[("HYDRANT_API_BIND", "unix:/tmp/api.sock")]), None);
        assert_eq!(unix.unwrap().unwrap().unix_mode, 0o600);

        for (mode, want) in [
            ("0660", 0o660),
            ("660", 0o660),
            ("0o640", 0o640),
            ("0777", 0o777),
        ] {
            let pairs = [
                ("HYDRANT_API_BIND", "unix:/tmp/api.sock"),
                ("HYDRANT_API_SOCKET_MODE", mode),
            ];
            let binds = ApiBinds::from_lookup(lookup(&pairs), None)
                .unwrap()
                .unwrap();
            assert_eq!(binds.unix_mode, want, "mode {mode}");
        }
    }

    #[test]
    fn management_is_unix_only_unless_enabled_for_tcp() {
        let tcp = ApiBinds::new(ApiBind::Tcp("127.0.0.1:3000".parse().unwrap()));
        assert!(!tcp.tcp_management);
        assert!(!tcp.serves_management());
        assert!(tcp.clone().with_tcp_management(true).serves_management());
        let both: ApiBinds = "127.0.0.1:3000,unix:/tmp/api.sock".parse().unwrap();
        assert!(both.serves_management());

        // the flag applies to the default binds too, not just ones set in HYDRANT_API_BIND
        let enabled = ApiBinds::from_lookup(
            lookup(&[("HYDRANT_API_TCP_MANAGEMENT", "true")]),
            Some(tcp.clone()),
        );
        assert!(enabled.unwrap().unwrap().tcp_management);
        let disabled = ApiBinds::from_lookup(
            lookup(&[("HYDRANT_API_TCP_MANAGEMENT", "false")]),
            Some(tcp),
        );
        assert!(!disabled.unwrap().unwrap().tcp_management);
    }

    #[test]
    fn from_lookup_refuses_bad_values() {
        let sock = ("HYDRANT_API_BIND", "unix:/tmp/api.sock");
        for bad in [
            &[("HYDRANT_API_BIND", "")][..],
            &[("HYDRANT_API_BIND", "none,unix:/tmp/api.sock")],
            &[sock, ("HYDRANT_API_SOCKET_MODE", "0800")],
            &[sock, ("HYDRANT_API_SOCKET_MODE", "1777")],
            &[sock, ("HYDRANT_API_SOCKET_MODE", "+660")],
            &[sock, ("HYDRANT_API_SOCKET_MODE", "rw")],
            &[sock, ("HYDRANT_API_TCP_MANAGEMENT", "yes")],
        ] {
            assert!(
                ApiBinds::from_lookup(lookup(bad), None).is_err(),
                "{bad:?} should be refused"
            );
        }
    }

    #[test]
    fn with_unix_mode_keeps_only_permission_bits() {
        let binds = ApiBinds::new(ApiBind::Unix("/tmp/api.sock".into())).with_unix_mode(0o4660);
        assert_eq!(binds.unix_mode, 0o660);
    }

    #[test]
    fn binds_can_be_unix_only() {
        let binds = ApiBinds::new(ApiBind::Unix("/tmp/api.sock".into()));
        assert_eq!(binds.iter().filter_map(ApiBind::tcp).count(), 0);
        assert_eq!(binds.iter().count(), 1);
    }
}

#[cfg(test)]
mod host_listing_tests {
    use super::*;
    use crate::config::Config;
    use crate::db::set_ks_count;
    use crate::pds_meta::HostStatus;

    fn test_config(path: &std::path::Path) -> Config {
        Config {
            database_path: path.to_path_buf(),
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn list_hosts_includes_seeded_hosts_without_cursors() -> Result<()> {
        let tmp = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(tmp.path())).await?;

        {
            let state = hydrant.state.clone();
            state
                .db
                .run(move |db| -> Result<()> {
                    let mut batch = db.inner.batch();
                    crate::db::pds_meta::set_status(
                        &mut batch,
                        &db.filter,
                        "offline.example",
                        HostStatus::Offline,
                    )?;
                    crate::db::pds_meta::set_status(
                        &mut batch,
                        &db.filter,
                        "active.example",
                        HostStatus::Active,
                    )?;
                    set_ks_count(
                        &mut batch,
                        db,
                        &keys::pds_account_count_key("offline.example"),
                        7,
                    );
                    set_ks_count(
                        &mut batch,
                        db,
                        &keys::pds_account_count_key("active.example"),
                        42,
                    );
                    db.cursors
                        .insert(
                            keys::firehose_cursor_key("active.example"),
                            123_i64.to_be_bytes(),
                        )
                        .into_diagnostic()?;
                    batch.commit().into_diagnostic()?;
                    db.persist()
                })
                .await?;

            crate::pds_meta::PdsMeta::update_host(
                &hydrant.state.pds_meta,
                "offline.example",
                |h| {
                    h.status = HostStatus::Offline;
                },
            );
            crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "active.example", |h| {
                h.status = HostStatus::Active;
            });
            hydrant.state.db.update_count("p|offline.example", 7);
            hydrant.state.db.update_count("p|active.example", 42);
        }

        let (hosts, next) = hydrant.list_hosts(None, 10).await?;
        assert!(next.is_none());

        let offline = hosts
            .iter()
            .find(|h| h.name == "offline.example")
            .expect("seeded host without cursor should be listed");
        let active = hosts
            .iter()
            .find(|h| h.name == "active.example")
            .expect("host with cursor should be listed");

        assert_eq!(offline.seq, 0);
        assert_eq!(offline.account_count, 7);
        assert_eq!(offline.status, HostStatus::Offline);

        assert_eq!(active.seq, 123);
        assert_eq!(active.account_count, 42);
        assert_eq!(active.status, HostStatus::Active);

        let status = hydrant.get_host_status("active.example").await?;
        assert_eq!(status.map(|host| host.seq), Some(123));

        Ok(())
    }

    #[tokio::test]
    async fn list_hosts_pagination_correctness() -> Result<()> {
        let tmp = tempfile::tempdir().into_diagnostic()?;
        let hydrant = Hydrant::new(test_config(tmp.path())).await?;

        // Seed some in DB: host2, host4, host6
        {
            let state = hydrant.state.clone();
            state
                .db
                .run(move |db| -> Result<()> {
                    db.cursors
                        .insert(keys::firehose_cursor_key("host2"), 2_i64.to_be_bytes())
                        .into_diagnostic()?;
                    db.cursors
                        .insert(keys::firehose_cursor_key("host4"), 4_i64.to_be_bytes())
                        .into_diagnostic()?;
                    db.cursors
                        .insert(keys::firehose_cursor_key("host6"), 6_i64.to_be_bytes())
                        .into_diagnostic()?;
                    db.persist()
                })
                .await?;
        }

        // Seed some in memory: host1, host3, host5
        crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host1", |h| {
            h.status = HostStatus::Active;
        });
        crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host3", |h| {
            h.status = HostStatus::Active;
        });
        crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "host5", |h| {
            h.status = HostStatus::Active;
        });

        // Lexicographically sorted: host1, host2, host3, host4, host5, host6
        // Page 1: limit = 2
        let (hosts1, next1) = hydrant.list_hosts(None, 2).await?;
        assert_eq!(hosts1.len(), 2);
        assert_eq!(hosts1[0].name.as_str(), "host1");
        assert_eq!(hosts1[1].name.as_str(), "host2");
        assert_eq!(next1.as_deref(), Some("host2"));

        // Page 2: limit = 2, cursor = host2
        let (hosts2, next2) = hydrant.list_hosts(next1.as_deref(), 2).await?;
        assert_eq!(hosts2.len(), 2);
        assert_eq!(hosts2[0].name.as_str(), "host3");
        assert_eq!(hosts2[1].name.as_str(), "host4");
        assert_eq!(next2.as_deref(), Some("host4"));

        // Page 3: limit = 2, cursor = host4
        let (hosts3, next3) = hydrant.list_hosts(next2.as_deref(), 2).await?;
        assert_eq!(hosts3.len(), 2);
        assert_eq!(hosts3[0].name.as_str(), "host5");
        assert_eq!(hosts3[1].name.as_str(), "host6");
        assert!(next3.is_none());

        Ok(())
    }
}