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

pub struct TtlAwareRouter {
    pub alpha: f64,
}

impl TtlAwareRouter {
    pub fn new(alpha: f64) -> Self {
        Self { alpha }
    }
}

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

    fn route(
        &mut self,
        req: &RequestRecord,
        instances: &[Instance],
        meta: &MetaStore,
        now: f64,
    ) -> RouteDecision {
        let n = instances.len();
        let scores = meta.score_prefix(&req.hash_ids, now, n);
        let mut best = 0u32;
        let mut best_key = (i64::MIN, f64::INFINITY); // maximize prefix, then minimize load
        let mut candidates = Vec::with_capacity(n);
        for inst in instances {
            let p = scores[inst.id as usize];
            let load = inst.kv_blocks_used as f64 + self.alpha * inst.queue_len() as f64;
            candidates.push(CandidateInfo {
                instance: inst.id,
                predicted_prefix: p,
                load_blocks: inst.kv_blocks_used,
                queue_len: inst.queue_len(),
            });
            let key = (p as i64, -load);
            // we want max prefix, min load -> compare (p, -load) lexicographically max
            if key > (best_key.0, -best_key.1) {
                best_key = (p as i64, load);
                best = inst.id;
            }
        }
        crate::router::local_route_decision(
            req.req_id,
            "ttl_aware",
            best,
            0.0,
            candidates,
            "max meta_store prefix, tie -> least loaded",
        )
    }
}