// 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;
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 = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);
    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 = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);
    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 ───────────────────────────────────────────────────────────

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 round-robin across NUMA nodes and pinned to their
/// node's CPUs. UMA is the degenerate case: one node, no pinning.
///
/// Workers are pre-spawned dormant and activated one by one as CPU efficiency
/// falls below `SPAWN_THRESHOLD`. This avoids over-provisioning on I/O-bound
/// or memory-bandwidth-bound workloads while saturating CPU-bound ones.
///
/// # 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_tx) → 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. A timer thread
    /// fires a CPU-efficiency check every `TIMER_SECS` seconds; each completed
    /// partition resets that timer (forcing an immediate check) and also
    /// triggers its own inline check. A new worker is activated whenever
    /// efficiency falls below `SPAWN_THRESHOLD`.
    ///
    /// `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 SPAWN_THRESHOLD: f64 = 0.95;
        const TIMER_SECS:      u64 = 30;

        let n_cores = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1);

        // ── Channels ──────────────────────────────────────────────────────────
        let (part_tx,     part_rx)     = unbounded::<usize>();
        let (activate_tx, activate_rx) = unbounded::<()>();
        // 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>>();

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

        let max_workers = self.max_workers();
        let n_nodes     = self.nodes.len();
        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, round-robin across NUMA nodes ──────
            for w in 0..max_workers {
                let node    = &self.nodes[w % n_nodes];
                let prx     = part_rx.clone();
                let etx     = event_tx.clone();
                let arx     = activate_rx.clone();
                let pool    = node.pool.clone();
                let cpu_ids = &node.cpu_ids;

                s.spawn(move || {
                    if arx.recv().is_err() { return; }
                    if !cpu_ids.is_empty() { pin_current_thread(cpu_ids); }
                    for i in &prx {
                        let t = Instant::now();
                        let r = match &pool {
                            Some(p) => p.install(|| f(i)),
                            None    => f(i),
                        };
                        etx.send(WorkerEvent::Completed(i, r, t.elapsed())).ok();
                    }
                });
            }
            // Drop controller's event_tx: event_rx closes when all workers +
            // timer have exited.
            drop(event_tx);

            // ── Controller ────────────────────────────────────────────────────
            activate_tx.send(()).ok();
            let mut n_active              = 1usize;
            let mut cpu_sample            = CpuSample::now();
            let mut eff_at_last_spawn     = 0.0f64; // 0 = no previous spawn to evaluate
            let mut completed             = 0usize;

            while completed < n_total {
                let Ok(event) = event_rx.recv() else { 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(
                            &activate_tx, &mut n_active, max_workers,
                            &mut cpu_sample, &mut eff_at_last_spawn,
                            n_cores, SPAWN_THRESHOLD, completed, n_total,
                        );
                    }
                    WorkerEvent::TimerTick => {
                        maybe_activate(
                            &activate_tx, &mut n_active, max_workers,
                            &mut cpu_sample, &mut eff_at_last_spawn,
                            n_cores, SPAWN_THRESHOLD, completed, n_total,
                        );
                    }
                }
            }

            // Dormant workers exit when activate_tx closes.
            drop(activate_tx);
            // 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,
}

fn maybe_activate(
    activate_tx:         &crossbeam_channel::Sender<()>,
    n_active:            &mut usize,
    max_workers:         usize,
    cpu_sample:          &mut CpuSample,
    eff_at_last_spawn:   &mut f64,
    n_cores:             usize,
    threshold:           f64,
    completed:           usize,
    n_total:             usize,
) {
    if *n_active >= max_workers || completed >= n_total { return; }

    let eff = cpu_sample.cpu_efficiency(n_cores);
    if eff >= threshold { return; }  // CPU already saturated

    // Check that the previous activation was beneficial enough.
    // Going from k-1 → k workers, the minimum acceptable speedup is (k-1+0.2)/(k-1).
    // For the very first extra worker (n_active == 1, no previous spawn), skip this
    // check: eff_at_last_spawn == 0 acts as the sentinel.
    let last_spawn_was_beneficial = if *eff_at_last_spawn < 1e-9 {
        true  // first additional worker: no prior data to evaluate
    } else {
        let k_before   = (*n_active - 1) as f64;
        let min_speedup = (k_before + 0.2) / k_before;
        let actual_speedup = eff / *eff_at_last_spawn;
        actual_speedup >= min_speedup
    };

    if last_spawn_was_beneficial {
        activate_tx.send(()).ok();
        *eff_at_last_spawn = eff;
        *n_active += 1;
        *cpu_sample = CpuSample::now();
        debug!(
            "activated worker {}/{} — efficiency {:.0}%",
            n_active, max_workers, eff * 100.0,
        );
    }
}
