package transform

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/binary"
	"io"
	"slices"
	"strconv"
	"strings"
	"unicode/utf16"

	"github.com/vulncheck-oss/go-exploit/output"
	"golang.org/x/text/cases"
	"golang.org/x/text/language"
)

func StringToUnicodeByteArray(s string) []byte {
	//nolint:prealloc
	var byteArray []byte
	for _, r := range utf16.Encode([]rune(s)) {
		byteArray = append(byteArray, byte(r&0xFF), byte(r>>8&0xFF))
	}

	return byteArray
}

// Given gzip encoded data, return the decompressed data.
func Inflate(compressed []byte) ([]byte, bool) {
	reader, err := gzip.NewReader(bytes.NewBuffer(compressed))
	if err != nil {
		output.PrintFrameworkError(err.Error())

		return []byte{}, false
	}
	defer reader.Close()

	inflated, err := io.ReadAll(reader)
	if err != nil {
		output.PrintFrameworkError(err.Error())

		return []byte{}, false
	}

	return inflated, true
}

// EncodeBase64 uses standard format non-URL safe base64 encoding to encode a string.
func EncodeBase64(s string) string {
	return base64.StdEncoding.EncodeToString([]byte(s))
}

// EncodeBase64URL encodes to URL safe base64 with padding.
func EncodeBase64URL(s string) string {
	return base64.URLEncoding.EncodeToString([]byte(s))
}

// EncodeBase64Chunks creates a slice of maxiumum size base64 strings, where the maxChunkSize is
// the calculated base64 chunk size and not of the original data. This is useful when you know you
// need to send some data in a chunked manner and the target contains a data size restriction, but
// you do not want to guess at the size of encoded data.
//
// If a chunk size maximum is requested that is larger than the encoded string will be the only
// chunk.
func EncodeBase64Chunks(s string, maxChunkSize uint) []string {
	// An example helps demonstrate why this is useful. Take the following string:
	//
	// 1234567890123456789012345678901234567890
	//
	// If you need to send this data to a target and the target limits you to 10 bytes of
	// base64 data, you cannot just split the string of base64 if the parser is strict, and you
	// also can't just split the raw data before encoding in a predictable manner due to
	// padding. For example:
	//
	// $ printf '1234567890' | base64 -w0 | wc -c
	// 16
	//
	// 1/3 is often stated, but misses padding, as you can see in the 10-10/3 example:
	//
	// $ printf '1234567' | base64 -w0 | wc -c
	// 12
	//
	// The optimal size is actually to ensure the block smaller fits, which 1234, 12345, and
	// 123456 all fit into. This means the optimal fit for the first block to use the most
	// space possible is 123456:
	//
	// $ printf '123456' | base64 -w0 | wc -c
	// 8
	//
	// While the n/3-1 rule works for most cases of pre-base64 encoded data, there is the need
	// to ensure you minimize requests by figuring out what the best block size is. That's what
	// all *this* (hand waving) does.

	// corner case, fail exit early
	if len(s) == 0 {
		return []string{}
	}
	// calculate the maximum base64 size with padding
	maxSize := func(n int) int {
		return (((4 * n / 3) + 3) & ^3)
	}
	// start with a chunk size that is 2/3 the size and subtract one, this normally gives the
	// closest fit, but because of "computer numbers" rounding can be iffy so this ensures that
	// the chunk size calculation iterates to the best block size.
	chunkSize := (len(s) / int(maxChunkSize)) - (len(s) / int(maxChunkSize) / 3) - 1
	for {
		if maxSize(chunkSize) > int(maxChunkSize) {
			chunkSize--

			break
		}
		chunkSize++
	}
	chunks := []string{}
	for c := range slices.Chunk([]byte(s), chunkSize) {
		chunks = append(chunks, base64.StdEncoding.EncodeToString(c))
	}

	return chunks
}

// DecodeBase64 decodes base64 with standard encoding.
func DecodeBase64(s string) string {
	decoded, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		output.PrintFrameworkError(err.Error())

		return ""
	}

	return string(decoded)
}

// DecodeBase64URL decodes URL safe base64 variant.
func DecodeBase64URL(s string) string {
	// Internally uses the non-padded version for decoding so this will allow forcing padded
	// data to be non-padded and support both transparently.
	s = strings.ReplaceAll(s, `=`, ``)
	decoded, err := base64.RawURLEncoding.DecodeString(s)
	if err != nil {
		output.PrintFrameworkError(err.Error())

		return ""
	}

	return string(decoded)
}

func Title(s string) string {
	return cases.Title(language.Und, cases.NoLower).String(s)
}

// URL encode every character in the provided string.
func URLEncodeString(inputString string) string {
	encodedChars := ""
	for _, char := range inputString {
		encodedChars += "%" + strconv.FormatInt(int64(char), 16)
	}

	return encodedChars
}

// PackLittleInt16 packs a little-endian 16-bit integer as a string.
func PackLittleInt16(n int) string {
	var packed strings.Builder

	err := binary.Write(&packed, binary.LittleEndian, int16(n))
	if err != nil {
		output.PrintFrameworkError(err.Error())
	}

	return packed.String()
}

// PackLittleInt32 packs a little-endian 32-bit integer as a string.
func PackLittleInt32(n int) string {
	var packed strings.Builder

	err := binary.Write(&packed, binary.LittleEndian, int32(n))
	if err != nil {
		output.PrintFrameworkError(err.Error())
	}

	return packed.String()
}

// PackLittleInt64 packs a little-endian 64-bit integer as a string.
func PackLittleInt64(n int) string {
	var packed strings.Builder

	err := binary.Write(&packed, binary.LittleEndian, int64(n))
	if err != nil {
		output.PrintFrameworkError(err.Error())
	}

	return packed.String()
}

// PackBigInt16 packs a big-endian 16-bit integer as a string.
func PackBigInt16(n int) string {
	var packed strings.Builder

	err := binary.Write(&packed, binary.BigEndian, int16(n))
	if err != nil {
		output.PrintFrameworkError(err.Error())
	}

	return packed.String()
}

// PackBigInt32 packs a big-endian 32-bit integer as a string.
func PackBigInt32(n int) string {
	var packed strings.Builder

	err := binary.Write(&packed, binary.BigEndian, int32(n))
	if err != nil {
		output.PrintFrameworkError(err.Error())
	}

	return packed.String()
}

// PackBigInt64 packs a big-endian 64-bit integer as a string.
func PackBigInt64(n int) string {
	var packed strings.Builder

	err := binary.Write(&packed, binary.BigEndian, int64(n))
	if err != nil {
		output.PrintFrameworkError(err.Error())
	}

	return packed.String()
}
