//! Multi-scalar multiplication over secp256k1 using k256 group arithmetic
//! (`ProjectivePoint` / `AffinePoint` / `Scalar`), instrumented with operation counters.
//!
//! Two algorithms are implemented here from scratch (not taken from k256):
//!   * `straus`    — Straus/Shamir interleaving with width-w signed-digit (wNAF) recoding,
//!                   one odd-multiple table per term, projective-projective additions.
//!   * `pippenger` — bucket method (Pippenger 1976), c-bit windows, bucket accumulation with
//!                   mixed (projective + affine) additions and the running-sum trick.
//! The k256 library routines `lincomb` (constant-time GLV + radix-16) and `lincomb_vartime`
//! (GLV + wNAF multiscalar) are also exposed for comparison; they are not instrumented.
//! Scalar recoding is done on `u64` limbs (exact integer arithmetic only).

use k256::elliptic_curve::group::{Group, GroupEncoding};
use k256::elliptic_curve::ops::{LinearCombination, Reduce};
use k256::elliptic_curve::ff::PrimeField;
use k256::{AffinePoint, CompressedPoint, FieldBytes, ProjectivePoint, Scalar};
use serde::Serialize;
use sha2::{Digest, Sha256};

#[derive(Serialize, Clone, Copy, Debug, Default)]
pub struct OpCounts {
    /// projective + projective additions (includes additions involving the identity)
    pub adds: u64,
    /// projective + affine (mixed) additions
    pub mixed_adds: u64,
    /// point doublings
    pub doublings: u64,
    /// additions where one operand was the identity (subset of adds + mixed_adds; cheap in math,
    /// full cost in k256's complete formulas)
    pub identity_adds: u64,
}

impl OpCounts {
    pub fn total(&self) -> u64 {
        self.adds + self.mixed_adds + self.doublings
    }
}

/// BIP340 `lift_x`: the point with the given x coordinate and even y, if any.
pub fn lift_x(x: &[u8; 32]) -> Option<AffinePoint> {
    let mut c = [0u8; 33];
    c[0] = 2;
    c[1..].copy_from_slice(x);
    let cp = CompressedPoint::from(c);
    AffinePoint::from_bytes(&cp).into_option()
}

pub fn tagged_hash_prefix(tag: &str) -> [u8; 32] {
    let th = Sha256::digest(tag.as_bytes());
    let mut out = [0u8; 32];
    out.copy_from_slice(&th);
    out
}

/// e = int(tagged_hash("BIP0340/challenge", r || pk || msg)) mod n
pub fn challenge(th: &[u8; 32], r: &[u8; 32], pk: &[u8; 32], msg: &[u8]) -> Scalar {
    let h = Sha256::new()
        .chain_update(th)
        .chain_update(th)
        .chain_update(r)
        .chain_update(pk)
        .chain_update(msg)
        .finalize();
    let mut b = [0u8; 32];
    b.copy_from_slice(&h);
    <Scalar as Reduce<FieldBytes>>::reduce(&FieldBytes::from(b))
}

pub fn scalar_from_bytes_strict(b: &[u8]) -> Option<Scalar> {
    let mut a = [0u8; 32];
    a.copy_from_slice(b);
    Scalar::from_repr(FieldBytes::from(a)).into_option()
}

/// Little-endian 64-bit limbs of a scalar (exact).
pub fn limbs(s: &Scalar) -> [u64; 4] {
    let b = s.to_bytes();
    let mut l = [0u64; 4];
    for i in 0..4 {
        l[3 - i] = u64::from_be_bytes(b[i * 8..i * 8 + 8].try_into().unwrap());
    }
    l
}

fn is_zero(x: &[u64; 5]) -> bool {
    x.iter().all(|&v| v == 0)
}

fn sub_small(x: &mut [u64; 5], d: u64) {
    let mut borrow = d;
    for limb in x.iter_mut() {
        let (v, b) = limb.overflowing_sub(borrow);
        *limb = v;
        borrow = b as u64;
        if borrow == 0 {
            break;
        }
    }
}

fn add_small(x: &mut [u64; 5], d: u64) {
    let mut carry = d;
    for limb in x.iter_mut() {
        let (v, c) = limb.overflowing_add(carry);
        *limb = v;
        carry = c as u64;
        if carry == 0 {
            break;
        }
    }
}

fn shr1(x: &mut [u64; 5]) {
    for i in 0..5 {
        let hi = if i + 1 < 5 { x[i + 1] << 63 } else { 0 };
        x[i] = (x[i] >> 1) | hi;
    }
}

