package set1

import (
	"os"
	"slices"
)

// 00aaaaaa 00aabbbb 00bbbbcc 00cccccc
// aaaaaaaa bbbbbbbb cccccccc
func Base64Decode(block []byte) (decoded []byte) {

	base64Mapping := map[byte]byte{}

	for i, b := range Base64Alphabet {
		base64Mapping[byte(b)] = byte(i)
	}

	decoded = []byte{}

	for i := 0; i < len(block); i += 4 {
		decodedBytes := []byte{0b00000000, 0b00000000, 0b00000000}

		encodedBytes := block[i : i+4]

		for i, b := range encodedBytes {
			if b == Base64PaddingCharacter {
				encodedBytes[i] = 0b00000000
			} else {
				encodedBytes[i] = base64Mapping[b]
			}
		}

		decodedBytes[0] = (encodedBytes[0] << 2) | (encodedBytes[1] >> 4)
		decodedBytes[1] = (encodedBytes[1] << 4) | (encodedBytes[2] >> 2)
		decodedBytes[2] = (encodedBytes[2] << 6) | encodedBytes[3]

		for _, b := range decodedBytes {
			if b != 0b00000000 {
				decoded = append(decoded, b)
			}
		}
	}

	return decoded
}

func BruteForceXORSingleCharacterPlainEncodedBlock(block []byte) (score int, key byte, result []byte) {
	for i := 0; i < 256; i++ {
		decoded := XORDecodePlain(block, byte(i))
		decodedScore := ScoreTextByAlphabeticFrequency(decoded)

		if decodedScore > score {
			score = decodedScore
			key = byte(i)
			result = decoded
		}
	}

	return
}

func HammingDistance(str1, str2 []byte) (distance int) {
	for i := 0; i < len(str1); i++ {
		if i >= len(str2) {
			distance += CountBits(str1[i])
		} else {
			distance += CountBits(str1[i] ^ str2[i])
		}
	}

	return
}

// 1111

// 1111 -1 = 1110
// 1111 & 1110 = 1110

// 1110 -1 = 1101
// 1110 & 1101 = 1100

// 1100 -1 = 1011
// 1100 & 1011 = 1000

// 1000 -1 = 0111
// 0111 & 1000 = 0000

// ------------------

// 1010

// 1010 -1 = 1001
// 1010 & 1001 = 1000

// 1000 -1 = 0111
// 1000 & 0111 = 0000

func CountBits(x byte) (count int) {
	// This does as many iterations as there are 1 bits in a bitset. Magic.
	for x != 0 {
		x &= x - 1
		count++
	}

	return
}

func ReadChallenge6CipherFromFile() (cipher []byte) {
	file, err := os.ReadFile("./res/set-1-challenge-6.txt")

	if err != nil {
		panic(err)
	}

	return file
}

func NormalizedHammingDistance(strings ...[]byte) (normalizedDistance float32) {
	for i := 1; i < len(strings); i++ {
		normalizedDistance += float32(HammingDistance(strings[i-1], strings[i]))
	}

	normalizedDistance = normalizedDistance / float32(len(strings[0])) / float32(len(strings))

	return
}

type KeySizeRank struct {
	Size int
	Rank float32
}

func RankKeySizeCandidates(minSize, maxSize int, cipher []byte) (keySizeCandidates []KeySizeRank) {
	keySizes := []KeySizeRank{}

	for keySize := minSize; keySize <= maxSize; keySize++ {
		keySizes = append(keySizes, KeySizeRank{
			Size: keySize,
			Rank: NormalizedHammingDistance(
				cipher[0:keySize],
				cipher[keySize:keySize*2],
			),
		})
	}

	slices.SortFunc(keySizes, func(ksr1 KeySizeRank, ksr2 KeySizeRank) int {
		if ksr1.Rank > ksr2.Rank {
			return 1
		} else if ksr1.Rank == ksr2.Rank {
			return 0
		} else {
			return -1
		}
	})

	return keySizes
}

func XORDecrypt(key []byte, cipher []byte) (result []byte) {
	for i := 0; i < len(cipher); i += len(key) {
		cipherBlock := cipher[i : i+len(key)]

		for ib, b := range cipherBlock {
			result = append(result, b^key[ib])
		}
	}

	return

}