package secret

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
	"strings"

	. "gitea.zaclys.com/bvaudour/gob/option"
	"golang.org/x/crypto/blowfish"
)

const (
	AES      = "aes"
	BLOWFISH = "blowfish"
)

type CipherFunc func([]byte) (cipher.Block, error)

var (
	encoders = map[string]CipherFunc{
		AES: aes.NewCipher,
		BLOWFISH: func(key []byte) (c cipher.Block, err error) {
			c, err = blowfish.NewCipher(key)
			return
		},
	}
	blockSize = map[string]int{
		AES:      aes.BlockSize,
		BLOWFISH: blowfish.BlockSize,
	}
)

func addBase64Padding(value string) string {
	m := len(value) % 4
	if m != 0 {
		value += strings.Repeat("=", 4-m)
	}

	return value
}

func encodeToBase64[T ~[]byte | ~string](value []byte) T {
	encoded := strings.Replace(base64.URLEncoding.EncodeToString(value), "=", "", -1)
	return T(encoded)
}

func decodeFromBase64[T ~[]byte | ~string](value T) Result[[]byte] {
	decoded, err := base64.URLEncoding.DecodeString(addBase64Padding(string(value)))
	if err != nil {
		return Err[[]byte](err)
	}
	return Ok([]byte(decoded))
}

func pad(src []byte, bs int) []byte {
	padding := bs - len(src)%bs
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}

func unpad(src []byte) Result[[]byte] {
	length := len(src)
	unpadding := int(src[length-1])

	if unpadding > length {
		return Err[[]byte](errors.New("unpad error. This could happen when incorrect encryption key is used"))
	}

	return Ok(src[:(length - unpadding)])
}

// Encrypt chiffre une donnée d’entrée (avec une clé aléatoire) selon
// l’algorithme demandé. Les algorithmes supportés sont :
// - blowfish
// - aes
func Encrypt[T ~[]byte | ~string](data, key T, algorithm string) Result[T] {
	c, ok := encoders[algorithm]
	if !ok {
		return Err[T](errAlgorithm(algorithm))
	}

	block, err := c([]byte(key))
	if err != nil {
		return Err[T](err)
	}

	bs := blockSize[algorithm]
	msg := pad([]byte(data), bs)
	cipherText := make([]byte, bs+len(msg))
	iv := cipherText[:bs]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return Err[T](err)
	}

	cfb := cipher.NewCFBEncrypter(block, iv)
	cfb.XORKeyStream(cipherText[bs:], msg)

	return Ok(encodeToBase64[T](cipherText))
}

// Decrypte déchiffre une donnée d’entrée (avec une clé aléatoire) selon
// l’algorithme demandé. Les algorithmes supportés sont :
// - blowfish
// - aes
func Decrypt[T ~[]byte | ~string](encrypted, key T, algorithm string) Result[T] {
	c, ok := encoders[algorithm]
	if !ok {
		return Err[T](errAlgorithm(algorithm))
	}

	block, err := c([]byte(key))
	if err != nil {
		return Err[T](err)
	}

	decoded := decodeFromBase64(encrypted)
	if err, ok := decoded.Err(); ok {
		return Err[T](err)
	}

	dec, _ := decoded.Ok()
	bs := blockSize[algorithm]

	if (len(dec) % bs) != 0 {
		return Err[T](errors.New("blocksize must be multiple of data length"))
	}

	iv := dec[:bs]
	msg := dec[bs:]

	cfb := cipher.NewCFBDecrypter(block, iv)
	cfb.XORKeyStream(msg, msg)

	unpadMsg := unpad(msg)
	if err, ok := unpadMsg.Err(); ok {
		return Err[T](err)
	}

	result, _ := unpadMsg.Ok()
	return Ok(T(result))
}