use num_bigint::{BigUint, RandBigInt, ToBigUint};
use num_integer::Integer;
use num_traits::One;
use rand::thread_rng;

/// Miller–Rabin primality test
/// https://en.wikipedia.org/wiki/Miller-Rabin_primality_test
///
/// `k` is the number of iterations (more iterations → higher confidence).
/// This is a basic implementation for odd candidate `n` greater than 2.
pub fn is_probably_prime(n: &BigUint, k: u32) -> bool {
    // Handle small numbers.
    if n == &2u32.to_biguint().unwrap() {
        return true;
    }

    if n < &2u32.to_biguint().unwrap() || n.is_even() {
        return false;
    }

    // Write n − 1 as d * 2^r.
    let one: BigUint = One::one();
    let two: BigUint = 2u32.to_biguint().unwrap();

    let n_minus_one = n - &one;

    let mut d = n_minus_one.clone();
    let mut r = 0u32;

    while d.is_even() {
        d /= &two;
        r += 1;
    }

    // Try k random witnesses.
    let mut rng = thread_rng();

    'witness_loop: for _ in 0..k {
        // Choose a random a in [2, n-2]
        let a = rng.gen_biguint_range(&two, &(n - &two));

        // Compute x = a^d mod n.
        let mut x = a.modpow(&d, n);

        if x == one || x == n_minus_one {
            continue 'witness_loop;
        }

        // Repeat squaring x up to r − 1 times.
        for _ in 0..(r - 1) {
            x = x.modpow(&two, n);
            if x == n_minus_one {
                continue 'witness_loop;
            }
        }

        // If we never hit n − 1, then composite.
        return false;
    }
    true
}

