package functions

import (
	"crypto/hmac"
	"crypto/md5"
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"math/big"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
	"github.com/google/uuid"
)

func GetCryptoFunctions() map[string]luajit.GoFunction {
	return map[string]luajit.GoFunction{
		"base64_encode": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("base64_encode: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("base64_encode: argument must be a string")
				return 2
			}
			encoded := base64.StdEncoding.EncodeToString([]byte(str))
			s.PushString(encoded)
			return 1
		},

		"base64_decode": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("base64_decode: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("base64_decode: argument must be a string")
				return 2
			}
			decoded, err := base64.StdEncoding.DecodeString(str)
			if err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("base64_decode: %v", err))
				return 2
			}
			s.PushString(string(decoded))
			return 1
		},

		"base64_url_encode": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("base64_url_encode: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("base64_url_encode: argument must be a string")
				return 2
			}
			encoded := base64.URLEncoding.EncodeToString([]byte(str))
			s.PushString(encoded)
			return 1
		},

		"base64_url_decode": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("base64_url_decode: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("base64_url_decode: argument must be a string")
				return 2
			}
			decoded, err := base64.URLEncoding.DecodeString(str)
			if err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("base64_url_decode: %v", err))
				return 2
			}
			s.PushString(string(decoded))
			return 1
		},

		"hex_encode": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("hex_encode: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("hex_encode: argument must be a string")
				return 2
			}
			encoded := hex.EncodeToString([]byte(str))
			s.PushString(encoded)
			return 1
		},

		"hex_decode": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("hex_decode: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("hex_decode: argument must be a string")
				return 2
			}
			decoded, err := hex.DecodeString(str)
			if err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("hex_decode: %v", err))
				return 2
			}
			s.PushString(string(decoded))
			return 1
		},

		"md5_hash": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("md5_hash: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("md5_hash: argument must be a string")
				return 2
			}
			hash := md5.Sum([]byte(str))
			s.PushString(hex.EncodeToString(hash[:]))
			return 1
		},

		"sha1_hash": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("sha1_hash: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("sha1_hash: argument must be a string")
				return 2
			}
			hash := sha1.Sum([]byte(str))
			s.PushString(hex.EncodeToString(hash[:]))
			return 1
		},

		"sha256_hash": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("sha256_hash: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("sha256_hash: argument must be a string")
				return 2
			}
			hash := sha256.Sum256([]byte(str))
			s.PushString(hex.EncodeToString(hash[:]))
			return 1
		},

		"sha512_hash": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("sha512_hash: %v", err))
				return 2
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("sha512_hash: argument must be a string")
				return 2
			}
			hash := sha512.Sum512([]byte(str))
			s.PushString(hex.EncodeToString(hash[:]))
			return 1
		},

		"hmac_sha256": func(s *luajit.State) int {
			if err := s.CheckExactArgs(2); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("hmac_sha256: %v", err))
				return 2
			}
			message, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("hmac_sha256: first argument must be a string")
				return 2
			}
			key, err := s.SafeToString(2)
			if err != nil {
				s.PushNil()
				s.PushString("hmac_sha256: second argument must be a string")
				return 2
			}
			h := hmac.New(sha256.New, []byte(key))
			h.Write([]byte(message))
			s.PushString(hex.EncodeToString(h.Sum(nil)))
			return 1
		},

		"hmac_sha1": func(s *luajit.State) int {
			if err := s.CheckExactArgs(2); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("hmac_sha1: %v", err))
				return 2
			}
			message, err := s.SafeToString(1)
			if err != nil {
				s.PushNil()
				s.PushString("hmac_sha1: first argument must be a string")
				return 2
			}
			key, err := s.SafeToString(2)
			if err != nil {
				s.PushNil()
				s.PushString("hmac_sha1: second argument must be a string")
				return 2
			}
			h := hmac.New(sha1.New, []byte(key))
			h.Write([]byte(message))
			s.PushString(hex.EncodeToString(h.Sum(nil)))
			return 1
		},

		"uuid_generate": func(s *luajit.State) int {
			id := uuid.New()
			s.PushString(id.String())
			return 1
		},

		"uuid_generate_v4": func(s *luajit.State) int {
			id := uuid.New()
			s.PushString(id.String())
			return 1
		},

		"uuid_validate": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushBoolean(false)
				return 1
			}
			str, err := s.SafeToString(1)
			if err != nil {
				s.PushBoolean(false)
				return 1
			}
			_, err = uuid.Parse(str)
			s.PushBoolean(err == nil)
			return 1
		},

		"random_bytes": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("random_bytes: %v", err))
				return 2
			}
			length, err := s.SafeToNumber(1)
			if err != nil || length < 0 || length != float64(int(length)) {
				s.PushNil()
				s.PushString("random_bytes: argument must be a non-negative integer")
				return 2
			}
			if length > 65536 {
				s.PushNil()
				s.PushString("random_bytes: length too large (max 65536)")
				return 2
			}
			bytes := make([]byte, int(length))
			if _, err := rand.Read(bytes); err != nil {
				s.PushNil()
				s.PushString("random_bytes: failed to generate random bytes")
				return 2
			}
			s.PushString(string(bytes))
			return 1
		},

		"random_hex": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("random_hex: %v", err))
				return 2
			}
			length, err := s.SafeToNumber(1)
			if err != nil || length < 0 || length != float64(int(length)) {
				s.PushNil()
				s.PushString("random_hex: argument must be a non-negative integer")
				return 2
			}
			if length > 32768 {
				s.PushNil()
				s.PushString("random_hex: length too large (max 32768)")
				return 2
			}
			bytes := make([]byte, int(length))
			if _, err := rand.Read(bytes); err != nil {
				s.PushNil()
				s.PushString("random_hex: failed to generate random bytes")
				return 2
			}
			s.PushString(hex.EncodeToString(bytes))
			return 1
		},

		"random_string": func(s *luajit.State) int {
			if err := s.CheckMinArgs(1); err != nil {
				s.PushNil()
				s.PushString(fmt.Sprintf("random_string: %v", err))
				return 2
			}
			length, err := s.SafeToNumber(1)
			if err != nil || length < 0 || length != float64(int(length)) {
				s.PushNil()
				s.PushString("random_string: argument must be a non-negative integer")
				return 2
			}
			if length > 65536 {
				s.PushNil()
				s.PushString("random_string: length too large (max 65536)")
				return 2
			}

			charset := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
			if s.GetTop() >= 2 {
				if customCharset, err := s.SafeToString(2); err == nil {
					charset = customCharset
				}
			}

			result := make([]byte, int(length))
			charsetLen := big.NewInt(int64(len(charset)))
			for i := range result {
				n, err := rand.Int(rand.Reader, charsetLen)
				if err != nil {
					s.PushNil()
					s.PushString("random_string: failed to generate random number")
					return 2
				}
				result[i] = charset[n.Int64()]
			}
			s.PushString(string(result))
			return 1
		},

		"secure_compare": func(s *luajit.State) int {
			if err := s.CheckExactArgs(2); err != nil {
				s.PushBoolean(false)
				return 1
			}
			a, err := s.SafeToString(1)
			if err != nil {
				s.PushBoolean(false)
				return 1
			}
			b, err := s.SafeToString(2)
			if err != nil {
				s.PushBoolean(false)
				return 1
			}
			s.PushBoolean(hmac.Equal([]byte(a), []byte(b)))
			return 1
		},
	}
}