package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"strconv"
	"time"
)

func calculateProgress(start, end, step, index int) float32 {
	steps := (end - start) / step
	current := index - start
	if current == 0 {
		return 0
	}
	return float32(current) / float32(steps)
}

func printProgress(progress <-chan float32, message string) {
	for {
		time.Sleep(time.Second)
		p, channelOpen := <-progress
		if !channelOpen {
			break
		}
		fmt.Printf("%s %.2f%%\n", message, p*100)
	}
}

func checkIsDivisibleByPrime(number, offset, stride int, primes *[]int, resultChannel chan bool, ctx context.Context) {
	for i := offset; i < len(*primes); i += stride {
		select {
			case <-ctx.Done():
				resultChannel <- false
				return
			default:
				// only have to check primes, because all other numbers are multiple of primes.
				// if a number is divisible by a multiple of a prime, it is by definition also divisible by the prime itself.
				if number % (*primes)[i] == 0 {
					resultChannel <- true
					return
				}
		}
	}
	resultChannel <- false
}

func generatePrimes(upperLimit int, loadList bool) []int {
	var primes []int

	bootTime := time.Now()

	if loadList {
		primes = loadPrimes()
	} else {
		primes = []int{2, 3, 5}
	}

	numRoutines := runtime.NumCPU()

	fmt.Printf("Startup time: %v\nCalculating with %d routines\n\n", time.Now().Sub(bootTime), numRoutines)

	progress := make(chan float32)
	go printProgress(progress, "Generating primes")

	startTime := time.Now()

	for i := (primes[len(primes) - 1] + 1); i <= upperLimit; i++ {
		select {
		case progress <- calculateProgress(6, upperLimit, 1, i):
		default:
		}

		isPrime := true

		checkCtx, checkCancel := context.WithCancel(context.Background())
		resultChannel := make(chan bool)

		for r := 0; r < numRoutines; r++ {
			go checkIsDivisibleByPrime(i, r, numRoutines, &primes, resultChannel, checkCtx)
		}

		for r := 0; r < numRoutines; r++ {
			if <-resultChannel {
				isPrime = false
				checkCancel()
				// can't break here, because all channel writes have to be read to avoid deadlocks
			}
		}

		if isPrime {
			primes = append(primes, i)
		}
	}
	close(progress)

	endTime := time.Now()

	fmt.Printf("Prime generation took %v\nLargest prime found: %d\nTotal prime number count: %d\n", endTime.Sub(startTime), primes[len(primes) - 1], len(primes))

	return primes
}

func calculatePrimeParts(number int, primes []int) []int {
	// don't calculate if number is a prime itself
	if primes[len(primes)-1] == number {
		return []int{number}
	}

	progress := make(chan float32)

	go printProgress(progress, "Calculating")

	var primeParts []int

	for i := len(primes) - 1; i >= 0; i-- {
		select {
		case progress <- 1.0 - calculateProgress(0, len(primes) - 1, 1, i):
		default:
		}
		flooredDiv := number / primes[i]
		if flooredDiv == 0 {
			continue
		}
		primeParts = append(primeParts, primes[i])
		number = number - (flooredDiv * primes[i])
		if number == 0 {
			break
		}
	}
	close(progress)
	if number != 0 {
		primeParts = append(primeParts, 1)
	}
	return primeParts
}

func main() {
	var primeList bool
	var dontLoad bool

	flag.BoolVar(&primeList, "p", false, "Only calculate and print prime list")
	flag.BoolVar(&dontLoad, "d", false, "Don't load precalculated primes, calculate from 0")
	flag.Parse()
	numStr := flag.Arg(0)

	if numStr == "" {
		fmt.Printf("Usage: %s [-p|d] <number>\n", os.Args[0])
		flag.PrintDefaults()
		return
	}

	number, err := strconv.Atoi(numStr)
	if err != nil {
		log.Fatal(err)
	}

	if primeList {
		onlyGenerate(number, dontLoad)
	} else {
		calculate(number, dontLoad)
	}
}

func onlyGenerate(number int, dontLoad bool) {
	primes := generatePrimes(number, !dontLoad)

	file, err := os.Create("prime.txt")
	if err != nil {
		log.Fatal(err)
	}

	writer := bufio.NewWriter(file)

	for _, prime := range primes {
		writer.WriteString(fmt.Sprintf("%d\n", prime))
	}

	writer.Flush()
	file.Close()
}

func calculate(number int, dontLoad bool) {
	primes := generatePrimes(number, !dontLoad)

	primeParts := calculatePrimeParts(number, primes)

	sum := 0

	for i, prime := range primeParts {
		sum += prime
		print(prime)

		// if not last element
		if i < len(primeParts) - 1 {
			print(" + ")
		}
	}
	print(" = ")
	println(sum)
}

func loadPrimes() []int {
	file, err := os.Open("prime.txt")
	if err != nil {
		log.Fatal(err)
	}

	var primes []int

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		nextPrime, err := strconv.Atoi(scanner.Text())
		if err != nil {
			log.Fatal(err)
		}
		primes = append(primes, nextPrime)
	}

	file.Close()

	return primes
}