package crypto

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

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
	"github.com/google/uuid"
	"golang.org/x/crypto/argon2"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/crypto/pbkdf2"
	"golang.org/x/crypto/scrypt"
)

func GetFunctionList() map[string]luajit.GoFunction {
	return map[string]luajit.GoFunction{
		"base64_encode":     base64_encode,
		"base64_decode":     base64_decode,
		"base64_url_encode": base64_url_encode,
		"base64_url_decode": base64_url_decode,
		"hex_encode":        hex_encode,
		"hex_decode":        hex_decode,
		"md5_hash":          md5_hash,
		"sha1_hash":         sha1_hash,
		"sha256_hash":       sha256_hash,
		"sha512_hash":       sha512_hash,
		"hmac_sha256":       hmac_sha256,
		"hmac_sha1":         hmac_sha1,
		"uuid_generate":     uuid_generate,
		"uuid_generate_v4":  uuid_generate_v4,
		"uuid_validate":     uuid_validate,
		"random_bytes":      random_bytes,
		"random_hex":        random_hex,
		"random_string":     random_string,
		"secure_compare":    secure_compare,
		"argon2_hash":       argon2_hash,
		"argon2_verify":     argon2_verify,
		"bcrypt_hash":       bcrypt_hash,
		"bcrypt_verify":     bcrypt_verify,
		"scrypt_hash":       scrypt_hash,
		"scrypt_verify":     scrypt_verify,
		"pbkdf2_hash":       pbkdf2_hash,
		"pbkdf2_verify":     pbkdf2_verify,
		"password_hash":     password_hash,
		"password_verify":   password_verify,
	}
}

func base64_encode(s *luajit.State) int {
	str := s.ToString(1)
	encoded := base64.StdEncoding.EncodeToString([]byte(str))
	s.PushString(encoded)
	return 1
}

func base64_decode(s *luajit.State) int {
	str := s.ToString(1)
	decoded, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		s.PushNil()
		s.PushString("invalid base64 data")
		return 2
	}
	s.PushString(string(decoded))
	return 1
}

func base64_url_encode(s *luajit.State) int {
	str := s.ToString(1)
	encoded := base64.URLEncoding.EncodeToString([]byte(str))
	s.PushString(encoded)
	return 1
}

func base64_url_decode(s *luajit.State) int {
	str := s.ToString(1)
	decoded, err := base64.URLEncoding.DecodeString(str)
	if err != nil {
		s.PushNil()
		s.PushString("invalid base64url data")
		return 2
	}
	s.PushString(string(decoded))
	return 1
}

func hex_encode(s *luajit.State) int {
	str := s.ToString(1)
	encoded := hex.EncodeToString([]byte(str))
	s.PushString(encoded)
	return 1
}

func hex_decode(s *luajit.State) int {
	str := s.ToString(1)
	decoded, err := hex.DecodeString(str)
	if err != nil {
		s.PushNil()
		s.PushString("invalid hex data")
		return 2
	}
	s.PushString(string(decoded))
	return 1
}

func md5_hash(s *luajit.State) int {
	str := s.ToString(1)
	hash := md5.Sum([]byte(str))
	s.PushString(hex.EncodeToString(hash[:]))
	return 1
}

func sha1_hash(s *luajit.State) int {
	str := s.ToString(1)
	hash := sha1.Sum([]byte(str))
	s.PushString(hex.EncodeToString(hash[:]))
	return 1
}

func sha256_hash(s *luajit.State) int {
	str := s.ToString(1)
	hash := sha256.Sum256([]byte(str))
	s.PushString(hex.EncodeToString(hash[:]))
	return 1
}

func sha512_hash(s *luajit.State) int {
	str := s.ToString(1)
	hash := sha512.Sum512([]byte(str))
	s.PushString(hex.EncodeToString(hash[:]))
	return 1
}

func hmac_sha256(s *luajit.State) int {
	message := s.ToString(1)
	key := s.ToString(2)
	h := hmac.New(sha256.New, []byte(key))
	h.Write([]byte(message))
	s.PushString(hex.EncodeToString(h.Sum(nil)))
	return 1
}

func hmac_sha1(s *luajit.State) int {
	message := s.ToString(1)
	key := s.ToString(2)
	h := hmac.New(sha1.New, []byte(key))
	h.Write([]byte(message))
	s.PushString(hex.EncodeToString(h.Sum(nil)))
	return 1
}

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

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

func uuid_validate(s *luajit.State) int {
	str := s.ToString(1)
	_, err := uuid.Parse(str)
	s.PushBoolean(err == nil)
	return 1
}

func random_bytes(s *luajit.State) int {
	length := int(s.ToNumber(1))
	if length < 0 || length > 65536 {
		s.PushNil()
		s.PushString("invalid length")
		return 2
	}
	bytes := make([]byte, length)
	if _, err := rand.Read(bytes); err != nil {
		s.PushNil()
		s.PushString("failed to generate random bytes")
		return 2
	}
	s.PushString(string(bytes))
	return 1
}

