//! Short Weierstrass curve y² = x³ + ax + b over F_p (p < 2^62), affine points in Montgomery form.
//!
//! Conventions match `src/wp34/ec.py`: affine coordinates, `None`/`Pt::INF` is the point at infinity,
//! the group law counts one group addition (`add`) or doubling (`dbl`) plus one field inversion (`inv`)
//! and the field multiplications (`mul`, squarings included) actually performed.

use crate::field::Field;
use serde::{Deserialize, Serialize};

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Counters {
    pub add: u64,
    pub dbl: u64,
    pub inv: u64,
    pub mul: u64,
}

impl Counters {
    #[inline]
    pub fn merge(&mut self, o: &Counters) {
        self.add += o.add;
        self.dbl += o.dbl;
        self.inv += o.inv;
        self.mul += o.mul;
    }
    pub fn group_ops(&self) -> u64 {
        self.add + self.dbl
    }
}

/// Affine point in Montgomery form; `inf` marks the point at infinity.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct Pt {
    pub x: u64,
    pub y: u64,
    pub inf: bool,
}

impl Pt {
    pub const INF: Pt = Pt { x: 0, y: 0, inf: true };
    #[inline(always)]
    pub fn new(x: u64, y: u64) -> Pt {
        Pt { x, y, inf: false }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Curve {
    pub f: Field,
    /// a, b in Montgomery form
    pub a: u64,
    pub b: u64,
    /// standard-form a, b (for JSON / checks)
    pub a_std: u64,
    pub b_std: u64,
}

impl Curve {
    pub fn new(p: u64, a: u64, b: u64) -> Curve {
        let f = Field::new(p);
        assert!(a < p && b < p, "a, b must be reduced mod p");
        // discriminant check: 4a³ + 27b² ≠ 0
        let am = f.to_mont(a);
        let bm = f.to_mont(b);
        let a3 = f.mul(f.mul(am, am), am);
        let d = f.add(f.mul(f.to_mont(4 % p), a3), f.mul(f.to_mont(27 % p), f.sqr(bm)));
        assert!(d != 0, "singular curve");
        Curve { f, a: am, b: bm, a_std: a, b_std: b }
    }

    #[inline]
    pub fn p(&self) -> u64 {
        self.f.p
    }

    pub fn point_from_std(&self, x: u64, y: u64) -> Pt {
        Pt::new(self.f.to_mont(x % self.f.p), self.f.to_mont(y % self.f.p))
    }

    pub fn point_to_std(&self, p: &Pt) -> Option<(u64, u64)> {
        if p.inf {
            None
        } else {
            Some((self.f.from_mont(p.x), self.f.from_mont(p.y)))
        }
    }

    pub fn is_on_curve(&self, pt: &Pt) -> bool {
        if pt.inf {
            return true;
        }
        let f = &self.f;
        let lhs = f.sqr(pt.y);
        let rhs = f.add(f.add(f.mul(f.sqr(pt.x), pt.x), f.mul(self.a, pt.x)), self.b);
        lhs == rhs
    }

    #[inline(always)]
    pub fn neg(&self, pt: &Pt) -> Pt {
        if pt.inf {
            *pt
        } else {
            Pt::new(pt.x, self.f.neg(pt.y))
        }
    }

    /// Full affine addition with all special cases (one inversion). Counted.
    pub fn add(&self, p: &Pt, q: &Pt, c: &mut Counters) -> Pt {
        if p.inf {
            return *q;
        }
        if q.inf {
            return *p;
        }
        let f = &self.f;
        if p.x == q.x {
            if f.add(p.y, q.y) == 0 {
                return Pt::INF;
            }
            return self.dbl(p, c);
        }
        c.add += 1;
        c.inv += 1;
        c.mul += 3;
        let lam = f.mul(f.sub(q.y, p.y), f.inv(f.sub(q.x, p.x)));
        self.finish_add(p, q.x, lam)
    }

    /// Doubling with all special cases (one inversion). Counted.
    pub fn dbl(&self, p: &Pt, c: &mut Counters) -> Pt {
        if p.inf || p.y == 0 {
            return Pt::INF;
        }
        let f = &self.f;
        c.dbl += 1;
        c.inv += 1;
        c.mul += 4;
        let num = f.add(f.mul(f.to_mont(3), f.sqr(p.x)), self.a);
        let lam = f.mul(num, f.inv(f.dbl(p.y)));
        self.finish_dbl(p, lam)
    }

    /// Given the slope λ of the chord through p and q (q.x = `qx`), compute p + q.
    /// Costs 2 field multiplications (λ², λ·(x1−x3)); counted by the caller.
    #[inline(always)]
    pub fn finish_add(&self, p: &Pt, qx: u64, lam: u64) -> Pt {
        let f = &self.f;
        let x3 = f.sub(f.sub(f.sqr(lam), p.x), qx);
        let y3 = f.sub(f.mul(lam, f.sub(p.x, x3)), p.y);
        Pt::new(x3, y3)
    }

    /// Given the tangent slope λ at p, compute 2p (2 multiplications, counted by caller).
    #[inline(always)]
    pub fn finish_dbl(&self, p: &Pt, lam: u64) -> Pt {
        let f = &self.f;
        let x3 = f.sub(f.sqr(lam), f.dbl(p.x));
        let y3 = f.sub(f.mul(lam, f.sub(p.x, x3)), p.y);
        Pt::new(x3, y3)
    }

    /// Numerator of the tangent slope: 3x² + a (one squaring + one multiplication, counted by caller).
    #[inline(always)]
    pub fn tangent_num(&self, p: &Pt) -> u64 {
        let f = &self.f;
        f.add(f.mul(f.to_mont(3), f.sqr(p.x)), self.a)
    }

    /// Left-to-right double-and-add scalar multiplication (k any u64). Counted.
    pub fn mul(&self, k: u64, p: &Pt, c: &mut Counters) -> Pt {
        if p.inf || k == 0 {
            return Pt::INF;
        }
        let mut r = Pt::INF;
        let bits = 64 - k.leading_zeros();
        for i in (0..bits).rev() {
            r = self.dbl(&r, c);
            if (k >> i) & 1 == 1 {
                r = self.add(&r, p, c);
            }
        }
        r
    }

    pub fn eq(&self, p: &Pt, q: &Pt) -> bool {
        if p.inf || q.inf {
            return p.inf && q.inf;
        }
        p.x == q.x && p.y == q.y
    }

    /// Random point by lifting random x (standard-form output converted to Montgomery).
    pub fn random_point<R: rand::Rng>(&self, rng: &mut R) -> Pt {
        let f = &self.f;
        loop {
            let x = rng.random_range(0..f.p);
            let xm = f.to_mont(x);
            let rhs = f.add(f.add(f.mul(f.sqr(xm), xm), f.mul(self.a, xm)), self.b);
            if let Some(y) = f.sqrt_plain(f.from_mont(rhs)) {
                let y = if rng.random_bool(0.5) { y } else { (f.p - y) % f.p };
                return Pt::new(xm, f.to_mont(y));
            }
        }
    }
}

/// Fixed-base comb table: 16 nibble positions × 16 entries, so that k·P costs ≤ 15 additions
/// (plus the table precomputation). Used to start random walks cheaply (u·G + v·Q).
pub struct CombTable {
    pub table: Vec<[Pt; 16]>, // table[i][d] = d·16^i·P
}

impl CombTable {
    pub fn new(curve: &Curve, p: &Pt, c: &mut Counters) -> CombTable {
        let mut table = Vec::with_capacity(16);
        let mut base = *p; // 16^i P
        for _ in 0..16 {
            let mut row = [Pt::INF; 16];
            for d in 1..16 {
                row[d] = curve.add(&row[d - 1], &base, c);
            }
            // next base = 16·base
            let mut nb = base;
            for _ in 0..4 {
                nb = curve.dbl(&nb, c);
            }
            base = nb;
            table.push(row);
        }
        CombTable { table }
    }

    pub fn mul(&self, curve: &Curve, k: u64, c: &mut Counters) -> Pt {
        let mut r = Pt::INF;
        for i in 0..16 {
            let d = ((k >> (4 * i)) & 0xF) as usize;
            if d != 0 {
                r = curve.add(&r, &self.table[i][d], c);
            }
        }
        r
    }
}