app "prime"
    packages { pf: "https://github.com/roc-lang/basic-cli/releases/download/0.7.1/Icc3xJoIixF3hCcfXrDwLCu4wQHtNdPyoJkEbkgIElA.tar.br" }
    imports [pf.Stdout, pf.Task, pf.Arg]
    provides [main] to pf

main : Task.Task {} I32
main =
    max <- Task.await parseArg
    numStrings =
        primesIn max
        |> List.map Num.toStr
        |> Str.joinWith "\n"

    Stdout.line "$(numStrings)"

parseArg =
    args <- Task.await (Arg.list)
    arg =
        args
        |> List.get 1
        |> Result.withDefault "50_000_000"
        |> Str.toU64
    when arg is
        Ok a -> a |> Task.ok
        Err InvalidNumStr -> crash "Passed argument is not a numberr"

primesIn : U64 -> List U64
primesIn = \max ->
    filterPrimesRec = \acc, l ->
        when l is
            [] -> acc
            [head, .. as tail] if isPrimePost acc head ->
                acc |> List.append head |> filterPrimesRec tail

            [_, .. as tail] -> filterPrimesRec acc tail
    [2]
    |> List.concat (List.range { start: At 3u64, end: At max, step: 2 })
    |> \l -> filterPrimesRec [] l

isPrimePost = \list, n ->
    isPrimeRec = \l, limit ->
        when l is
            [_] -> Bool.true
            [head, .. as tail] ->
                if Num.isMultipleOf n head then
                    Bool.false
                else if head > limit then
                    Bool.true
                else
                    isPrimeRec tail limit

            [] -> Bool.true
    sqrtF = Num.sqrt (Num.toF32 n)
    sqrtI = Num.floor sqrtF

    if Num.isZero (sqrtF - Num.toF32 sqrtI) then
        Bool.false
    else
        isPrimeRec list (sqrtI + 1)

countPrimesIn : U64 -> Nat
countPrimesIn = \limit ->
    List.range { start: At 2, end: At limit }
    |> List.countIf \n ->
        if n < 4 then
            Bool.true
        else if Num.isEven n then
            Bool.false
        else
            List.range { start: At 3, end: At (Num.ceiling (Num.sqrt (Num.toF64 n))), step: 2 }
            |> List.any \a -> Num.isMultipleOf n a
            |> Bool.not

expect
    (
        countPrimesIn 1000000
        |> \expected ->
            dbg expected

            expected
    )
    == (
        List.len (primesIn 1000000)
        |> \actual ->
            dbg actual

            actual

    )