func random_hex(s *luajit.State) int {
	length := int(s.ToNumber(1))
	if length < 0 || length > 32768 {
		s.PushNil()
		s.PushString("invalid length")
		return 2
	}
	bytes := make([]byte, length)
	if _, err := rand.Read(bytes); err != nil {
		s.PushNil()
		s.PushString("failed to generate random bytes")
		return 2
	}
	s.PushString(hex.EncodeToString(bytes))
	return 1
}

func random_string(s *luajit.State) int {
	length := int(s.ToNumber(1))
	if length < 0 || length > 65536 {
		s.PushNil()
		s.PushString("invalid length")
		return 2
	}

	charset := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	if s.GetTop() >= 2 && !s.IsNil(2) {
		charset = s.ToString(2)
	}

	if len(charset) == 0 {
		s.PushNil()
		s.PushString("empty charset")
		return 2
	}

	result := make([]byte, 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("failed to generate random number")
			return 2
		}
		result[i] = charset[n.Int64()]
	}
	s.PushString(string(result))
	return 1
}

func secure_compare(s *luajit.State) int {
	a := s.ToString(1)
	b := s.ToString(2)
	s.PushBoolean(hmac.Equal([]byte(a), []byte(b)))
	return 1
}

func argon2_hash(s *luajit.State) int {
	password := s.ToString(1)
	time := uint32(1)
	memory := uint32(64 * 1024)
	threads := uint8(4)
	keyLen := uint32(32)

	if s.GetTop() >= 2 && !s.IsNil(2) {
		time = uint32(s.ToNumber(2))
	}
	if s.GetTop() >= 3 && !s.IsNil(3) {
		memory = uint32(s.ToNumber(3))
	}
	if s.GetTop() >= 4 && !s.IsNil(4) {
		threads = uint8(s.ToNumber(4))
	}
	if s.GetTop() >= 5 && !s.IsNil(5) {
		keyLen = uint32(s.ToNumber(5))
	}

	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		s.PushNil()
		s.PushString("failed to generate salt")
		return 2
	}

	hash := argon2.IDKey([]byte(password), salt, time, memory, threads, keyLen)
	encodedSalt := base64.RawStdEncoding.EncodeToString(salt)
	encodedHash := base64.RawStdEncoding.EncodeToString(hash)

	result := fmt.Sprintf("$argon2id$v=19$m=%d,t=%d,p=%d$%s$%s",
		memory, time, threads, encodedSalt, encodedHash)

	s.PushString(result)
	return 1
}

func argon2_verify(s *luajit.State) int {
	password := s.ToString(1)
	hash := s.ToString(2)

	parts := strings.Split(hash, "$")
	if len(parts) != 6 || parts[1] != "argon2id" {
		s.PushBoolean(false)
		return 1
	}

	var memory, time uint32
	var threads uint8
	if _, err := fmt.Sscanf(parts[3], "m=%d,t=%d,p=%d", &memory, &time, &threads); err != nil {
		s.PushBoolean(false)
		return 1
	}

	salt, err := base64.RawStdEncoding.DecodeString(parts[4])
	if err != nil {
		s.PushBoolean(false)
		return 1
	}

	expectedHash, err := base64.RawStdEncoding.DecodeString(parts[5])
	if err != nil {
		s.PushBoolean(false)
		return 1
	}

	actualHash := argon2.IDKey([]byte(password), salt, time, memory, threads, uint32(len(expectedHash)))
	s.PushBoolean(hmac.Equal(actualHash, expectedHash))
	return 1
}

func bcrypt_hash(s *luajit.State) int {
	password := s.ToString(1)
	cost := 12

	if s.GetTop() >= 2 && !s.IsNil(2) {
		cost = int(s.ToNumber(2))
		if cost < 4 || cost > 31 {
			s.PushNil()
			s.PushString("invalid cost (must be 4-31)")
			return 2
		}
	}

	hash, err := bcrypt.GenerateFromPassword([]byte(password), cost)
	if err != nil {
		s.PushNil()
		s.PushString("bcrypt hash failed")
		return 2
	}

	s.PushString(string(hash))
	return 1
}

func bcrypt_verify(s *luajit.State) int {
	password := s.ToString(1)
	hash := s.ToString(2)

	err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(password))
	s.PushBoolean(err == nil)
	return 1
}

func scrypt_hash(s *luajit.State) int {
	password := s.ToString(1)
	N := 32768   // CPU cost
	r := 8       // block size
	p := 1       // parallelization
	keyLen := 32 // key length

	if s.GetTop() >= 2 && !s.IsNil(2) {
		N = int(s.ToNumber(2))
	}
	if s.GetTop() >= 3 && !s.IsNil(3) {
		r = int(s.ToNumber(3))
	}
	if s.GetTop() >= 4 && !s.IsNil(4) {
		p = int(s.ToNumber(4))
	}
	if s.GetTop() >= 5 && !s.IsNil(5) {
		keyLen = int(s.ToNumber(5))
	}

	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		s.PushNil()
		s.PushString("failed to generate salt")
		return 2
	}

	hash, err := scrypt.Key([]byte(password), salt, N, r, p, keyLen)
	if err != nil {
		s.PushNil()
		s.PushString("scrypt hash failed")
		return 2
	}

	encodedSalt := base64.RawStdEncoding.EncodeToString(salt)
	encodedHash := base64.RawStdEncoding.EncodeToString(hash)

	result := fmt.Sprintf("$scrypt$N=%d,r=%d,p=%d$%s$%s", N, r, p, encodedSalt, encodedHash)
	s.PushString(result)
	return 1
}

