package ymt_v3

import (
	"bytes"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/tjfoc/gmsm/sm4"
)

// EncryptType 加密类型
type EncryptType string

const (
	// EncryptTypeAES AES-ECB加密
	EncryptTypeAES EncryptType = "aes"
	// EncryptTypeSM4 SM4-CBC加密
	EncryptTypeSM4 EncryptType = "sm4"
)

// ==================== RSA签名 ====================

// SignWithRSA 使用RSA私钥对数据进行签名
func SignWithRSA(signStr string, privateKeyPEM string) (string, error) {
	block, _ := pem.Decode([]byte(privateKeyPEM))
	if block == nil {
		return "", fmt.Errorf("failed to decode PEM private key")
	}

	privateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return "", fmt.Errorf("failed to parse private key: %v", err)
		}
	}

	rsaPrivateKey, ok := privateKey.(*rsa.PrivateKey)
	if !ok {
		return "", fmt.Errorf("not a RSA private key")
	}

	h := sha256.New()
	h.Write([]byte(signStr))
	hashed := h.Sum(nil)

	signature, err := rsa.SignPKCS1v15(rand.Reader, rsaPrivateKey, crypto.SHA256, hashed)
	if err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(signature), nil
}

// VerifyWithRSA 使用RSA公钥验证签名
func VerifyWithRSA(signStr, signature string, publicKeyPEM string) error {
	block, _ := pem.Decode([]byte(publicKeyPEM))
	if block == nil {
		return fmt.Errorf("failed to decode PEM public key")
	}

	publicKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return fmt.Errorf("failed to parse public key: %v", err)
	}

	rsaPublicKey, ok := publicKey.(*rsa.PublicKey)
	if !ok {
		return fmt.Errorf("not a RSA public key")
	}

	h := sha256.New()
	h.Write([]byte(signStr))
	hashed := h.Sum(nil)

	sigBytes, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return fmt.Errorf("failed to decode signature: %v", err)
	}

	return rsa.VerifyPKCS1v15(rsaPublicKey, crypto.SHA256, hashed, sigBytes)
}

// ==================== AES ECB加密 ====================

// AESECBEncrypt AES-ECB模式加密
func AESECBEncrypt(plaintext []byte, key []byte) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("创建AES cipher失败: %v", err)
	}

	padded := pkcs7Padding(plaintext, aes.BlockSize)

	ciphertext := make([]byte, len(padded))
	for i := 0; i < len(padded); i += aes.BlockSize {
		block.Encrypt(ciphertext[i:i+aes.BlockSize], padded[i:i+aes.BlockSize])
	}

	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// AESECBDecrypt AES-ECB模式解密
func AESECBDecrypt(encryptedData string, key []byte) ([]byte, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(encryptedData)
	if err != nil {
		return nil, fmt.Errorf("解码失败: %v", err)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("创建AES cipher失败: %v", err)
	}

	if len(ciphertext)%aes.BlockSize != 0 {
		return nil, fmt.Errorf("密文长度不是块大小的整数倍")
	}

	plaintext := make([]byte, len(ciphertext))
	for i := 0; i < len(ciphertext); i += aes.BlockSize {
		block.Decrypt(plaintext[i:i+aes.BlockSize], ciphertext[i:i+aes.BlockSize])
	}

	plaintext, err = pkcs7UnPadding(plaintext)
	if err != nil {
		return nil, fmt.Errorf("去除填充失败: %v", err)
	}

	return plaintext, nil
}

// ==================== SM4 CBC加密 ====================

// SM4CBCEncrypt SM4-CBC模式加密
// IV会前置到密文中，解密时自动提取
func SM4CBCEncrypt(plaintext []byte, key []byte) (string, error) {
	if len(key) != 16 {
		return "", fmt.Errorf("SM4密钥长度必须为16字节")
	}

	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("创建SM4 cipher失败: %v", err)
	}

	padded := pkcs7Padding(plaintext, block.BlockSize())

	iv := make([]byte, block.BlockSize())
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", fmt.Errorf("生成IV失败: %v", err)
	}

	mode := cipher.NewCBCEncrypter(block, iv)
	ciphertext := make([]byte, len(padded))
	mode.CryptBlocks(ciphertext, padded)

	// IV前置到密文
	result := append(iv, ciphertext...)
	return base64.StdEncoding.EncodeToString(result), nil
}

// SM4CBCDecrypt SM4-CBC模式解密
// 从密文中提取前置的IV进行解密
func SM4CBCDecrypt(encryptedData string, key []byte) ([]byte, error) {
	if len(key) != 16 {
		return nil, fmt.Errorf("SM4密钥长度必须为16字节")
	}

	data, err := base64.StdEncoding.DecodeString(encryptedData)
	if err != nil {
		return nil, fmt.Errorf("解码失败: %v", err)
	}

	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("创建SM4 cipher失败: %v", err)
	}

	blockSize := block.BlockSize()
	if len(data) < blockSize*2 {
		return nil, fmt.Errorf("数据长度不足")
	}
	iv := data[:blockSize]
	ciphertext := data[blockSize:]

	mode := cipher.NewCBCDecrypter(block, iv)
	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)

	plaintext, err = pkcs7UnPadding(plaintext)
	if err != nil {
		return nil, fmt.Errorf("去除填充失败: %v", err)
	}

	return plaintext, nil
}

// ==================== PKCS7填充 ====================

