// NUMA-aware partition runner via hwlocality.
//
// Detects NUMA topology using hwloc (cross-platform: Linux, macOS, etc.) and
// builds one Rayon ThreadPool per NUMA node with threads pinned to that node's
// CPUs. Linux first-touch policy then places graph allocations in local DRAM
// automatically — no explicit memory binding needed.
//
// UMA systems (single socket, Apple Silicon, etc.) are the degenerate case:
// one synthetic node containing all cores, no pool, no pinning.

use std::sync::Arc;
use std::time::{Duration, Instant};

use crossbeam_channel::unbounded;
#[cfg(feature = "numa")]
use hwlocality::Topology;
#[cfg(feature = "numa")]
use hwlocality::cpu::binding::CpuBindingFlags;
#[cfg(feature = "numa")]
use hwlocality::cpu::cpuset::CpuSet;
#[cfg(feature = "numa")]
use hwlocality::object::types::ObjectType;
use obisys::{CpuSample, IoSample};
use tracing::debug;

// ── Public interface ──────────────────────────────────────────────────────────

pub struct NumaSetup {
    /// One entry per NUMA node. `None` on UMA systems (no pool, no pinning).
    pub pools: Vec<Option<Arc<rayon::ThreadPool>>>,
    /// CPU indices for each NUMA node, in node order.
    pub cpus_per_node: Vec<Vec<usize>>,
}

impl NumaSetup {
    /// Maximum worker slots per node (one per physical core in the node).
    pub fn workers_per_node(&self) -> usize {
        self.cpus_per_node
            .first()
            .map(|c| c.len().max(1))
            .unwrap_or(1)
    }
}

/// Detect NUMA topology and build per-node Rayon pools.
/// Always succeeds: falls back to a single synthetic UMA node on failure.
#[cfg(feature = "numa")]
pub fn build() -> NumaSetup {
    if let Ok(topology) = Topology::new() {
        let nodes: Vec<Vec<usize>> = topology
            .objects_with_type(ObjectType::NUMANode)
            .filter_map(|obj| obj.cpuset())
            .map(|cpuset| {
                cpuset
                    .iter_set()
                    .map(|idx| usize::from(idx))
                    .collect::<Vec<_>>()
            })
            .filter(|v| !v.is_empty())
            .collect();

        if nodes.len() > 1 {
            if let Some(pools) = nodes
                .iter()
                .map(|cpus| build_pool(cpus).map(|p| Some(Arc::new(p))))
                .collect::<Option<Vec<_>>>()
            {
                debug!(
                    "NUMA topology: {} node(s), {} core(s)/node",
                    nodes.len(),
                    nodes.first().map_or(0, |v| v.len()),
                );
                return NumaSetup {
                    pools,
                    cpus_per_node: nodes,
                };
            }
        }
    }

    // UMA fallback: single synthetic node, all cores, no pool, no pinning.
    let n_cores = obisys::effective_parallelism();
    debug!("UMA: single synthetic node, {} core(s)", n_cores);
    NumaSetup {
        pools: vec![None],
        cpus_per_node: vec![(0..n_cores).collect()],
    }
}

#[cfg(not(feature = "numa"))]
pub fn build() -> NumaSetup {
    let n_cores = obisys::effective_parallelism();
    debug!("UMA: single synthetic node, {} core(s)", n_cores);
    NumaSetup {
        pools: vec![None],
        cpus_per_node: vec![(0..n_cores).collect()],
    }
}

/// Bind the calling thread to `cpu_indices` using hwloc.
/// Silently returns on any error so the thread still runs, just unbound.
#[cfg(feature = "numa")]
pub fn pin_current_thread(cpu_indices: &[usize]) {
    let Ok(topology) = Topology::new() else {
        return;
    };
    let mut cpuset = CpuSet::new();
    for &idx in cpu_indices {
        cpuset.set(idx);
    }
    let _ = topology.bind_cpu(&cpuset, CpuBindingFlags::THREAD);
}

#[cfg(not(feature = "numa"))]
pub fn pin_current_thread(_cpu_indices: &[usize]) {}

// ── Internal helpers ──────────────────────────────────────────────────────────

#[cfg(feature = "numa")]
fn build_pool(cpus: &[usize]) -> Option<rayon::ThreadPool> {
    let cpus = cpus.to_vec();
    rayon::ThreadPoolBuilder::new()
        .num_threads(cpus.len())
        .spawn_handler(move |thread| {
            let cpus = cpus.clone();
            std::thread::Builder::new().spawn(move || {
                pin_current_thread(&cpus);
                thread.run();
            })?;
            Ok(())
        })
        .build()
        .ok()
}

