package totp

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base32"
	"encoding/binary"
	"fmt"
	"strings"
	"sync"
	"time"
)

// TOTPConfig TOTP配置
type TOTPConfig struct {
	Secret    string        // TOTP密钥
	Digits    int           // TOTP验证码长度
	TimeStep  time.Duration // TOTP时间步长
	Algorithm string        // TOTP算法
}

var (
	defaultConfig = TOTPConfig{
		Secret:    "",
		Digits:    6,
		TimeStep:  30 * time.Second,
		Algorithm: "SHA1",
	}
	mu sync.RWMutex
)

// SetTOTPSecret 设置TOTP密钥
func SetTOTPSecret(secret string) {
	mu.Lock()
	defer mu.Unlock()
	defaultConfig.Secret = secret
}

// GetTOTPSecret 获取TOTP密钥
func GetTOTPSecret() string {
	mu.RLock()
	defer mu.RUnlock()
	return defaultConfig.Secret
}

// GenerateTOTPCode 生成TOTP验证码
func GenerateTOTPCode() (string, error) {
	mu.RLock()
	config := defaultConfig
	mu.RUnlock()

	if config.Secret == "" {
		return "", fmt.Errorf("TOTP密钥未设置")
	}

	// 确保密钥是Base32编码的
	secret := strings.ToUpper(config.Secret)
	secret = strings.ReplaceAll(secret, " ", "")
	secretBytes, err := base32.StdEncoding.DecodeString(secret)
	if err != nil {
		return "", fmt.Errorf("解码TOTP密钥失败: %w", err)
	}

	// 获取当前时间戳并转换为TOTP时间计数器
	timeCounter := uint64(time.Now().Unix()) / uint64(config.TimeStep.Seconds())

	// 将时间计数器转换为字节数组
	timeBytes := make([]byte, 8)
	binary.BigEndian.PutUint64(timeBytes, timeCounter)

	// 使用HMAC-SHA1计算TOTP值
	h := hmac.New(sha1.New, secretBytes)
	h.Write(timeBytes)
	hash := h.Sum(nil)

	// 根据RFC 6238，我们使用哈希的最后一个字节的低4位作为偏移
	offset := hash[len(hash)-1] & 0x0f

	// 从哈希中提取4字节并转换为整数
	binary := binary.BigEndian.Uint32(hash[offset : offset+4])

	// 屏蔽最高位并获取指定位数的数字
	totp := binary & 0x7fffffff % uint32(pow10(config.Digits))

	// 将数字格式化为字符串，填充前导零
	return fmt.Sprintf("%0*d", config.Digits, totp), nil
}

// ValidateTOTPCode 验证TOTP验证码
func ValidateTOTPCode(code string) bool {
	generatedCode, err := GenerateTOTPCode()
	if err != nil {
		return false
	}
	return code == generatedCode
}

// pow10 计算10的n次方
func pow10(n int) int {
	result := 1
	for i := 0; i < n; i++ {
		result *= 10
	}
	return result
}