use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::cluster::meta_store::MetaStore;
use crate::instance::Instance;
use crate::router::{CandidateInfo, RouteDecision, Router};
use crate::trace::RequestRecord;
use crate::types::InstanceId;

pub struct RandomRouter {
    rng: ChaCha8Rng,
}

impl RandomRouter {
    pub fn new(seed: u64) -> Self {
        Self {
            rng: ChaCha8Rng::seed_from_u64(seed),
        }
    }
}

impl Router for RandomRouter {
    fn name(&self) -> &'static str {
        "random"
    }

    fn route(
        &mut self,
        req: &RequestRecord,
        instances: &[Instance],
        _meta: &MetaStore,
        _now: f64,
    ) -> RouteDecision {
        let n = instances.len();
        let chosen = self.rng.gen_range(0..n) as InstanceId;
        crate::router::local_route_decision(
            req.req_id,
            "random",
            chosen,
            0.0,
            vec![CandidateInfo {
                instance: chosen,
                predicted_prefix: 0,
                load_blocks: instances[chosen as usize].kv_blocks_used,
                queue_len: instances[chosen as usize].queue_len(),
            }],
            "uniform random",
        )
    }
}

#[derive(Default)]
pub struct RoundRobinRouter {
    next: u32,
}

impl RoundRobinRouter {
    pub fn new() -> Self {
        Self::default()
    }
}

impl Router for RoundRobinRouter {
    fn name(&self) -> &'static str {
        "round_robin"
    }

    fn route(
        &mut self,
        req: &RequestRecord,
        instances: &[Instance],
        _meta: &MetaStore,
        _now: f64,
    ) -> RouteDecision {
        let n = instances.len() as u32;
        let chosen = self.next % n;
        self.next = self.next.wrapping_add(1);
        crate::router::local_route_decision(
            req.req_id,
            "round_robin",
            chosen,
            0.0,
            vec![CandidateInfo {
                instance: chosen,
                predicted_prefix: 0,
                load_blocks: instances[chosen as usize].kv_blocks_used,
                queue_len: instances[chosen as usize].queue_len(),
            }],
            "round robin",
        )
    }
}