package utils

import (
	"bytes"
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"log"

	"github.com/go-redis/redis/v8"
	"golang.org/x/crypto/pbkdf2"
)

var ctx = context.Background()

// 配置信息
const (
	redisAddr     = "110.41.7.101:6388"
	redisPassword = "Test@Device@2024"
	redisDB       = 9
)

// 初始化Redis客户端
func initRedisClient() *redis.Client {
	client := redis.NewClient(&redis.Options{
		Addr:     redisAddr,
		Password: redisPassword,
		DB:       redisDB,
	})

	// 测试连接
	_, err := client.Ping(ctx).Result()
	if err != nil {
		log.Fatalf("Failed to connect to Redis: %v", err)
	}

	return client
}

// GetEncryptMode 获取加密方式
func GetEncryptMode(deviceID string) (string, error) {
	// 初始化Redis客户端
	client := initRedisClient()
	defer client.Close()
	encryptKey := fmt.Sprintf("iot%s", deviceID)
	encrypt, err := client.HGet(ctx, encryptKey, "encryption").Result()
	if err != nil {
		if err == redis.Nil {
			return "plaintext", nil
		}
		// 其他错误返回错误信息
		return "plaintext", fmt.Errorf("failed to get encryption mode from Redis: %w", err)
	}
	return encrypt, nil
}

// 获取远程控制所需的加密参数
func getRemoteControlParams(client *redis.Client, deviceID string) ([]byte, []byte, []byte, error) {
	// 从Redis获取sn1、sn2和PSK
	sn1Key := fmt.Sprintf("iot%s", deviceID)
	sn2Key := fmt.Sprintf("iot%s", deviceID)
	pskKey := fmt.Sprintf("iot%s", deviceID)

	sn1Hex, err := client.HGet(ctx, sn1Key, "sn1").Result()
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to get SN1 from Redis: %v", err)
	}

	sn2Hex, err := client.HGet(ctx, sn2Key, "sn2").Result()
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to get SN2 from Redis: %v", err)
	}

	pskHex, err := client.HGet(ctx, pskKey, "psk").Result()
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to get PSK from Redis: %v", err)
	}

	// 将sn1和sn2从16进制字符串转换为字节数组
	sn1, err := hex.DecodeString(sn1Hex)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("sn1 解码失败: %v", err)
	}

	sn2, err := hex.DecodeString(sn2Hex)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("sn2 解码失败: %v", err)
	}

	psk, err := hex.DecodeString(pskHex)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("PSK 解码失败: %v", err)
	}
	// 验证长度
	if len(sn1) != 8 || len(sn2) != 8 {
		return nil, nil, nil, fmt.Errorf("SN1/SN2 must be 8 bytes after hex decoding")
	}

	return sn1, sn2, psk, nil
}

// 生成PBKDF2密钥
func deriveKey(password, salt []byte, iterations, keyLen int) []byte {
	return pbkdf2.Key(password, salt, iterations, keyLen, sha256.New)
}

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

// PKCS5去填充
func pkcs5Unpadding(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("data is empty")
	}

	paddingLen := int(data[len(data)-1])
	if paddingLen > len(data) || paddingLen == 0 {
		return nil, fmt.Errorf("invalid padding length")
	}

	for i := len(data) - paddingLen; i < len(data); i++ {
		if data[i] != byte(paddingLen) {
			return nil, fmt.Errorf("invalid padding bytes")
		}
	}

	return data[:len(data)-paddingLen], nil
}

// AES-128-CBC加密
func encryptAES128CBC(plaintext, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 填充数据
	paddedData := pkcs5Padding(plaintext, aes.BlockSize)

	// 创建CBC模式
	mode := cipher.NewCBCEncrypter(block, iv)

	// 加密
	ciphertext := make([]byte, len(paddedData))
	mode.CryptBlocks(ciphertext, paddedData)

	return ciphertext, nil
}

// AES-128-CBC解密
func decryptAES128CBC(ciphertext, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 创建CBC模式
	mode := cipher.NewCBCDecrypter(block, iv)

	// 解密
	decrypted := make([]byte, len(ciphertext))
	mode.CryptBlocks(decrypted, ciphertext)

	// 去填充
	return pkcs5Unpadding(decrypted)
}

// 计算HMAC-SHA256
func calculateHMAC(data, secret []byte) []byte {
	h := hmac.New(sha256.New, secret)
	h.Write(data)
	return h.Sum(nil)
}

// 构建CoAP消息（包含HMAC）
func buildCoAPMessage(coapHeader, payload, hmac []byte) []byte {
	// 注意：如果是CoAP over TCP，需要调整coapHeader中的长度字段以包含HMAC
	return append(append(coapHeader, payload...), hmac...)
}

// EncryptBodyAES 数据加密,只加密body信息
func EncryptBodyAES(devId string, payload []byte) []byte {
	// 初始化Redis客户端
	client := initRedisClient()
	defer client.Close()
	// 获取加密参数
	sn1, sn2, psk, err := getRemoteControlParams(client, devId)
	if err != nil {
		log.Fatalf("Failed to get encryption parameters: %v", err)
	}
	// 拼接盐值
	salt := append(sn1, sn2...)

	// 生成PBKDF2密钥（32字节）
	digest := deriveKey(psk, salt, 1, 32)
	// 提取AES密钥和IV
	aesKey := digest[:16]
	iv := digest[16:32]
	ciphertext, err := encryptAES128CBC(payload, aesKey, iv)
	if err != nil {
		log.Fatalf("Encryption error: %v", err)
	}

	return ciphertext
}