/// Width-w non-adjacent form: digits d_i in {0, ±1, ±3, …, ±(2^{w-1}-1)}, least significant first.
pub fn wnaf(k: &[u64; 4], w: u32) -> Vec<i8> {
    debug_assert!((2..=7).contains(&w));
    let mut x = [k[0], k[1], k[2], k[3], 0u64];
    let mask = (1u64 << w) - 1;
    let half = 1i64 << (w - 1);
    let mut out = Vec::with_capacity(258);
    while !is_zero(&x) {
        if x[0] & 1 == 1 {
            let mut d = (x[0] & mask) as i64;
            if d >= half {
                d -= 1 << w;
            }
            if d > 0 {
                sub_small(&mut x, d as u64);
            } else {
                add_small(&mut x, (-d) as u64);
            }
            out.push(d as i8);
        } else {
            out.push(0);
        }
        shr1(&mut x);
    }
    out
}

/// Straus interleaved wNAF multi-scalar multiplication: Σ k_i P_i.
pub fn straus(terms: &[(ProjectivePoint, Scalar)], w: u32, cnt: &mut OpCounts) -> ProjectivePoint {
    let tsize = 1usize << (w - 2); // odd multiples 1,3,…,2^{w-1}-1
    let mut tables: Vec<Vec<ProjectivePoint>> = Vec::with_capacity(terms.len());
    let mut digits: Vec<Vec<i8>> = Vec::with_capacity(terms.len());
    for (p, k) in terms {
        let dbl = p.double();
        cnt.doublings += 1;
        let mut t = Vec::with_capacity(tsize);
        t.push(*p);
        for i in 1..tsize {
            t.push(t[i - 1] + dbl);
            cnt.adds += 1;
        }
        tables.push(t);
        digits.push(wnaf(&limbs(k), w));
    }
    let maxlen = digits.iter().map(|d| d.len()).max().unwrap_or(0);
    let mut acc = ProjectivePoint::IDENTITY;
    let mut started = false;
    for i in (0..maxlen).rev() {
        if started {
            acc = acc.double();
            cnt.doublings += 1;
        }
        for (j, d) in digits.iter().enumerate() {
            if i < d.len() && d[i] != 0 {
                let idx = (d[i].unsigned_abs() / 2) as usize;
                let pt = tables[j][idx];
                if !started {
                    cnt.identity_adds += 1;
                }
                acc += if d[i] > 0 { pt } else { -pt };
                cnt.adds += 1;
                started = true;
            }
        }
    }
    acc
}

fn window(l: &[u64; 4], shift: u32, c: u32) -> usize {
    let li = (shift / 64) as usize;
    let off = shift % 64;
    let mut v = l[li] >> off;
    if off + c > 64 && li + 1 < 4 {
        v |= l[li + 1] << (64 - off);
    }
    (v & ((1u64 << c) - 1)) as usize
}

/// Window width heuristic for Pippenger as a function of the number of terms.
pub fn pippenger_window(n_terms: usize) -> u32 {
    match n_terms {
        0..=7 => 3,
        8..=31 => 4,
        32..=127 => 5,
        128..=511 => 6,
        512..=2047 => 7,
        2048..=8191 => 8,
        _ => 9,
    }
}

/// Pippenger bucket method: Σ k_i P_i with affine inputs (mixed additions into buckets).
pub fn pippenger(terms: &[(AffinePoint, Scalar)], c: u32, cnt: &mut OpCounts) -> ProjectivePoint {
    let n_windows = (256 + c - 1) / c;
    let lim: Vec<[u64; 4]> = terms.iter().map(|(_, k)| limbs(k)).collect();
    let nb = (1usize << c) - 1;
    let mut buckets = vec![ProjectivePoint::IDENTITY; nb];
    let mut touched = vec![false; nb];
    let mut acc = ProjectivePoint::IDENTITY;
    let mut started = false;
    for wi in (0..n_windows).rev() {
        if started {
            for _ in 0..c {
                acc = acc.double();
                cnt.doublings += 1;
            }
        }
        for (b, t) in buckets.iter_mut().zip(touched.iter_mut()) {
            *b = ProjectivePoint::IDENTITY;
            *t = false;
        }
        let shift = wi * c;
        let mut any = false;
        for (i, (p, _)) in terms.iter().enumerate() {
            let d = window(&lim[i], shift, c);
            if d != 0 {
                if !touched[d - 1] {
                    cnt.identity_adds += 1;
                    touched[d - 1] = true;
                }
                buckets[d - 1] += p;
                cnt.mixed_adds += 1;
                any = true;
            }
        }
        if !any {
            continue;
        }
        // running-sum trick: Σ_j j·B_j = Σ_j (Σ_{i ≥ j} B_i)
        let mut running = ProjectivePoint::IDENTITY;
        let mut sum = ProjectivePoint::IDENTITY;
        let mut run_started = false;
        for (j, b) in buckets.iter().enumerate().rev() {
            if touched[j] {
                if !run_started {
                    cnt.identity_adds += 1;
                }
                running += b;
                cnt.adds += 1;
                run_started = true;
            }
            if run_started {
                sum += running;
                cnt.adds += 1;
            }
        }
        if started {
            acc += sum;
            cnt.adds += 1;
        } else {
            acc = sum;
            started = true;
        }
    }
    acc
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum MsmAlgo {
    Straus,
    Pippenger,
    K256LincombVartime,
    K256LincombCt,
}

impl MsmAlgo {
    pub fn name(&self) -> &'static str {
        match self {
            MsmAlgo::Straus => "straus_wnaf5",
            MsmAlgo::Pippenger => "pippenger",
            MsmAlgo::K256LincombVartime => "k256_lincomb_vartime",
            MsmAlgo::K256LincombCt => "k256_lincomb_ct",
        }
    }
}