// ── PartitionRunner ─────────────────────────────────────────────────────────

/// Growth step (fraction of a node's worker capacity added per activation
/// event, see [`NodeActivation::grow`]).
const GROWTH_DIVISOR: usize = 8;
/// Minimum CPU efficiency growth to activate more workers, as a fraction of
/// the size of the *last growth step* (e.g. `0.2` after adding 8 workers
/// requires the next check to show at least +1.6 cores of growth — 20 % of
/// the ~8 cores those 8 workers should contribute if the workload is truly
/// CPU-bound). Scaling by the last step's size — not the cumulative total —
/// keeps the bar meaningful regardless of how many workers are already
/// active, instead of demanding an ever-larger absolute jump as the pool
/// grows.
const CPU_SPAWN_THRESHOLD: f64 = 0.2;
/// Minimum I/O throughput growth (relative) to activate more workers.
const IO_SPAWN_THRESHOLD: f64 = 0.2;

struct NodeConfig {
    pool: Option<Arc<rayon::ThreadPool>>,
    cpu_ids: Vec<usize>,
    max_workers: usize,
}

/// Generic NUMA-aware runner for partition-level parallel work.
///
/// Workers are distributed evenly across NUMA nodes and pinned to their
/// node's CPUs. UMA is the degenerate case: one node, no pinning.
///
/// Workers are pre-spawned dormant, one activation channel per node so
/// growth always targets a specific node rather than whichever dormant
/// worker happens to wake up first on a shared channel. Growth (both the
/// initial count and each subsequent step) is expressed as a fraction of
/// `workers_per_node`, applied identically to every node, so the pace of
/// ramp-up depends on node size rather than node count — a single-NUMA-node
/// (UMA) machine ramps just as fast as an 8-node one.
///
/// # Termination
///
/// ```text
/// drop(part_tx)      → part_rx drains → workers exit → drop their result_tx
/// drop(result_tx)    → result_rx closes → controller loop exits
/// drop(activate_txs) → dormant workers exit cleanly
/// ```
pub struct PartitionRunner {
    nodes: Vec<NodeConfig>,
}

impl PartitionRunner {
    /// Total worker slots across all nodes.
    pub fn max_workers(&self) -> usize {
        self.nodes.iter().map(|n| n.max_workers).sum()
    }

    /// Detect topology and build. Always succeeds.
    pub fn new() -> Self {
        let ns = build();
        let wpn = ns.workers_per_node();
        debug!(
            "PartitionRunner: {} node(s) × {} worker(s)/node max",
            ns.pools.len(),
            wpn,
        );
        let nodes = ns
            .pools
            .into_iter()
            .zip(ns.cpus_per_node)
            .map(|(pool, cpu_ids)| NodeConfig {
                pool,
                cpu_ids,
                max_workers: wpn,
            })
            .collect();
        Self { nodes }
    }

