// Package util provide some utility functions
package util

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"os"
	"strings"
	"time"
)

// GenerateClientToken define func of generate a secure token for client identification
func GenerateClientToken(host string, serviceName string) (string, error) {
	uniqueID := fmt.Sprintf("%d", time.Now().UnixNano())
	clientInfo := fmt.Sprintf("host=%s;service=%s;id=%s", host, serviceName, uniqueID)
	secretKey := os.Getenv("TOKEN_SECRET_KEY")
	if secretKey == "" {
		return "", fmt.Errorf("TOKEN_SECRET_KEY environment variable not set")
	}

	mac := hmac.New(sha256.New, []byte(secretKey))
	mac.Write([]byte(clientInfo))
	signature := mac.Sum(nil)

	token := fmt.Sprintf("%s.%s",
		base64.URLEncoding.EncodeToString([]byte(clientInfo)),
		base64.URLEncoding.EncodeToString(signature))

	return token, nil
}

// verifyClientToken define func of verify the client token
func verifyClientToken(token string, secretKey string) (bool, string, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 2 {
		return false, "", fmt.Errorf("invalid token format")
	}

	encodedClientInfo := parts[0]
	encodedSignature := parts[1]

	clientInfoBytes, err := base64.URLEncoding.DecodeString(encodedClientInfo)
	if err != nil {
		return false, "", fmt.Errorf("failed to decode client info: %w", err)
	}
	signatureBytes, err := base64.URLEncoding.DecodeString(encodedSignature)
	if err != nil {
		return false, "", fmt.Errorf("failed to decode signature: %w", err)
	}

	mac := hmac.New(sha256.New, []byte(secretKey))
	mac.Write(clientInfoBytes)
	expectedSignature := mac.Sum(nil)

	if !hmac.Equal(signatureBytes, expectedSignature) {
		return false, "", nil // 签名不匹配，token 无效
	}

	clientInfo := string(clientInfoBytes)
	return true, clientInfo, nil
}

// --- 示例使用 ---

func main() {
	// 模拟获取客户端信息
	// 实际应用中，host 可以通过获取本地主机名，serviceName 可以是当前服务的标识。
	host, err := os.Hostname()
	if err != nil {
		fmt.Println("Error getting hostname:", err)
		host = "unknown-host"
	}
	serviceName := "my-web-service" // 替换为你的服务名

	// 设置一个示例密钥（在实际应用中，请从环境变量加载）
	// export TOKEN_SECRET_KEY="your-super-secret-key-here"
	os.Setenv("TOKEN_SECRET_KEY", "my-super-secret-and-long-key-for-hmac")

	token, err := GenerateClientToken(host, serviceName)
	if err != nil {
		fmt.Println("Error generating token:", err)
		return
	}

	fmt.Println("Generated Token:", token)

	// --- 验证 Token (示例) ---
	// 在接收方，你需要实现一个验证函数
	isValid, clientInfo, err := verifyClientToken(token, os.Getenv("TOKEN_SECRET_KEY"))
	if err != nil {
		fmt.Println("Error verifying token:", err)
	} else if isValid {
		fmt.Println("Token is valid!")
		fmt.Println("Client Info:", clientInfo) // 包含 host, service, id
	} else {
		fmt.Println("Token is invalid.")
	}
}