package main

import (
	"errors"
	"strconv"
	"strings"
	"testing"
)

// primesByTrialDivision is a deliberately naive independent implementation, so
// a bug in the sieve is not mirrored by the thing checking it.
func primesByTrialDivision(limit int) []int {
	var primes []int
	for n := 2; n <= limit; n++ {
		isPrime := true
		for d := 2; d*d <= n; d++ {
			if n%d == 0 {
				isPrime = false
				break
			}
		}
		if isPrime {
			primes = append(primes, n)
		}
	}
	return primes
}

func primesFrom(isPrime []bool) []int {
	var primes []int
	for n, prime := range isPrime {
		if prime {
			primes = append(primes, n)
		}
	}
	return primes
}

func TestSieveMatchesTrialDivision(t *testing.T) {
	const limit = 5000

	got := primesFrom(sieve(limit))
	want := primesByTrialDivision(limit)

	if len(got) != len(want) {
		t.Fatalf("sieve(%d) produced %d primes, trial division produced %d", limit, len(got), len(want))
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("sieve(%d) primes[%d] = %d, want %d", limit, i, got[i], want[i])
		}
	}
}

func TestSieveKnownSmallPrimes(t *testing.T) {
	want := []int{2, 3, 5, 7, 11, 13, 17, 19, 23, 29}

	got := primesFrom(sieve(30))

	if len(got) != len(want) {
		t.Fatalf("sieve(30) = %v, want %v", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Errorf("sieve(30)[%d] = %d, want %d", i, got[i], want[i])
		}
	}
}

// The prime-counting function at powers of ten is a published, independently
// known result, so it pins the sieve at a scale trial division cannot reach.
func TestSievePrimeCounts(t *testing.T) {
	cases := []struct {
		limit int
		want  int
	}{
		{10, 4},
		{100, 25},
		{1_000, 168},
		{10_000, 1_229},
		{100_000, 9_592},
		{1_000_000, 78_498},
	}

	for _, tc := range cases {
		t.Run(strconv.Itoa(tc.limit), func(t *testing.T) {
			if got := len(primesFrom(sieve(tc.limit))); got != tc.want {
				t.Errorf("pi(%d) = %d, want %d", tc.limit, got, tc.want)
			}
		})
	}
}

func TestSieveBelowTwoHasNoPrimes(t *testing.T) {
	for _, limit := range []int{-1, 0, 1} {
		if got := primesFrom(sieve(limit)); len(got) != 0 {
			t.Errorf("sieve(%d) = %v, want no primes", limit, got)
		}
	}
}

func TestSieveExactlyTwo(t *testing.T) {
	got := primesFrom(sieve(2))
	if len(got) != 1 || got[0] != 2 {
		t.Errorf("sieve(2) = %v, want [2]", got)
	}
}

// 0 and 1 are not prime, and the slice is indexed by the number itself, so
// those positions must stay false rather than shifting every later index.
func TestSieveIndexingIsByNumber(t *testing.T) {
	isPrime := sieve(10)

	if len(isPrime) != 11 {
		t.Fatalf("len(sieve(10)) = %d, want 11", len(isPrime))
	}
	if isPrime[0] || isPrime[1] {
		t.Error("0 and 1 must not be marked prime")
	}
	if !isPrime[7] {
		t.Error("7 must be marked prime")
	}
	if isPrime[9] {
		t.Error("9 must not be marked prime")
	}
}

func TestWritePrimesFormat(t *testing.T) {
	var buf strings.Builder

	count, err := writePrimes(&buf, sieve(20))
	if err != nil {
		t.Fatalf("writePrimes returned error: %v", err)
	}

	const want = "2\n3\n5\n7\n11\n13\n17\n19\n"
	if buf.String() != want {
		t.Errorf("output = %q, want %q", buf.String(), want)
	}
	if count != 8 {
		t.Errorf("count = %d, want 8", count)
	}
}

func TestWritePrimesEmpty(t *testing.T) {
	var buf strings.Builder

	count, err := writePrimes(&buf, sieve(1))
	if err != nil {
		t.Fatalf("writePrimes returned error: %v", err)
	}
	if buf.String() != "" {
		t.Errorf("output = %q, want empty", buf.String())
	}
	if count != 0 {
		t.Errorf("count = %d, want 0", count)
	}
}

type errWriter struct{ err error }

func (w errWriter) Write([]byte) (int, error) { return 0, w.err }

func TestWritePrimesPropagatesWriteError(t *testing.T) {
	wantErr := errors.New("disk full")

	// Enough primes to overflow bufio's buffer, so the failure surfaces during
	// writing rather than only at the final flush.
	_, err := writePrimes(errWriter{err: wantErr}, sieve(200_000))

	if !errors.Is(err, wantErr) {
		t.Errorf("writePrimes error = %v, want %v", err, wantErr)
	}
}