    /// Run `f(i)` for every index in `order`.
    ///
    /// Workers are pre-spawned dormant and activated adaptively, per node:
    /// `(workers_per_node / INITIAL_DIVISOR).max(1)` are woken immediately on
    /// every node, then `(workers_per_node / GROWTH_DIVISOR).max(1)` more per
    /// node each time the check below fires. A timer thread fires that check
    /// every `TIMER_SECS` seconds; each completed partition resets that timer
    /// (forcing an immediate check) and also triggers its own inline check. A
    /// growth step happens whenever CPU efficiency grows by at least
    /// `CPU_SPAWN_THRESHOLD` of what the last growth step should have
    /// contributed, or I/O throughput grows by at least `IO_SPAWN_THRESHOLD`
    /// (relative) since the last check — whichever resource is the actual
    /// bottleneck still shows headroom.
    ///
    /// `on_done(i, result, elapsed)` is called from the controller thread as
    /// each partition completes — suitable for progress bars and result
    /// aggregation.
    ///
    /// Returns the first error produced by `f`, if any.
    pub fn run<F, R, E, C>(&self, order: &[usize], f: F, mut on_done: C) -> Result<(), E>
    where
        F: Fn(usize) -> Result<R, E> + Send + Sync,
        R: Send,
        E: Send,
        C: FnMut(usize, R, Duration) + Send,
    {
        let n_total = order.len();
        if n_total == 0 {
            return Ok(());
        }

        const TIMER_SECS: u64 = 30;
        const INITIAL_DIVISOR: usize = 4;

        // ── Channels ──────────────────────────────────────────────────────────
        let (part_tx, part_rx) = unbounded::<usize>();
        // reset_tx: controller → timer ("reset the 30 s window")
        let (reset_tx, reset_rx) = unbounded::<()>();
        // event_tx: workers + timer → controller (unified event stream)
        let (event_tx, event_rx) = unbounded::<WorkerEvent<R, E>>();
        // One activation channel per node: growth always targets a specific
        // node, rather than whichever dormant worker happens to win the race
        // on a channel shared across all nodes.
        let (activate_txs, activate_rxs): (Vec<_>, Vec<_>) =
            (0..self.nodes.len()).map(|_| unbounded::<()>()).unzip();

        for &i in order {
            part_tx.send(i).ok();
        }
        drop(part_tx);

        let max_workers = self.max_workers();
        let node_caps: Vec<usize> = self.nodes.iter().map(|n| n.max_workers).collect();
        let f = &f;

        let mut first_err: Option<E> = None;

        std::thread::scope(|s| {
            // ── Timer thread ──────────────────────────────────────────────────
            // Sends TimerTick every TIMER_SECS seconds. Resets its window each
            // time reset_rx receives a message (i.e. on partition completion).
            let timer_tx = event_tx.clone();
            s.spawn(move || {
                let period = Duration::from_secs(TIMER_SECS);
                loop {
                    crossbeam_channel::select! {
                        recv(reset_rx) -> r => {
                            if r.is_err() { break; } // reset_tx dropped → exit
                        }
                        default(period) => {
                            if timer_tx.send(WorkerEvent::TimerTick).is_err() { break; }
                        }
                    }
                }
            });

            // ── Pre-spawn workers dormant, grouped by node ────────────────────
            // Each worker listens on its own node's activation channel only.
            for (node, arx) in self.nodes.iter().zip(activate_rxs.iter()) {
                let cpu_ids = &node.cpu_ids;
                for _ in 0..node.max_workers {
                    let prx = part_rx.clone();
                    let etx = event_tx.clone();
                    let arx = arx.clone();
                    let pool = node.pool.clone();

                    s.spawn(move || {
                        let tid = std::thread::current().id();
                        debug!(?tid, "PartitionRunner worker: waiting on activation");
                        if arx.recv().is_err() {
                            debug!(?tid, "PartitionRunner worker: activation channel closed, exiting");
                            return;
                        }
                        debug!(?tid, "PartitionRunner worker: activated");
                        if !cpu_ids.is_empty() {
                            pin_current_thread(cpu_ids);
                        }
                        for i in &prx {
                            debug!(?tid, partition = i, "PartitionRunner worker: picked partition");
                            let t = Instant::now();
                            let r = match &pool {
                                Some(p) => p.install(|| f(i)),
                                None => f(i),
                            };
                            debug!(?tid, partition = i, "PartitionRunner worker: partition done");
                            etx.send(WorkerEvent::Completed(i, r, t.elapsed())).ok();
                        }
                        debug!(?tid, "PartitionRunner worker: no more partitions, exiting");
                    });
                }
            }
            // Drop controller's event_tx: event_rx closes when all workers +
            // timer have exited.
            drop(event_tx);

            // ── Controller ────────────────────────────────────────────────────
            let mut activation = NodeActivation::new(&activate_txs, &node_caps, max_workers);
            activation.activate_initial(INITIAL_DIVISOR, n_total);
            debug!(n_total, activated = activation.total(), "PartitionRunner controller: initial activation");

            let mut cpu_sample = CpuSample::now();
            let mut io_sample = IoSample::now();
            let mut completed = 0usize;

            while completed < n_total {
                debug!(completed, n_total, "PartitionRunner controller: waiting for an event");
                let Ok(event) = event_rx.recv() else {
                    debug!("PartitionRunner controller: event channel closed, stopping");
                    break;
                };
                match event {
                    WorkerEvent::Completed(i, r, dur) => {
                        match r {
                            Ok(v) => on_done(i, v, dur),
                            Err(e) => {
                                if first_err.is_none() {
                                    first_err = Some(e);
                                }
                            }
                        }
                        completed += 1;
                        // Reset the 30 s timer.
                        reset_tx.send(()).ok();
                        // Inline check: same logic as a timer tick.
                        maybe_activate(
                            &mut activation,
                            &mut cpu_sample,
                            &mut io_sample,
                            completed,
                            n_total,
                        );
                    }
                    WorkerEvent::TimerTick => {
                        maybe_activate(
                            &mut activation,
                            &mut cpu_sample,
                            &mut io_sample,
                            completed,
                            n_total,
                        );
                    }
                }
            }

            // Dormant workers exit once every sender for their node's channel
            // is dropped — `activate_txs` holds the only ones.
            drop(activate_txs);
            // Timer thread exits when reset_tx closes.
            drop(reset_tx);
        });

        match first_err {
            Some(e) => Err(e),
            None => Ok(()),
        }
    }
}

