#![no_std]
#![no_main]

use core::fmt::Write;
use core::panic::PanicInfo;
use lwcpu::{NumberDisplay, CharDisplay};
use riscv_rt::entry;

const SEGMENT_MAX: usize = 8;
const SEGMENT_SIZE: usize = SEGMENT_MAX*SEGMENT_MAX;

// Calculates the primes in the first segment using the simple sieve
// Returns the primes in the primes array with n being the number of found primes
fn simple_sieve(primes: &mut [usize; SEGMENT_SIZE], n: &mut usize) {
    let mut segment = [true; SEGMENT_SIZE];
    segment[0] = false;
    segment[1] = false;

    for x in 2..SEGMENT_MAX {
        if segment[x] {
            for y in ((x * x)..SEGMENT_SIZE).step_by(x) {
                segment[y] = false;
            }
        }
    }
    let mut index = 0;
    for (x,prime) in segment.iter().enumerate() {
        if *prime {
            primes[index] = x;
            index += 1;
        }
    }
    *n = index;
}

fn segment_sieve(primes: &[usize], prime_multiple: &mut [usize]) {
    let mut low: usize = SEGMENT_SIZE;
    let mut high: usize = SEGMENT_SIZE*2;
    for _ in 1..SEGMENT_SIZE { // there are [SEGMENT_SIZE] number of segments
        let mut segment = [true; SEGMENT_SIZE];
        for (i,p) in primes.iter().enumerate() {
            let mut multiple = prime_multiple[i];
            while multiple < high {
                segment[multiple-low] = false;
                multiple += *p;
            }
            prime_multiple[i] = multiple;
        }

        for (x,prime) in segment.iter().enumerate() {
            if *prime {
                NumberDisplay::display_number((low+x) as u32);
            }
        }

        low = high;
        high += SEGMENT_SIZE;
    }
}

#[panic_handler]
fn panic_handler(_info: &PanicInfo) -> ! {
    NumberDisplay::display_number(0xDEADC0DE);
    let mut text = CharDisplay{};
    let _ = write!(text, "Panic: {}", _info.message());
    loop {}
}

#[entry]
fn main() -> ! {
    // There will be no more than [SEGMENT_SIZE] primes in the segment
    let mut primes_buffer = [0; SEGMENT_SIZE];
    let mut n = 0;
    simple_sieve(&mut primes_buffer, &mut n);

    let primes = &primes_buffer[..n];

    let mut prime_multiple = [0; SEGMENT_SIZE];
    for (i,x) in primes.iter().enumerate() {
        NumberDisplay::display_number(*x as u32);
        let mut multiple = *x;
        while multiple < SEGMENT_SIZE {
            multiple += x;
        }
        prime_multiple[i] = multiple;
    }
    segment_sieve(primes, &mut prime_multiple);
    loop {}
}