use rug::{Complete, Integer, rand::RandState};
use std::io;

fn solovey_strassen_test(n: &Integer, k: u32) -> bool {
    let mut rand = RandState::new();
    for _ in 0..k {
        let bound = n.clone() - 3u32;
        let a = bound.random_below(&mut rand) + 2u32;

        let gcd = a.gcd_ref(n).complete();
        if gcd > 1 {
            return false;
        }

        let x1 = legendre(&a, n);
        let x2 = jacobi(&a, n);
        if x1 != x2 {
            return false;
        }
    }
    return true;
}

fn jacobi(a: &Integer, p: &Integer) -> Integer {
    let factors = factorize(&p);
    let mut result = Integer::from(1u32);
    for factor in &factors {
        let x = legendre(a, &factor);
        result *= x;
    }
    return result;
}

fn legendre(a: &Integer, p: &Integer) -> Integer {
    let power = (p.clone() - 1u32)
                .div_exact_ref(&Integer::from(2u32))
                .complete();
    let result = Integer::from(a.pow_mod_ref(&power, &p).unwrap());
    if result == 0 || result == 1 {
        result
    }
    else {
        Integer::from(-1)
    }
}

fn main() {
    println!("Введите положительное нечётное число n:");
    let mut p = String::new();
    let stdin = io::stdin();
    match stdin.read_line(&mut p) {
        Ok(_) => (),
        Err(_) => {
            println!("Произошла ошибка при чтении из стандартного ввода");
            return ();
        }
    }
    let p = match Integer::parse(p) {
        Ok(parsed) => parsed.complete(),
        Err(_) => {
            println!("Не удалось считать число");
            return ();
        }
    };
    if p.is_divisible(&Integer::from(2u32)) {
        println!("Введено чётное число");
        return ();
    }

    println!("Введите количество раундов k:");
    let mut k = String::new();
    match stdin.read_line(&mut k) {
        Ok(_) => (),
        Err(_) => {
            println!("Произошла ошибка при чтении из стандартного ввода");
            return ();
        }
    }
    let k = match k.trim_end().parse::<u32>() {
        Ok(parsed) => parsed,
        Err(_) => {
            println!("Не удалось считать число");
            return ();
        }
    };

    let is_prime = solovey_strassen_test(&p, k);
    if is_prime {
        let probability = 1f64 - 2f64.powf(-(k as f64));
        println!("Число {} является простым с вероятностью {}", &p, probability);
    }
    else {
        println!("Число не является простым");
    }
}

fn factorize(n: &Integer) -> Vec<Integer> {
    let mut factors = Vec::new();
    let mut front = vec![n.clone()];

    while !front.is_empty() {
        let num = front.pop().unwrap();
        match lehman(&num) {
            Some(factor) => {
                let other = num.div_exact_ref(&factor).complete();
                front.push(factor);
                front.push(other);
            },
            None => factors.push(num)
        };
    }

    return factors;
}

fn lehman(n: &Integer) -> Option<Integer> {
    let one = &Integer::from(1u32);
    let mut a = Integer::from(2u32);
    let root3 = n.root_ref(3).complete();

    while a <= root3 {
        let (_, r) = n.div_rem_ref(&a).complete();
        if r == 0 {
            return Some(a);
        }
        a += one;
    }

    let root6 = n.root_ref(6).complete();
    let mut k = Integer::from(1u32);
    while k <= root3 {
        let mut d = Integer::ZERO;
        let sqrtk4 = k.sqrt_ref().complete() * 4;
        let (r6sk4, _) = root6.div_rem_ref(&sqrtk4).complete();
        while d <= (&r6sk4 + one).complete() {
            let kn4: Integer = k.clone() * n.clone() * 4;
            let sqrt_kn4 = kn4.sqrt_ref().complete();

            let number = (sqrt_kn4.clone() + d.clone()).square() - kn4.clone();
            if number.is_perfect_square() {
                let big_a = sqrt_kn4 + d.clone();
                let big_b = (big_a.square_ref() - kn4).sqrt();

                let gcd1 = (big_a.clone() + big_b.clone()).gcd_ref(n).complete();
                let gcd2 = (big_a.clone() - big_b.clone()).gcd_ref(n).complete();

                if &gcd1 > one && &gcd1 < n {
                    return Some(gcd1);
                }
                else if &gcd2 > one && &gcd2 < n {
                    return Some(gcd2);
                }
            }

            d += one;
        }
        k += one;
    }

    return None;
}