//! _grumbles_ ketama doesn't support uds _grumbles_
//!
//! (also this lets us add an option to choose in the future)

use std::sync::Arc;

use pingora::lb::selection::BackendSelection;
use pingora::lb::selection::consistent::{KetamaConfig, KetamaHashing};
use pingora::lb::{self, Backend};
use pingora::prelude::*;

pub enum SelectionChoice {
    Ketama(Arc<KetamaHashing>),
    RoundRobin(Arc<RoundRobin>),
}

#[allow(
    dead_code,
    reason = "required for trait impl, will use the in the future"
)]
pub enum SelectionConfig {
    Ketama(KetamaConfig),
    RoundRobin(()),
}
pub enum SelectionIter {
    Ketama(<KetamaHashing as BackendSelection>::Iter),
    RoundRobin(<RoundRobin as BackendSelection>::Iter),
}
impl lb::selection::BackendIter for SelectionIter {
    fn next(&mut self) -> Option<&Backend> {
        match self {
            SelectionIter::Ketama(k) => k.next(),
            SelectionIter::RoundRobin(rr) => rr.next(),
        }
    }
}
impl BackendSelection for SelectionChoice {
    type Iter = SelectionIter;

    type Config = SelectionConfig;

    fn build(backends: &std::collections::BTreeSet<lb::Backend>) -> Self {
        // need the arc here for iter
        if backends
            .iter()
            .any(|b| matches!(b.addr, pingora::protocols::l4::socket::SocketAddr::Unix(_)))
        {
            // ketama does not support unix domain sockets
            Self::RoundRobin(Arc::new(RoundRobin::build(backends)))
        } else {
            Self::Ketama(Arc::new(KetamaHashing::build(backends)))
        }
    }

    fn iter(self: &std::sync::Arc<Self>, key: &[u8]) -> Self::Iter
    where
        Self::Iter: lb::selection::BackendIter,
    {
        match &**self {
            SelectionChoice::Ketama(k) => SelectionIter::Ketama(k.iter(key)),
            SelectionChoice::RoundRobin(rr) => SelectionIter::RoundRobin(rr.iter(key)),
        }
    }

    fn build_with_config(
        backends: &std::collections::BTreeSet<Backend>,
        config: &Self::Config,
    ) -> Self {
        // need the arc here for iter
        match config {
            SelectionConfig::Ketama(ketama) => {
                Self::Ketama(Arc::new(KetamaHashing::build_with_config(backends, ketama)))
            }
            SelectionConfig::RoundRobin(rr) => {
                Self::RoundRobin(Arc::new(RoundRobin::build_with_config(backends, rr)))
            }
        }
    }
}