/// One BIP340 signature to batch-verify (public key already lifted; `lift_x(R)` and the
/// challenge hash are computed inside the batch routine, as in a real verifier).
pub struct SigItem {
    pub pk_x: [u8; 32],
    pub p: AffinePoint,
    pub msg: [u8; 32],
    pub sig: [u8; 64],
}

pub struct BatchOutcome {
    pub valid: bool,
    pub counts: OpCounts,
    pub msm_ns: u128,
    pub n_terms: usize,
    pub window: u32,
}

/// Batch verification  Σ a_i s_i G − Σ a_i R_i − Σ a_i e_i P_i = O   with a_1 = 1,
/// a_i random 128-bit (BIP340 §"Batch Verification"; Bernstein et al. 2012 for Ed25519).
pub fn batch_verify(items: &[&SigItem], weights: &[[u8; 16]], algo: MsmAlgo, th: &[u8; 32], pip_c: u32) -> BatchOutcome {
    let n = items.len();
    let mut terms: Vec<(AffinePoint, Scalar)> = Vec::with_capacity(2 * n + 1);
    let mut s_g = Scalar::ZERO;
    let mut cnt = OpCounts::default();
    let fail = |n_terms| BatchOutcome { valid: false, counts: OpCounts::default(), msm_ns: 0, n_terms, window: 0 };
    for (i, it) in items.iter().enumerate() {
        let r: [u8; 32] = it.sig[..32].try_into().unwrap();
        let s = match scalar_from_bytes_strict(&it.sig[32..]) {
            Some(s) => s,
            None => return fail(0),
        };
        let rp = match lift_x(&r) {
            Some(p) => p,
            None => return fail(0),
        };
        let e = challenge(th, &r, &it.pk_x, &it.msg);
        let a = if i == 0 { Scalar::ONE } else { Scalar::from(u128::from_le_bytes(weights[i])) };
        s_g += a * s;
        // negate the points (free in affine) rather than the scalars, so a_i stays 128-bit
        terms.push((-rp, a));
        terms.push((-it.p, a * e));
    }
    terms.push((AffinePoint::GENERATOR, s_g));
    let n_terms = terms.len();
    let t0 = std::time::Instant::now();
    let (res, window) = match algo {
        MsmAlgo::Straus => {
            let proj: Vec<(ProjectivePoint, Scalar)> = terms.iter().map(|(p, k)| (ProjectivePoint::from(*p), *k)).collect();
            (straus(&proj, 5, &mut cnt), 5)
        }
        MsmAlgo::Pippenger => {
            let c = if pip_c == 0 { pippenger_window(n_terms) } else { pip_c };
            (pippenger(&terms, c, &mut cnt), c)
        }
        MsmAlgo::K256LincombVartime => {
            let proj: Vec<(ProjectivePoint, Scalar)> = terms.iter().map(|(p, k)| (ProjectivePoint::from(*p), *k)).collect();
            (<ProjectivePoint as LinearCombination<[(ProjectivePoint, Scalar)]>>::lincomb_vartime(proj.as_slice()), 0)
        }
        MsmAlgo::K256LincombCt => {
            let proj: Vec<(ProjectivePoint, Scalar)> = terms.iter().map(|(p, k)| (ProjectivePoint::from(*p), *k)).collect();
            (<ProjectivePoint as LinearCombination<[(ProjectivePoint, Scalar)]>>::lincomb(proj.as_slice()), 0)
        }
    };
    let msm_ns = t0.elapsed().as_nanos();
    BatchOutcome { valid: bool::from(res.is_identity()), counts: cnt, msm_ns, n_terms, window }
}