func scrypt_verify(s *luajit.State) int {
	password := s.ToString(1)
	hash := s.ToString(2)

	parts := strings.Split(hash, "$")
	if len(parts) != 5 || parts[1] != "scrypt" {
		s.PushBoolean(false)
		return 1
	}

	var N, r, p int
	if _, err := fmt.Sscanf(parts[2], "N=%d,r=%d,p=%d", &N, &r, &p); err != nil {
		s.PushBoolean(false)
		return 1
	}

	salt, err := base64.RawStdEncoding.DecodeString(parts[3])
	if err != nil {
		s.PushBoolean(false)
		return 1
	}

	expectedHash, err := base64.RawStdEncoding.DecodeString(parts[4])
	if err != nil {
		s.PushBoolean(false)
		return 1
	}

	actualHash, err := scrypt.Key([]byte(password), salt, N, r, p, len(expectedHash))
	if err != nil {
		s.PushBoolean(false)
		return 1
	}

	s.PushBoolean(hmac.Equal(actualHash, expectedHash))
	return 1
}

func pbkdf2_hash(s *luajit.State) int {
	password := s.ToString(1)
	iterations := 100000
	keyLen := 32

	if s.GetTop() >= 2 && !s.IsNil(2) {
		iterations = int(s.ToNumber(2))
	}
	if s.GetTop() >= 3 && !s.IsNil(3) {
		keyLen = int(s.ToNumber(3))
	}

	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		s.PushNil()
		s.PushString("failed to generate salt")
		return 2
	}

	hash := pbkdf2.Key([]byte(password), salt, iterations, keyLen, sha256.New)
	encodedSalt := base64.RawStdEncoding.EncodeToString(salt)
	encodedHash := base64.RawStdEncoding.EncodeToString(hash)

	result := fmt.Sprintf("$pbkdf2-sha256$i=%d$%s$%s", iterations, encodedSalt, encodedHash)
	s.PushString(result)
	return 1
}

func pbkdf2_verify(s *luajit.State) int {
	password := s.ToString(1)
	hash := s.ToString(2)

	parts := strings.Split(hash, "$")
	if len(parts) != 5 || parts[1] != "pbkdf2-sha256" {
		s.PushBoolean(false)
		return 1
	}

	var iterations int
	if _, err := fmt.Sscanf(parts[2], "i=%d", &iterations); err != nil {
		s.PushBoolean(false)
		return 1
	}

	salt, err := base64.RawStdEncoding.DecodeString(parts[3])
	if err != nil {
		s.PushBoolean(false)
		return 1
	}

	expectedHash, err := base64.RawStdEncoding.DecodeString(parts[4])
	if err != nil {
		s.PushBoolean(false)
		return 1
	}

	actualHash := pbkdf2.Key([]byte(password), salt, iterations, len(expectedHash), sha256.New)
	s.PushBoolean(hmac.Equal(actualHash, expectedHash))
	return 1
}

func password_hash(s *luajit.State) int {
	password := s.ToString(1)
	algorithm := "argon2id" // default

	if s.GetTop() >= 2 && !s.IsNil(2) {
		algorithm = s.ToString(2)
	}

	switch algorithm {
	case "argon2id":
		s.PushString(password)
		return argon2_hash(s)
	case "bcrypt":
		s.PushString(password)
		if s.GetTop() >= 3 {
			s.PushNumber(s.ToNumber(3))
		}
		return bcrypt_hash(s)
	case "scrypt":
		s.PushString(password)
		return scrypt_hash(s)
	case "pbkdf2":
		s.PushString(password)
		return pbkdf2_hash(s)
	default:
		s.PushNil()
		s.PushString("unsupported algorithm: " + algorithm)
		return 2
	}
}

func password_verify(s *luajit.State) int {
	hash := s.ToString(2)

	// Auto-detect algorithm from hash format
	if strings.HasPrefix(hash, "$argon2id$") {
		return argon2_verify(s)
	} else if strings.HasPrefix(hash, "$2a$") || strings.HasPrefix(hash, "$2b$") || strings.HasPrefix(hash, "$2y$") {
		return bcrypt_verify(s)
	} else if strings.HasPrefix(hash, "$scrypt$") {
		return scrypt_verify(s)
	} else if strings.HasPrefix(hash, "$pbkdf2-sha256$") {
		return pbkdf2_verify(s)
	}

	s.PushBoolean(false)
	return 1
}