// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .

package encryption

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net"
	"time"

	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
)

// GenerateEncryptionKey 生成加密密钥
// 返回 keyID 和 actualKey
func GenerateEncryptionKey(remoteIP string, userAgent string, policy *serverconfigs.EncryptionKeyPolicy) (keyID string, actualKey string) {
	if policy == nil {
		policy = serverconfigs.NewEncryptionKeyPolicy()
	}

	// 1. 时间分片
	timeBucket := time.Now().Unix() / policy.TimeBucket

	// 2. IP 归一化（CIDR）
	ipPrefix := GetCIDRPrefix(remoteIP, policy.IPCIDR)

	// 3. UA 简化
	uaHash := userAgent
	if policy.UASimplify {
		uaHash = serverconfigs.SimplifyUserAgent(userAgent)
	}
	uaHashValue := sha256.Sum256([]byte(uaHash))

	// 4. 生成 KeyID（用于缓存）
	keyID = fmt.Sprintf("%s_%d_%x", ipPrefix, timeBucket, uaHashValue[:8])

	// 5. 生成实际密钥（HMAC-SHA256）
	hmacKey := hmac.New(sha256.New, []byte(policy.ServerSecret))
	hmacKey.Write([]byte(fmt.Sprintf("%s_%s_%d", ipPrefix, uaHash, timeBucket)))
	finalKey := hex.EncodeToString(hmacKey.Sum(nil))

	return keyID, finalKey
}

// GetCIDRPrefix 获取 IP 的 CIDR 前缀
func GetCIDRPrefix(ipString string, prefixLen int) string {
	ip := net.ParseIP(ipString)
	if ip == nil {
		return "0.0.0.0"
	}

	ipv4 := ip.To4()
	if ipv4 != nil {
		// IPv4
		if prefixLen > 32 {
			prefixLen = 32
		}
		if prefixLen < 0 {
			prefixLen = 0
		}

		// 计算掩码
		mask := net.CIDRMask(prefixLen, 32)
		if mask == nil {
			return "0.0.0.0"
		}

		// 应用掩码
		maskedIP := ipv4.Mask(mask)
		return maskedIP.String()
	}

	// IPv6
	ipv6 := ip.To16()
	if ipv6 != nil {
		if prefixLen > 128 {
			prefixLen = 128
		}
		if prefixLen < 0 {
			prefixLen = 0
		}

		// 计算掩码
		mask := net.CIDRMask(prefixLen, 128)
		if mask == nil {
			return "::"
		}

		// 应用掩码
		maskedIP := ipv6.Mask(mask)
		return maskedIP.String()
	}

	return "0.0.0.0"
}