pub fn generate_prime_mod_3_4(bits: u64, rounds: u32) -> BigUint {
    let mut rng = thread_rng();
    let four: BigUint = 4u32.to_biguint().unwrap();
    let three: BigUint = 3u32.to_biguint().unwrap();

    loop {
        // Generate a random number of the required bit-length.
        let mut candidate = rng.gen_biguint(bits);

        // Force the most significant bit to ensure it is exactly `bits` long.
        candidate.set_bit(bits - 1, true);
        // Force the number to be odd.
        candidate.set_bit(0, true);

        // Adjust candidate so that candidate mod 4 == 3.
        // Compute candidate mod 4.
        let rem = &candidate % &four;
        if rem != three {
            // Calculate the adjustment needed.
            // We want candidate + adj ≡ 3 (mod 4); that is, adj ≡ (3 - rem) mod 4.
            let adj = if three >= rem {
                &three - &rem
            } else {
                &three + &four - &rem
            };
            candidate += &adj;
        }

        // Now candidate mod 4 should equal 3.
        if &candidate % &four != three {
            continue; // Should not happen, but be safe.
        }

        // Use Miller–Rabin to test candidate for primality.
        if is_probably_prime(&candidate, rounds) {
            return candidate;
        }
        // Otherwise, try another candidate.
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// This test has chance of false positives since it's probabilistic.
    #[test]
    fn test_is_probably_prime() {
        let primes: Vec<BigUint> = vec![
            BigUint::parse_bytes(b"723646214863847842402314246121044767400617866176733021174245260448070467161519753555151305391831172396032179266088879736498934532967238875067731186605319314486487094813782345277515046149035823394700558031365128080643117834402421935144013956523482034192169360458395261772557972018417296402072764848759", 10).unwrap(),
            BigUint::parse_bytes(b"268263962333296278340388301081833650583348564229009436402694247537863120457689419730666587700766950987474095807568632415133217325374788947918148879191904084190506645642271822238922848768059332086841470078498514866531550241226838886983034780850983546266212727522552823301120544245546076442247019621281", 10).unwrap(),
            BigUint::parse_bytes(b"780316979179005788838703471892068793371544768939561076378203820641657870400087138649016384860083841124919558376637164999621109684339044217326208775683279745886968384731662626620658025388083028243454834224644661932974516055783125538555017192574377520720286008648938568231979020364719643484647961902129", 10).unwrap(),
            BigUint::parse_bytes(b"783671025198992682112959634964787905698619719557738373194519536366273453103042115144772610507228390203926825435595621395808116267173935555322692156198469029023864294375800526944184191931093378457858882431220722817775243587676744651835166948761514278051800658105050374249135721906202152414030775826309", 10).unwrap(),
            BigUint::parse_bytes(b"114288406837499406547552732741383061052894863572677735626102568465969919170114402741547085384743856605779595326952640244199273564193341729111945397295523754661407891360178372372116312851760319552575377963248694758420962490016988707765874984334711407905393818994588867816690596598910287994599226773623", 10).unwrap(),
            BigUint::parse_bytes(b"961901548037893583081258840295866034390773310176366253775445667810599851064250586050396339468859430600645234941487066774638210060293941746865942337333205982508422664447189855574952850359167644652778348230933876780699478432665987194022576502478048244264535150875567551113136844280086912348903874971289", 10).unwrap(),
            BigUint::parse_bytes(b"810327961783413779923342902853291047828538610393840061747638982695517616363172610276265437838324082720309171557216227906072585908285037423898327771759304451581073313718250825722376959089099901152629294365285556214142316224671295775214712025501417340129162614998356447407556284614940830419343154596069", 10).unwrap(),
            BigUint::parse_bytes(b"454714144122862528239309790149196885616593324828166685793258814905365447620370501511767118061263049961557700030750456903650018330037078841617668619030424125469578524711042124583419218861229437961029104448364686795365893249168551162671874789888134646816513322705015243576954880991642886515757494254667", 10).unwrap(),
            BigUint::parse_bytes(b"299912922453745379330945382700853905095713797966780365239274055194104040679024227632579691012245229549449540626544013218573948479074373395273033811581117143614653378585775286148244890322917470335843727405772749436431799738228530579087880286816355000236580795379259714642327676881458540815945300666161", 10).unwrap(),
            BigUint::parse_bytes(b"376003947134918788492667049006858112525261477602024892881243407277975729638079808355736777378039431492585839429794604176393826090578267002411236819553297345047777740428504027180799354733010246832276383653403857249078308266896127980484004668688670243017355407377733246174535348005452263544141849401979", 10).unwrap(),
        ];

        for p in primes.iter() {
            assert!(is_probably_prime(&p, 64));
        }
    }

    #[test]
    fn test_is_probably_prime_composites() {
        let composites: Vec<BigUint> = vec![
            BigUint::parse_bytes(b"376003947134918788492667049006858112525261477602024892881243407277975729638079808355736777378039431492585839429794604176393826090578267002411236819553297345047777740428504027180799354733010246832276383653403857249078308266896127980484004668688670243017355407377733246174535348005452263544141849401978", 10).unwrap(),
            BigUint::parse_bytes(b"215231561558920190548208026347283855893349614579652886715109041246408707968244796209282350597895810733095104924634951652047270639587011501367107805226413278625593440341692931130648759932592989139786464960914139157114834353651728654113090571685005692273343539839816029944416477359630133856900278127196", 10).unwrap(),
            BigUint::parse_bytes(b"749740332669263695852915686823796627478248903975029332069160620969447428491611565267586339403720492925006681992807328998947128762023955011291116754751950431300094251835880417511251321687004478207880237612242291350451356839420408200814283441229483559638270099621526259859618448556951589488935694922070", 10).unwrap(),
            BigUint::parse_bytes(b"169634000834057905687114185586149242948139045677318813563962611831640519981755545740448510103217026033366710697221717847963076219297087916994382041061663714750161811176003440851268811309999453614154110754445690358783609141283608528146586702857821814537749897378979649737926839924019514134727091160596", 10).unwrap(),
            BigUint::parse_bytes(b"391405423376058365603029493927540848246012397355482118997842261901662598117428494623814607612916906669639217808344280832395767938093281898200653708872133101732068607931282089810121182927184303240249371585726676226781755240078640338182784311458749444597839572385973551700960543208703821367720798778072", 10).unwrap(),
            BigUint::parse_bytes(b"145039293606441299096899867755014610341218018203667264841882003871966393556065187657468909971071812669502147831045878456015615749981859567162502647235340765709384621383450245978156289943326007863506929089750917534169817946092363165262671439744821303778355013113872010222371955198429114197736184765704", 10).unwrap(),
            BigUint::parse_bytes(b"338068116382585626410927100767745326752329150652235801367160858233547248828235927245043642851496668587705981602011646708907509708706737409743516100542296115834435744682149323258036975223060910642424686044835405231352594163800211063010578142263828931218684092895126045712869105445627601057623777552090", 10).unwrap(),
            BigUint::parse_bytes(b"733953124218193165395301551993038545406364629775221343783594613066305176245028773006556151249847554247355736690261867825939967080139988834452933842378528382325603700037881199289040333697671258970839307337751547799788219379733886335171402446580393040976829094147103849993165402411253055490791297710398", 10).unwrap(),
            BigUint::parse_bytes(b"913203980385282094763740857489413119990150673887971726908661546534253494037072205386715229605079865032590651982202933180796243044488005878787438030787026922448019409247593693512864665427511952738234654431289764435314410753974607679676812065736404760888000826732781882360969951775729150509866946286836", 10).unwrap(),
            BigUint::parse_bytes(b"232672468411571864128377684111610045835673835338188777956578101157986193934981310713021354012943502884801943521144084387464428198189142123559463576210566925881741632803965311463508844160091019108613782710387285819179522210515150555890570342556044250014755588708967402313183475521240406597634471243814", 10).unwrap(),
        ];

        for c in composites.iter() {
            assert!(!is_probably_prime(&c, 64));
        }
    }

    #[test]
    fn test_generate_prime_mod_3_4() {
        let bits = 2048;
        let rounds = 64;
        let prime = generate_prime_mod_3_4(bits, rounds);
        assert!(is_probably_prime(&prime, rounds));
        assert_eq!(prime.bits(), bits);
        assert_eq!(
            &prime % 4u32.to_biguint().unwrap(),
            3u32.to_biguint().unwrap()
        );
    }
}