//! Baby-step giant-step (Shanks 1971) with x-coordinate keys and ± lookup.
//!
//! Baby table: x(jG) → j for 0 < j < m (m = ⌈√n⌉ or --memory). Giant steps R_i = Q − i·(mG),
//! i = 0, 1, …, ⌈n/m⌉; a hit x(R_i) = x(jG) means R_i = ±jG, i.e. d = i·m ± j (both verified).
//! Both phases are chains advanced by a constant point, so m/S independent chains are advanced
//! together with Montgomery's simultaneous inversion (S = `chains`); the giant phase is split
//! across threads.

use crate::curve::{Counters, Pt};
use crate::instance::{isqrt_ceil, Instance};
use crate::stats::FxBuild;
use hashbrown::HashMap;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use std::time::Instant;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct BsgsConfig {
    /// number of baby steps m (default ⌈√n⌉)
    pub memory: Option<u64>,
    pub threads: usize,
    /// independent chains per thread for batch inversion
    pub chains: usize,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct BsgsResult {
    pub d: Option<u64>,
    pub ops: Counters,
    /// giant steps actually computed (all threads)
    pub iterations: u64,
    pub baby_steps: u64,
    pub giant_steps_planned: u64,
    pub table_entries: u64,
    pub table_bytes: u64,
    pub threads: usize,
    pub wall_s: f64,
    pub cpu_s: f64,
    pub build_s: f64,
}

/// Advance every chain by the constant point `s` (P_c ← P_c + s) with one batched inversion.
/// Chains at infinity or with x(P_c) = x(s) take the slow path.
fn advance_chains(inst: &Instance, pts: &mut [Pt], s: &Pt, dens: &mut [u64], scratch: &mut [u64], special: &mut [bool], c: &mut Counters) {
    let f = &inst.curve.f;
    let k = pts.len();
    for i in 0..k {
        let p = &pts[i];
        let sp = p.inf || p.x == s.x;
        special[i] = sp;
        dens[i] = if sp { f.one } else { f.sub(s.x, p.x) };
    }
    let (m, inv) = f.batch_inv(&mut dens[..k], &mut scratch[..k]);
    c.mul += m;
    c.inv += inv;
    for i in 0..k {
        if special[i] {
            pts[i] = inst.curve.add(&pts[i], s, c);
        } else {
            let lam = f.mul(f.sub(s.y, pts[i].y), dens[i]);
            pts[i] = inst.curve.finish_add(&pts[i], s.x, lam);
            c.add += 1;
            c.mul += 3;
        }
    }
}

pub fn solve_bsgs(inst: &Instance, cfg: &BsgsConfig) -> BsgsResult {
    let t0 = Instant::now();
    let cpu0 = crate::stats::thread_cpu_s();
    let n = inst.n.n;
    let mut c = Counters::default();
    let m = cfg.memory.unwrap_or(isqrt_ceil(n as u128) as u64).max(1).min(n);
    let t_total = (n as u128).div_ceil(m as u128) as u64 + 1;
    let threads = cfg.threads.max(1);
    let empty = |d: Option<u64>, c: Counters, t0: Instant, cpu0: f64| BsgsResult {
        d,
        ops: c,
        iterations: 0,
        baby_steps: m,
        giant_steps_planned: t_total,
        table_entries: 0,
        table_bytes: 0,
        threads,
        wall_s: t0.elapsed().as_secs_f64(),
        cpu_s: crate::stats::thread_cpu_s() - cpu0,
        build_s: 0.0,
    };
    if inst.q.inf {
        return empty(Some(0), c, t0, cpu0);
    }

    // ---- baby steps --------------------------------------------------------------------------
    let s = cfg.chains.max(1).min(m as usize);
    let len = (m as usize).div_ceil(s);
    let mut table: HashMap<u64, u64, FxBuild> = HashMap::with_capacity_and_hasher(m as usize + 1, Default::default());
    let mut pts: Vec<Pt> = (0..s).map(|ci| inst.curve.mul((ci * len) as u64, &inst.g, &mut c)).collect();
    let mut dens = vec![0u64; s];
    let mut scratch = vec![0u64; s];
    let mut special = vec![false; s];
    for step in 0..len {
        let mut active = 0;
        for ci in 0..s {
            let j = (ci * len + step) as u64;
            if j >= m {
                break;
            }
            active = ci + 1;
            if !pts[ci].inf {
                table.entry(pts[ci].x).or_insert(j);
            }
        }
        if active == 0 {
            break;
        }
        if step + 1 < len {
            advance_chains(inst, &mut pts[..active], &inst.g, &mut dens, &mut scratch, &mut special, &mut c);
        }
    }
    let build_s = t0.elapsed().as_secs_f64();
    let table_entries = table.len() as u64;
    let table_bytes = table.capacity() as u64 * (std::mem::size_of::<(u64, u64)>() as u64 + 1);

    // ---- giant steps (threaded) ---------------------------------------------------------------
    let mg = inst.curve.mul(m, &inst.g, &mut c);
    let neg_mg = inst.curve.neg(&mg);
    let found: Mutex<Option<u64>> = Mutex::new(None);
    let stop = AtomicBool::new(false);
    let per: Mutex<Vec<(Counters, u64, f64)>> = Mutex::new(Vec::new());
    let table = &table;
    let block = (t_total as usize).div_ceil(threads);
    std::thread::scope(|sc| {
        for t in 0..threads {
            let found = &found;
            let stop = &stop;
            let per = &per;
            let i_begin = (t * block) as u64;
            let i_end = (((t + 1) * block) as u64).min(t_total);
            if i_begin >= i_end {
                continue;
            }
            sc.spawn(move || {
                let cpu0 = crate::stats::thread_cpu_s();
                let mut c = Counters::default();
                let count = (i_end - i_begin) as usize;
                let s = cfg.chains.max(1).min(count);
                let len = count.div_ceil(s);
                // chain ci starts at R_{i_begin + ci·len} = Q − (i_begin + ci·len)·mG
                let mut pts: Vec<Pt> = (0..s)
                    .map(|ci| {
                        let i0 = i_begin + (ci * len) as u64;
                        let off = inst.curve.mul(i0 % n, &neg_mg, &mut c);
                        inst.curve.add(&inst.q, &off, &mut c)
                    })
                    .collect();
                let mut dens = vec![0u64; s];
                let mut scratch = vec![0u64; s];
                let mut special = vec![false; s];
                let mut iters = 0u64;
                let mut solved: Option<u64> = None;
                'outer: for step in 0..len {
                    if stop.load(Ordering::Relaxed) {
                        break;
                    }
                    let mut active = 0;
                    for ci in 0..s {
                        let i = i_begin + (ci * len + step) as u64;
                        if i >= i_end {
                            break;
                        }
                        active = ci + 1;
                        iters += 1;
                        let r = pts[ci];
                        if r.inf {
                            // Q = i·m·G
                            let d = inst.n.mul(i % n, m % n);
                            if inst.verify(d, &mut c) {
                                solved = Some(d);
                                break 'outer;
                            }
                            continue;
                        }
                        if let Some(&j) = table.get(&r.x) {
                            let im = inst.n.mul(i % n, m % n);
                            for d in [inst.n.add(im, j % n), inst.n.sub(im, j % n)] {
                                if inst.verify(d, &mut c) {
                                    solved = Some(d);
                                    break 'outer;
                                }
                            }
                        }
                    }
                    if active == 0 {
                        break;
                    }
                    if step + 1 < len {
                        advance_chains(inst, &mut pts[..active], &neg_mg, &mut dens, &mut scratch, &mut special, &mut c);
                    }
                }
                if let Some(d) = solved {
                    let mut f = found.lock().unwrap();
                    if f.is_none() {
                        *f = Some(d);
                    }
                    stop.store(true, Ordering::Relaxed);
                }
                per.lock().unwrap().push((c, iters, crate::stats::thread_cpu_s() - cpu0));
            });
        }
    });
    let mut iterations = 0u64;
    let mut cpu_s = crate::stats::thread_cpu_s() - cpu0;
    for (pc, it, cpu) in per.lock().unwrap().iter() {
        c.merge(pc);
        iterations += it;
        cpu_s += cpu;
    }
    let d = *found.lock().unwrap();
    BsgsResult {
        d,
        ops: c,
        iterations,
        baby_steps: m,
        giant_steps_planned: t_total,
        table_entries,
        table_bytes,
        threads,
        wall_s: t0.elapsed().as_secs_f64(),
        cpu_s,
        build_s,
    }
}