package frand

import (
	crand "crypto/rand"
	"encoding/binary"
	"math"
	"math/big"
	"math/rand/v2"
	"strconv"
	"sync"
)

// An RNG is a cryptographically-strong RNG constructed from the ChaCha stream
// cipher.
type RNG struct {
	r *rand.ChaCha8
}

// Read fills b with random data. It always returns len(b), nil.
func (r *RNG) Read(b []byte) (int, error) {
	return r.r.Read(b)
}

// Bytes is a helper function that allocates and returns n bytes of random data.
func (r *RNG) Bytes(n int) []byte {
	b := make([]byte, n)
	r.Read(b)
	return b
}

// Entropy256 is a helper function that returns 256 bits of random data.
func (r *RNG) Entropy256() (entropy [32]byte) {
	r.Read(entropy[:])
	return
}

// Entropy192 is a helper function that returns 192 bits of random data.
func (r *RNG) Entropy192() (entropy [24]byte) {
	r.Read(entropy[:])
	return
}

// Entropy128 is a helper function that returns 128 bits of random data.
func (r *RNG) Entropy128() (entropy [16]byte) {
	r.Read(entropy[:])
	return
}

// Uint64n returns a uniform random uint64 in [0,n). It panics if n == 0.
func (r *RNG) Uint64n(n uint64) uint64 {
	if n == 0 {
		panic("frand: argument to Uint64n is 0")
	}
	// To eliminate modulo bias, keep selecting at random until we fall within
	// a range that is evenly divisible by n.
	// NOTE: since n is at most math.MaxUint64, max is minimized when:
	//    n = math.MaxUint64/2 + 1 -> max = math.MaxUint64 - math.MaxUint64/2
	// This gives an expected 2 tries before choosing a value < max.
	max := math.MaxUint64 - math.MaxUint64%n
again:
	i := r.r.Uint64()
	if i >= max {
		goto again
	}
	return i % n
}

// Intn returns a uniform random int in [0,n). It panics if n <= 0.
func (r *RNG) Intn(n int) int {
	if n <= 0 {
		panic("frand: argument to Intn is <= 0: " + strconv.Itoa(n))
	}
	// NOTE: since n is at most math.MaxUint64/2, max is minimized when:
	//    n = math.MaxUint64/4 + 1 -> max = math.MaxUint64 - math.MaxUint64/4
	// This gives an expected 1.333 tries before choosing a value < max.
	return int(r.Uint64n(uint64(n)))
}

// BigIntn returns a uniform random *big.Int in [0,n). It panics if n <= 0.
func (r *RNG) BigIntn(n *big.Int) *big.Int {
	i, _ := crand.Int(r, n)
	return i
}

// Float64 returns a random float64 in [0,1).
func (r *RNG) Float64() float64 {
	return float64(r.Uint64n(1<<53)) / (1 << 53)
}

// Perm returns a random permutation of the integers [0,n). It panics if n < 0.
func (r *RNG) Perm(n int) []int {
	m := make([]int, n)
	for i := 1; i < n; i++ {
		j := r.Intn(i + 1)
		m[i] = m[j]
		m[j] = i
	}
	return m
}

// Shuffle randomly permutes n elements by repeatedly calling swap in the range
// [0,n). It panics if n < 0.
func (r *RNG) Shuffle(n int, swap func(i, j int)) {
	for i := n - 1; i > 0; i-- {
		swap(i, r.Intn(i+1))
	}
}

// NewCustom returns a new RNG instance seeded with the provided entropy. It
// panics if len(seed) != 32. The bufsize and rounds parameters are ignored.
func NewCustom(seed []byte, bufsize int, rounds int) *RNG {
	if len(seed) != 32 {
		panic("frand: invalid seed size")
	}
	return &RNG{rand.NewChaCha8(([32]byte)(seed))}
}

// "master" RNG, seeded from crypto/rand; RNGs returned by New derive their seed
// from this RNG. This means we only ever need to read system entropy a single
// time, at startup.
var masterRNG = func() *RNG {
	var seed [32]byte
	if _, err := crand.Read(seed[:]); err != nil {
		panic("not enough system entropy to seed master RNG")
	}
	return &RNG{rand.NewChaCha8(seed)}
}()
var masterMu sync.Mutex

// New returns a new RNG instance. The instance is seeded with entropy from
// crypto/rand, albeit indirectly; its seed is generated by a global "master"
// RNG, which itself is seeded from crypto/rand. This means the frand package
// only reads system entropy once, at startup.
func New() *RNG {
	masterMu.Lock()
	defer masterMu.Unlock()
	return &RNG{rand.NewChaCha8(masterRNG.Entropy256())}
}

// Global versions of each RNG method, leveraging a pool of RNGs.

var rngpool = sync.Pool{
	New: func() interface{} {
		return New()
	},
}

// Read fills b with random data. It always returns len(b), nil.
func Read(b []byte) (int, error) {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Read(b)
}

// Bytes is a helper function that allocates and returns n bytes of random data.
func Bytes(n int) []byte {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Bytes(n)
}

// Entropy256 is a helper function that returns 256 bits of random data.
func Entropy256() [32]byte {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Entropy256()
}

// Entropy192 is a helper function that returns 192 bits of random data.
func Entropy192() [24]byte {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Entropy192()
}

// Entropy128 is a helper function that returns 128 bits of random data.
func Entropy128() [16]byte {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Entropy128()
}

// Uint64n returns a uniform random uint64 in [0,n). It panics if n == 0.
func Uint64n(n uint64) uint64 {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Uint64n(n)
}

// Intn returns a uniform random int in [0,n). It panics if n <= 0.
func Intn(n int) int {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Intn(n)
}

// BigIntn returns a uniform random *big.Int in [0,n). It panics if n <= 0.
func BigIntn(n *big.Int) *big.Int {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.BigIntn(n)
}

// Float64 returns a random float64 in [0,1).
func Float64() float64 {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Float64()
}

// Perm returns a random permutation of the integers [0,n). It panics if n < 0.
func Perm(n int) []int {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	return r.Perm(n)
}

// Shuffle randomly permutes n elements by repeatedly calling swap in the range
// [0,n). It panics if n < 0.
func Shuffle(n int, swap func(i, j int)) {
	r := rngpool.Get().(*RNG)
	defer rngpool.Put(r)
	r.Shuffle(n, swap)
}

// Reader is a global, shared instance of a cryptographically strong pseudo-
// random generator. Reader is safe for concurrent use by multiple goroutines.
var Reader rngReader

type rngReader struct{}

func (rngReader) Read(b []byte) (int, error) { return Read(b) }

// A Source is a math/rand-compatible source of entropy. It is safe for
// concurrent use by multiple goroutines.
type Source struct {
	rng *RNG
	mu  sync.Mutex
}

// Seed uses the provided seed value to initialize the Source to a
// deterministic state.
func (s *Source) Seed(i int64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var seed [32]byte
	binary.LittleEndian.PutUint64(seed[:], uint64(i))
	s.rng.r.Seed(seed)
}

// Int63 returns a non-negative random 63-bit integer as an int64.
func (s *Source) Int63() int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	return int64(s.rng.Uint64n(1 << 63))
}

// Uint64 returns a random 64-bit integer.
func (s *Source) Uint64() uint64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.rng.Uint64n(math.MaxUint64)
}

// NewSource returns a source for use with the math/rand package.
func NewSource() *Source { return &Source{rng: New()} }