package main

import (
	"crypto/rand"
	"fmt"
	"io"
	"math/big"
	mrand "math/rand"
	"time"
)

type job struct {
	p      *big.Int // Number to check
	i      int      // Iterations
	result bool
	id     int
}

func main() {
	s1 := mrand.NewSource(time.Now().UnixMicro())
	r1 := mrand.New(s1)
	start := time.Now()

	p, _ := prime(r1, 8096, 1000)

	finish := time.Now()
	elapsed := finish.Sub(start)
	fmt.Printf("p: %v\n", p)
	fmt.Printf("That took: %dms\n", elapsed.Milliseconds())

	start = time.Now()
	p2, _ := rand.Prime(r1, 8096)
	finish = time.Now()
	elapsed = finish.Sub(start)
	fmt.Printf("p2: %v\n", p2)
	fmt.Printf("Normal generation took: %dms\n", elapsed.Milliseconds())

}

func prime(rand io.Reader, bits int, threads int) (p *big.Int, err error) {
	bytes := make([]byte, (bits+7)/8)

	p = new(big.Int)
	jobs := make(chan job)
	results := make(chan job)
	done := new(bool)
	*done = false

	for i := 0; i < threads; i++ {
		go worker(i, jobs, results, done)
	}

	for i := 0; ; i++ {
		_, err = io.ReadFull(rand, bytes)
		if err != nil {
			return nil, err
		}
		// This makes sure that the least significant bit is always 1
		// meaning that the number is always odd, since even numbers
		// aren't prime
		bytes[len(bytes)-1] |= 1

		p.SetBytes(bytes)

		newjob := &job{
			p:  p,
			i:  10,
			id: i,
		}

		jobs <- *newjob

		r := <-results
		if r.result {
			fmt.Printf("Found one, closing the channel and returning\n")
			close(jobs)
			*done = true
			close(results)
			return r.p, nil
		}
	}

}

func worker(id int, jobs <-chan job, results chan<- job, done *bool) {
	//fmt.Printf("Worker %d starting now!\n", id)
	for j := range jobs {
		j.result = j.p.ProbablyPrime(j.i)
		//fmt.Printf("Worker\t %d checked a number\n", id)
		if *done {
			return
		}
		results <- j
	}
}