// EncryptDataAESMac 数据进行加密,返回最终数据
func EncryptDataAESMac(devId string, payload []byte, coapHeader []byte) []byte {
	// 初始化Redis客户端
	client := initRedisClient()
	defer client.Close()
	// 获取加密参数
	sn1, sn2, psk, err := getRemoteControlParams(client, devId)
	if err != nil {
		log.Fatalf("Failed to get encryption parameters: %v", err)
	}
	// 拼接盐值
	salt := append(sn1, sn2...)
	// 生成PBKDF2密钥（32字节）
	digest := deriveKey(psk, salt, 1, 32)
	// 提取AES密钥和IV
	aesKey := digest[:16]
	iv := digest[16:32]
	ciphertext, err := encryptAES128CBC(payload, aesKey, iv)
	if err != nil {
		log.Fatalf("Encryption error: %v", err)
	}

	// 生成HMAC密钥
	hmacSecret := deriveKey(digest[:16], salt, 1, 32)
	// 计算HMAC (CoAP头 + 加密后的数据)
	inputData := append(coapHeader, ciphertext...)
	mac := calculateHMAC(inputData, hmacSecret)

	// 构建最终消息
	finalMessage := buildCoAPMessage(coapHeader, ciphertext, mac)

	return finalMessage
}

// CheckMac 判断报文是否完整
func CheckMac(devId string, payload []byte, coapHeader []byte) (bool, error) {
	// 初始化Redis客户端
	client := initRedisClient()
	defer client.Close()
	// 获取加密参数
	sn1, sn2, psk, err := getRemoteControlParams(client, devId)
	if err != nil {
		log.Fatalf("Failed to get encryption parameters: %v", err)
	}
	// 拼接盐值
	salt := append(sn1, sn2...)
	// 生成PBKDF2密钥（32字节）
	digest := deriveKey(psk, salt, 1, 32)
	// 生成HMAC密钥
	hmacSecret := deriveKey(digest[:16], salt, 1, 32)
	//获取加密后的mac值
	if len(payload) < 32 {
		log.Fatalf("Decryption error: %v", err)
	}
	msg := len(payload) - 32
	enMac := payload[msg:]
	enpaylod := payload[:msg]

	// 计算HMAC (CoAP头 + 加密后的数据)
	inputData := append(coapHeader, enpaylod...)
	mac := calculateHMAC(inputData, hmacSecret)

	if !hmac.Equal(mac, enMac) {
		return false, fmt.Errorf("HMAC verification failed")
	}
	return true, nil
}

// DecryptData 数据解密
func DecryptData(devId string, payload []byte) ([]byte, error) {
	// 初始化Redis客户端
	client := initRedisClient()
	defer client.Close()
	// 获取加密参数
	sn1, sn2, psk, err := getRemoteControlParams(client, devId)
	if err != nil {
		log.Fatalf("Failed to get encryption parameters: %v", err)
	}
	// 拼接盐值
	salt := append(sn1, sn2...)
	// 生成PBKDF2密钥（32字节）
	digest := deriveKey(psk, salt, 1, 32)
	// 提取AES密钥和IV
	aesKey := digest[:16]
	iv := digest[16:32]
	//获取加密后的mac值
	if len(payload) < 32 {
		log.Fatalf("Decryption error: %v", err)
	}
	//解密
	decrypted, err := decryptAES128CBC(payload, aesKey, iv)
	if err != nil {
		log.Fatalf("Decryption error: %v", err)
	}
	fmt.Printf("Decrypted Payload: %s\n", decrypted)

	return decrypted, err
}

// DecryptDataAESMac 对数据进行解密
func DecryptDataAESMac(devId string, payload []byte, coapHeader []byte) []byte {
	// 初始化Redis客户端
	client := initRedisClient()
	defer client.Close()
	// 获取加密参数
	sn1, sn2, psk, err := getRemoteControlParams(client, devId)
	if err != nil {
		log.Fatalf("Failed to get encryption parameters: %v", err)
	}
	// 拼接盐值
	salt := append(sn1, sn2...)
	// 生成PBKDF2密钥（32字节）
	digest := deriveKey(psk, salt, 1, 32)
	// 提取AES密钥和IV
	aesKey := digest[:16]
	iv := digest[16:32]
	// 生成HMAC密钥
	hmacSecret := deriveKey(digest[:16], salt, 1, 32)

	//获取加密后的mac值
	if len(payload) < 32 {
		log.Fatalf("Decryption error: %v", err)
	}
	msg := len(payload) - 32
	enMac := payload[msg:]
	enpaylod := payload[:msg]

	// 计算HMAC (CoAP头 + 加密后的数据)
	inputData := append(coapHeader, enpaylod...)
	mac := calculateHMAC(inputData, hmacSecret)

	if !hmac.Equal(mac, enMac) {
		log.Fatalf("HMAC verification failed")
	}

	//解密
	decrypted, err := decryptAES128CBC(enpaylod, aesKey, iv)
	if err != nil {
		log.Fatalf("Decryption error: %v", err)
	}
	fmt.Printf("Decrypted Payload: %s\n", decrypted)

	return decrypted
}