use std::env;
use std::fmt;
use std::cmp;
use rand::Rng;
use num_integer::Integer;
use num_traits::identities::{Zero, One};
use num_bigint::{BigInt, RandomBits};

#[derive(Clone)]
enum Point {
    NonZero((BigInt, BigInt)),
    Zero
}

impl fmt::Display for Point {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Point::NonZero((x, y)) => write!(f, "({x}, {y})"),
            Point::Zero => write!(f, "∞")
        }
    }
}

fn inverse_mod(a: &BigInt, n: &BigInt) -> Result<BigInt, BigInt> {
    let egcd = a.extended_gcd(n);
    if egcd.gcd.is_one() { Ok(egcd.x) } else { Err(egcd.gcd) }
}

fn ec_sum(p: &Point, q: &Point, a: &BigInt, n: &BigInt) -> Result<Point, BigInt> {
    // TODO: this should use references
    let (px, py) = if let Point::NonZero((x, y)) = p {
        (x, y)
    } else {
        return Ok(q.clone());
    };

    let (qx, qy) = if let Point::NonZero((x, y)) = q {
        (x, y)
    } else {
        return Ok(p.clone());
    };

    if ((px - qx) % n).is_zero() && ((py + qy) % n).is_zero() {
        return Ok(Point::Zero);
    }

    let k = if px != qx {
        (py - qy) * inverse_mod(&(px - qx), n)?
    } else {
        (3 * px * px + a) * inverse_mod(&(py + qy), n)?
    };

    let x = &k * &k - px - qx;
    let y = k * (px - &x) - py;

    Ok(Point::NonZero((x.clone() % n, y.clone() % n)))
}

fn ec_mul(m: &BigInt, p: &Point, a: &BigInt, n: &BigInt) -> Result<Point, BigInt> {
    if m.is_zero() {
        return Ok(Point::Zero);
    }

    if (m % BigInt::from(2)).is_zero() {
        ec_mul(&(m / 2), &ec_sum(p, p, a, n)?, a, n)
    } else {
        ec_sum(p, &ec_mul(&(m - 1), p, a, n)?, a, n)
    }
}


// Elliptic curve factorization method
// If n is prime, this method goes into an infinite loop
fn find_factor(n: &BigInt) -> BigInt {
    let mut rng = rand::thread_rng();
    let bound = cmp::max(n.sqrt().sqrt().sqrt() + 2, BigInt::from(256));

    loop {
        let a = rng.sample::<BigInt, _>(RandomBits::new(256)) % n;
        let x = rng.sample::<BigInt, _>(RandomBits::new(256)) % n;
        let y = rng.sample::<BigInt, _>(RandomBits::new(256)) % n;
        let mut p = Point::NonZero((x, y));

        let mut m = BigInt::from(2);
        while m < bound && !matches!(p, Point::Zero) {
            match ec_mul(&m, &p, &a, n) {
                Err(f) => {
                    println!("Factor {f} found with a = {a}, m = {m} and p = {p}");
                    return f;
                },
                Ok(new_p) => p = new_p
            }
            m += 1;
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if let Ok(n) = args[1].parse::<BigInt>() {
        let f = find_factor(&n);
        println!("{} = {} * {}", &n, &f, &n / &f);
    } else {
        println!("Invalid argument {} (cannot convert to number)", args[1]);
    }
}