package util

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"go-nkode/hashset"
	"math/big"
	r "math/rand"
	"time"
)

type ShuffleTypes interface {
	[]int | int | []uint64 | uint64
}

// fisherYatesShuffle shuffles a slice of bytes in place using the Fisher-Yates algorithm.
func fisherYatesShuffle[T ShuffleTypes](b *[]T) error {
	for i := len(*b) - 1; i > 0; i-- {
		bigJ, err := rand.Int(rand.Reader, big.NewInt(int64(i+1)))
		if err != nil {
			return err
		}
		j := bigJ.Int64()
		(*b)[i], (*b)[j] = (*b)[j], (*b)[i]
	}
	return nil
}

func FisherYatesShuffle[T ShuffleTypes](b *[]T) error {
	return fisherYatesShuffle(b)
}

func RandomBytes(n int) ([]byte, error) {
	b := make([]byte, n)
	_, err := rand.Read(b)
	if err != nil {
		return nil, err
	}
	return b, nil
}

func RandomPermutation(n int) ([]int, error) {
	perm := IdentityArray(n)
	err := fisherYatesShuffle(&perm)
	if err != nil {
		return nil, err
	}
	return perm, nil
}

func GenerateRandomUInt64() (uint64, error) {
	randBytes, err := RandomBytes(8)
	if err != nil {
		return 0, err
	}
	val := binary.LittleEndian.Uint64(randBytes)
	return val, nil
}

func GenerateRandomNonRepeatingUint64(listLen int) ([]uint64, error) {
	if listLen > int(1)<<32 {
		return nil, errors.New("list length must be less than 2^32")
	}
	listSet := make(hashset.Set[uint64])
	for {
		if listSet.Size() == listLen {
			break
		}
		randNum, err := GenerateRandomUInt64()
		if err != nil {
			return nil, err
		}
		listSet.Add(randNum)
	}

	data := listSet.ToSlice()
	return data, nil
}

func XorLists(l0 []uint64, l1 []uint64) ([]uint64, error) {
	if len(l0) != len(l1) {
		return nil, errors.New(fmt.Sprintf("list len mismatch %d, %d", len(l0), len(l1)))
	}

	xorList := make([]uint64, len(l0))
	for idx := 0; idx < len(l0); idx++ {
		xorList[idx] = l0[idx] ^ l1[idx]
	}
	return xorList, nil
}

func EncodeBase64Str(data []uint64) string {
	dataBytes := Uint64ArrToByteArr(data)
	encoded := base64.StdEncoding.EncodeToString(dataBytes)
	return encoded
}

func DecodeBase64Str(encoded string) ([]uint64, error) {
	dataBytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return nil, err
	}
	data := ByteArrToUint64Arr(dataBytes)
	return data, nil
}

func Uint64ArrToByteArr(intArr []uint64) []byte {
	byteArr := make([]byte, len(intArr)*8)
	for idx, val := range intArr {
		startIdx := idx * 8
		endIdx := (idx + 1) * 8
		binary.LittleEndian.PutUint64(byteArr[startIdx:endIdx], val)
	}
	return byteArr
}

func ByteArrToUint64Arr(byteArr []byte) []uint64 {
	intArr := make([]uint64, len(byteArr)/8)
	for idx := 0; idx < len(intArr); idx++ {
		startIdx := idx * 8
		endIdx := (idx + 1) * 8
		intArr[idx] = binary.LittleEndian.Uint64(byteArr[startIdx:endIdx])
	}
	return intArr
}

func IndexOf[T uint64 | int](arr []T, el T) int {
	for idx, val := range arr {
		if val == el {
			return idx
		}
	}
	return -1
}

func IdentityArray(arrLen int) []int {
	identityArr := make([]int, arrLen)

	for idx := range identityArr {
		identityArr[idx] = idx
	}
	return identityArr
}

func ListToMatrix(listArr []int, numbCols int) ([][]int, error) {
	if len(listArr)%numbCols != 0 {
		panic(fmt.Sprintf("Array is not evenly divisible by number of columns: %d mod %d = %d", len(listArr), numbCols, len(listArr)%numbCols))
	}
	numbRows := len(listArr) / numbCols
	matrix := make([][]int, numbRows)
	for idx := range matrix {
		startIdx := idx * numbCols
		endIdx := (idx + 1) * numbCols
		matrix[idx] = listArr[startIdx:endIdx]
	}
	return matrix, nil
}

func MatrixTranspose(matrix [][]int) ([][]int, error) {
	if matrix == nil || len(matrix) == 0 {
		return nil, fmt.Errorf("matrix cannot be nil or empty")
	}

	rows := len(matrix)
	cols := len((matrix)[0])

	// Check if the matrix is not rectangular
	for _, row := range matrix {
		if len(row) != cols {
			return nil, fmt.Errorf("all rows must have the same number of columns")
		}
	}

	transposed := make([][]int, cols)
	for i := range transposed {
		transposed[i] = make([]int, rows)
	}

	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			transposed[j][i] = (matrix)[i][j]
		}
	}

	return transposed, nil
}

func MatrixToList(matrix [][]int) []int {
	var flat []int
	for _, row := range matrix {
		flat = append(flat, row...)
	}
	return flat
}

func Choice[T any](items []T) T {
	r.Seed(time.Now().UnixNano()) // Seed the random number generator
	return items[r.Intn(len(items))]
}