// ── Internal event type ───────────────────────────────────────────────────────

enum WorkerEvent<R, E> {
    Completed(usize, Result<R, E>, Duration),
    TimerTick,
}

/// Tracks how many of each node's dormant workers have been woken, and
/// grows every node by the same amount at each step (capped by that node's
/// remaining dormant workers and by the run's total budget) so load stays
/// balanced across nodes at every point in time — never just "one more
/// worker somewhere". Also remembers the size of the last real growth step
/// (`last_step`), used to scale the CPU activation threshold to what that
/// step could plausibly have contributed (see `maybe_activate`).
struct NodeActivation<'a> {
    txs: &'a [crossbeam_channel::Sender<()>],
    caps: &'a [usize],
    active: Vec<usize>,
    total: usize,
    max: usize,
    last_step: usize,
}

impl<'a> NodeActivation<'a> {
    fn new(txs: &'a [crossbeam_channel::Sender<()>], caps: &'a [usize], max: usize) -> Self {
        Self {
            txs,
            caps,
            active: vec![0; txs.len()],
            total: 0,
            max,
            last_step: 0,
        }
    }

    fn total(&self) -> usize {
        self.total
    }
    fn last_step(&self) -> usize {
        self.last_step
    }
    fn max(&self) -> usize {
        self.max
    }
    fn is_full(&self) -> bool {
        self.total >= self.max
    }

    /// Wake up to `(node_cap / divisor).max(1)` dormant workers on every
    /// node, capped by `n_total`. Called once at startup, unconditionally.
    fn activate_initial(&mut self, divisor: usize, n_total: usize) {
        self.grow(divisor, n_total);
    }

    /// Same per-node sizing as [`activate_initial`](Self::activate_initial),
    /// applied as a growth step. Returns the number of workers actually
    /// activated (may be less than requested once a node or the total
    /// budget is exhausted). Updates `last_step` when it actually grew.
    fn grow(&mut self, divisor: usize, n_total: usize) -> usize {
        let before = self.total;
        for idx in 0..self.txs.len() {
            let wanted = (self.caps[idx] / divisor).max(1);
            let room = self.caps[idx].saturating_sub(self.active[idx]);
            let grow = wanted.min(room).min(n_total.saturating_sub(self.total));
            for _ in 0..grow {
                self.txs[idx].send(()).ok();
            }
            self.active[idx] += grow;
            self.total += grow;
        }
        let grew = self.total - before;
        if grew > 0 {
            self.last_step = grew;
        }
        grew
    }
}

fn maybe_activate(
    activation: &mut NodeActivation,
    cpu_sample: &mut CpuSample,
    io_sample: &mut IoSample,
    completed: usize,
    n_total: usize,
) {
    if activation.is_full() || completed >= n_total {
        return;
    }

    // Expect roughly 1 core of extra efficiency per worker activated in the
    // last growth step (CPU-bound case); require at least CPU_SPAWN_THRESHOLD
    // (20 %) of that expected gain before growing again. Scaling by the last
    // step's size — not the cumulative total — keeps the bar meaningful
    // regardless of how many workers are already active: growing by 8 should
    // always take ~+1.6 cores to confirm, whether that's the 2nd growth step
    // or the 20th.
    let cpu_threshold = CPU_SPAWN_THRESHOLD * activation.last_step() as f64;

    // Call both unconditionally (no `||` short-circuit): each sampler must
    // advance its own window every tick, regardless of what the other one
    // reports, or it would starve behind whichever signal fires first.
    let cpu_wants_more = cpu_sample.do_i_activate(cpu_threshold);
    let io_wants_more = io_sample.do_i_activate(IO_SPAWN_THRESHOLD * activation.last_step() as f64);
    if !(cpu_wants_more || io_wants_more) {
        return;
    }

    let grew = activation.grow(GROWTH_DIVISOR, n_total);
    if grew > 0 {
        debug!(
            "activated {} worker(s) — {}/{} active",
            grew,
            activation.total(),
            activation.max()
        );
    }
}
