require "./euler"

module Euler
  class Prime
    include Iterator(NumType | BigInt)

    def initialize()
      @sieve_size = 16
      @sieved_up_to = 0

      @primes = Array(Bool).new(@sieve_size, true)
      @primes[0] = false
      @primes[1] = false

      @index = 0
      iterative_eratosthenes_sieve
    end

    private def check_and_increase_sieve
      if @index > @primes.size - 1
        @sieve_size *= 2
        until @primes.size == @sieve_size
          @primes << true
        end
      end

      iterative_eratosthenes_sieve
    end

    private def iterative_eratosthenes_sieve
      2.step(to: Math.sqrt(@sieve_size)) do |i|
        if @primes[i]
          (i * i >= @sieved_up_to ? i * i : @sieved_up_to + i - (@sieved_up_to % i)).step(to: @sieve_size-1, by: i) do |j|
            @primes[j] = false
          end
        end
      end
      @sieved_up_to = @sieve_size-1
    end

    def next
      @index += 1
      check_and_increase_sieve

      until @primes[@index]
        @index += 1
        check_and_increase_sieve
      end
      @index
    end

    def self.eratosthenes_sieve(n)
      sieve = Array.new(n + 1, true)
      sieve[0] = false
      sieve[1] = false

      2.step(to: Math.sqrt(n)) do |i|
        if sieve[i]
          (i * i).step(to: n, by: i) do |j|
            sieve[j] = false
          end
        end
      end

      result = sieve.map_with_index { |b, i| b ? i : nil }.compact
    end
  end
end