// pkcs7Padding PKCS7填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// pkcs7UnPadding 去除PKCS7填充
func pkcs7UnPadding(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, fmt.Errorf("数据为空")
	}
	padding := int(data[length-1])
	if padding > length || padding == 0 {
		return nil, fmt.Errorf("无效的填充")
	}
	for i := length - padding; i < length; i++ {
		if data[i] != byte(padding) {
			return nil, fmt.Errorf("无效的填充")
		}
	}
	return data[:length-padding], nil
}

// ==================== 业务参数加密/解密 ====================

// EncryptBizParams 加密业务参数
// 1. 将业务参数去掉"零"值的参数，再由小到大按照字母排序，转成json字符串得到plaintext
// 2. 使用应用key将plaintext字符串加密，得到ciphertext
func EncryptBizParams(params interface{}, key []byte, encType EncryptType) (string, error) {
	plaintext, err := marshalWithoutZero(params)
	if err != nil {
		return "", fmt.Errorf("序列化业务参数失败: %v", err)
	}

	switch encType {
	case EncryptTypeAES:
		return AESECBEncrypt([]byte(plaintext), key)
	case EncryptTypeSM4:
		return SM4CBCEncrypt([]byte(plaintext), key)
	default:
		return "", fmt.Errorf("不支持的加密类型: %s", encType)
	}
}

// DecryptBizParams 解密业务参数
func DecryptBizParams(ciphertext string, key []byte, encType EncryptType) ([]byte, error) {
	switch encType {
	case EncryptTypeAES:
		return AESECBDecrypt(ciphertext, key)
	case EncryptTypeSM4:
		return SM4CBCDecrypt(ciphertext, key)
	default:
		return nil, fmt.Errorf("不支持的加密类型: %s", encType)
	}
}

// EncryptBizParamsRaw 直接加密字节数据（用于回调验签）
func EncryptBizParamsRaw(plaintext []byte, key []byte, encType EncryptType) (string, error) {
	switch encType {
	case EncryptTypeAES:
		return AESECBEncrypt(plaintext, key)
	case EncryptTypeSM4:
		return SM4CBCEncrypt(plaintext, key)
	default:
		return "", fmt.Errorf("不支持的加密类型: %s", encType)
	}
}

// marshalWithoutZero 将结构体序列化为JSON，去掉零值字段并按字母排序
func marshalWithoutZero(v interface{}) (string, error) {
	val := reflect.ValueOf(v)
	if val.Kind() == reflect.Ptr {
		val = val.Elem()
	}

	if val.Kind() != reflect.Struct {
		if val.Kind() == reflect.Map {
			// 类型断言为 map[string]interface{}
			m, ok := v.(map[string]interface{})
			if !ok {
				// 尝试从 reflect 转换
				m = make(map[string]interface{})
				for _, key := range val.MapKeys() {
					m[fmt.Sprintf("%v", key.Interface())] = val.MapIndex(key).Interface()
				}
			}
			return marshalMapSorted(m)
		}
		data, err := json.Marshal(v)
		return string(data), err
	}

	result := make(map[string]interface{})
	typ := val.Type()

	for i := 0; i < val.NumField(); i++ {
		field := val.Field(i)
		fieldType := typ.Field(i)

		jsonTag := fieldType.Tag.Get("json")
		if jsonTag == "" || jsonTag == "-" {
			continue
		}
		name := strings.Split(jsonTag, ",")[0]

		if isZeroValue(field) {
			continue
		}

		result[name] = field.Interface()
	}

	return marshalMapSorted(result)
}

// marshalMapSorted 将map按键排序后序列化为JSON
func marshalMapSorted(m map[string]interface{}) (string, error) {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var buf bytes.Buffer
	buf.WriteByte('{')
	for i, k := range keys {
		if i > 0 {
			buf.WriteByte(',')
		}
		keyBytes, _ := json.Marshal(k)
		buf.Write(keyBytes)
		buf.WriteByte(':')
		valBytes, err := json.Marshal(m[k])
		if err != nil {
			return "", err
		}
		buf.Write(valBytes)
	}
	buf.WriteByte('}')

	return buf.String(), nil
}

// isZeroValue 判断reflect.Value是否为零值
func isZeroValue(v reflect.Value) bool {
	switch v.Kind() {
	case reflect.String:
		return v.String() == ""
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return v.Int() == 0
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return v.Uint() == 0
	case reflect.Float32, reflect.Float64:
		return v.Float() == 0
	case reflect.Bool:
		return !v.Bool()
	case reflect.Slice, reflect.Map, reflect.Ptr:
		return v.IsNil()
	default:
		return false
	}
}

// ==================== 签名相关 ====================

// BuildSignStr 构建签名字符串
// 规则：app_id + timestamp + ciphertext
func BuildSignStr(appID, timestamp, ciphertext string) string {
	return appID + timestamp + ciphertext
}

// GenerateTimestamp 生成时间戳，格式 yyyy-MM-dd HH:mm:ss
func GenerateTimestamp() string {
	return time.Now().Format("2006-01-02 15:04:05")
}

// ==================== 回调验签 ====================

// VerifyNotifySign 验证回调通知签名
func VerifyNotifySign(header *NotifyHeader, data interface{}, key []byte, publicKeyPEM string, encType EncryptType) error {
	plaintext, err := marshalWithoutZero(data)
	if err != nil {
		return fmt.Errorf("序列化回调数据失败: %v", err)
	}

	ciphertext, err := EncryptBizParamsRaw([]byte(plaintext), key, encType)
	if err != nil {
		return fmt.Errorf("加密回调数据失败: %v", err)
	}

	signStr := BuildSignStr(header.Appid, header.Timestamp, ciphertext)

	return VerifyWithRSA(signStr, header.Sign, publicKeyPEM)
}