import { describe, expect, it } from "@suaveplan/testing/runner";
import type { Rng } from "@suaveplan/types";
import {
    isProbablePrime,
    SMALL_PRIMES,
    uniformBigint,
} from "./miller-rabin.js";

function seededRng(seed = 1): Rng {
    let state = seed >>> 0;
    return {
        nextUint32: () => {
            state = (Math.imul(state, 1664525) + 1013904223) >>> 0;
            return state;
        },
        nextFloat: () => {
            state = (Math.imul(state, 1664525) + 1013904223) >>> 0;
            return state / 0x1_0000_0000;
        },
    };
}

describe("isProbablePrime", () => {
    it("accepts known small primes", { timeout: 5000 }, () => {
        const rng = seededRng();
        for (const p of [3n, 7n, 97n, 251n, 2147483647n]) {
            expect(isProbablePrime(p, rng)).toBe(true);
        }
    });

    it("rejects known composites including Carmichael numbers", {
        timeout: 5000,
    }, () => {
        const rng = seededRng();
        for (const c of [4n, 9n, 15n, 91n, 561n, 1105n, 1729n]) {
            expect(isProbablePrime(c, rng)).toBe(false);
        }
    });

    it("rejects values < 2 and accepts 2 / 3", { timeout: 5000 }, () => {
        const rng = seededRng();
        expect(isProbablePrime(0n, rng)).toBe(false);
        expect(isProbablePrime(1n, rng)).toBe(false);
        expect(isProbablePrime(2n, rng)).toBe(true);
        expect(isProbablePrime(3n, rng)).toBe(true);
    });

    it("deterministic with seeded RNG", { timeout: 5000 }, () => {
        const a = isProbablePrime(101n, seededRng(42), 10);
        const b = isProbablePrime(101n, seededRng(42), 10);
        expect(a).toBe(b);
    });

    it("trial-division short-circuits even composites", {
        timeout: 5000,
    }, () => {
        // n ≡ 0 (mod 251) but odd — 251 * 9 = 2259
        expect(isProbablePrime(2259n, seededRng())).toBe(false);
    });

    it("SMALL_PRIMES contains 251 and starts at 3", () => {
        expect(SMALL_PRIMES[0]).toBe(3n);
        expect(SMALL_PRIMES.at(-1)).toBe(251n);
    });

    it("non-trivial sqrt of 1 path: Carmichael 561 across many seeds", {
        timeout: 10000,
    }, () => {
        // Witness a=2 against n=561 squares to 1 (non-trivial sqrt
        // of 1) — covers the `if (x === 1n) return true` branch.
        // Iterate many seeds so at least one round picks the right
        // witness early in the sequence.
        for (let seed = 1; seed <= 50; seed++) {
            expect(isProbablePrime(561n, seededRng(seed), 32)).toBe(false);
        }
    });

    it("non-trivial sqrt of 1 path: n=109969 with deterministic witness a=35", {
        timeout: 5000,
    }, () => {
        // n=109969 = 313*317 (composite, NOT divisible by any
        // SMALL_PRIMES entry up to 251 so Miller-Rabin runs in full).
        // n-1=109968=2^4 * 6873 → s=4, d=6873. Witness a=35:
        //   35^6873 mod 109969 = X (≠ 1, ≠ n-1)
        //   X^2 mod 109969 = ... eventually = 1 inside the squaring
        // loop — covers `if (x === 1n) return true` (the non-trivial
        // sqrt of 1 branch at line 122).
        //
        // uniformBigint with min=2n, max=109967n: range=109966
        // (17 bits, 3 bytes). topMask = (1 << (17 - 16)) - 1 = 1.
        // We need v=33 (so a = min + v = 2 + 33 = 35). Bytes BE
        // [0, 0, 33] → first byte masked to 0 (0 & 1 = 0), full
        // value = 33.
        const witnessRng: Rng = {
            nextUint32: () => 33,
            nextFloat: () => 0,
            nextBytes: (n: number) => {
                const out = new Uint8Array(n);
                if (n >= 1) out[n - 1] = 33;
                return out;
            },
        };
        expect(isProbablePrime(109969n, witnessRng, 1)).toBe(false);
    });
});

describe("uniformBigint", () => {
    it("returns min when range degenerate", { timeout: 5000 }, () => {
        expect(uniformBigint(5n, 4n, seededRng())).toBe(5n);
    });

    it("returns values within [min, max] across many draws", {
        timeout: 5000,
    }, () => {
        const rng = seededRng(99);
        for (let i = 0; i < 200; i++) {
            const v = uniformBigint(10n, 100n, rng);
            expect(v >= 10n && v <= 100n).toBe(true);
        }
    });

    it("covers the modular-reduction fallback path", { timeout: 5000 }, () => {
        // Force the fallback: provide an Rng that always returns bytes
        // with the top bit set so every rejection-sampled candidate
        // exceeds the range. The fallback then folds via modular
        // reduction.
        let i = 0;
        const overflowRng: Rng = {
            nextUint32: () => 0xffffffff,
            nextFloat: () => 0,
            nextBytes: (n: number) => {
                i++;
                return new Uint8Array(n).fill(0xff);
            },
        };
        const v = uniformBigint(0n, 100n, overflowRng);
        expect(v >= 0n && v <= 100n).toBe(true);
        expect(i).toBeGreaterThan(0);
    });
});