//! High-performance shortest path queries.
//!
//! This module implements the bidirectional Dijkstra search used to find the shortest path
//! within the customized hierarchy.
//!
//! To achieve maximum throughput, the query state ([`Query`]) is stored in a thread-local
//! variable (`QUERY_BUFFER`). It uses an `epoch` counter to implicitly clear the visited
//! state between queries. This completely eliminates the need for dynamic memory allocation
//! and `Vec` clearing during individual routing requests.

use crate::cch::Cch;
use crate::cch::customize::Weights;

#[repr(C)]
#[derive(Clone, Copy)]
pub struct QuerySlot {
    pub dist: f32,
    pub epoch: u32,
    pub parent_node: u32,
    pub parent_edge: u32,
}

pub struct Query {
    pub fw: Vec<QuerySlot>,
    pub bw: Vec<QuerySlot>,
    pub epoch: u32,
}

impl Query {
    pub fn new(n: usize) -> Self {
        let s = QuerySlot {
            dist: f32::INFINITY,
            epoch: 0,
            parent_node: u32::MAX,
            parent_edge: u32::MAX,
        };
        Self {
            fw: vec![s; n],
            bw: vec![s; n],
            epoch: 0,
        }
    }

    pub fn next(&mut self) {
        self.epoch = self.epoch.wrapping_add(1);
        if self.epoch == 0 {
            self.fw.iter_mut().chain(&mut self.bw).for_each(|s| s.epoch = 0);
            self.epoch = 1;
        }
    }
}

pub struct QueryResult {
    pub distance: f32,
    pub from_rank: u32,
    pub to_rank: u32,
    pub meeting: u32,
}

impl Cch {
    #[inline(always)]
    unsafe fn relax(&self, u: u32, weights: &[f32], slots: *mut QuerySlot, epoch: u32) -> Option<u32> {
        let u_idx = u as usize;
        unsafe {
            let u_slot = &*slots.add(u_idx);
            if u_slot.epoch == epoch {
                let d_u = u_slot.dist;
                let start = *self.topology.first_out.get_unchecked(u_idx) as usize;
                let end = *self.topology.first_out.get_unchecked(u_idx + 1) as usize;

                for i in start..end {
                    let (v_idx, w) = (*self.topology.head.get_unchecked(i) as usize, *weights.get_unchecked(i));
                    let v_s = &mut *slots.add(v_idx);
                    let dist = d_u + w;

                    if v_s.epoch != epoch || dist < v_s.dist {
                        *v_s = QuerySlot {
                            dist,
                            epoch,
                            parent_node: u,
                            parent_edge: i as u32,
                        };
                    }
                }
            }
            *self.topology.etree.get_unchecked(u_idx)
        }
    }

    pub fn query(&self, q: &mut Query, w: &Weights, from: u32, to: u32) -> Option<QueryResult> {
        let (f_r, t_r) = (self.ranks[from as usize], self.ranks[to as usize]);
        q.next();
        let (e, fw_p, bw_p) = (q.epoch, q.fw.as_mut_ptr(), q.bw.as_mut_ptr());

        unsafe {
            let init = |p: *mut QuerySlot| {
                *p = QuerySlot {
                    dist: 0.0,
                    epoch: e,
                    parent_node: u32::MAX,
                    parent_edge: u32::MAX,
                }
            };
            init(fw_p.add(f_r as usize));
            init(bw_p.add(t_r as usize));

            let mut curr = Some(f_r);
            while let Some(c) = curr {
                curr = self.relax(c, &w.up, fw_p, e);
            }

            let (mut best_d, mut meeting, mut curr) = (f32::INFINITY, u32::MAX, Some(t_r));
            while let Some(c) = curr {
                let (fs, bs) = (&*fw_p.add(c as usize), &*bw_p.add(c as usize));
                if fs.epoch == e && bs.epoch == e && fs.dist + bs.dist < best_d {
                    best_d = fs.dist + bs.dist;
                    meeting = c;
                }
                curr = self.relax(c, &w.down, bw_p, e);
            }

            (best_d < f32::INFINITY).then_some(QueryResult {
                distance: best_d,
                from_rank: f_r,
                to_rank: t_r,
                meeting,
            })
        }
    }
}
