// Package security 安全模块，提供验证码、登录限流等安全防护能力
package security

import (
	"encoding/base64"
	"net/http"
	"strings"

	"github.com/mojocn/base64Captcha"
)

// CaptchaInstance 全局验证码实例
var CaptchaInstance *base64Captcha.Captcha

// InitCaptcha 初始化数字验证码
// @param width 验证码宽度
// @param height 验证码高度
// @author sunct
func InitCaptcha(width, height int) {
	driver := base64Captcha.NewDriverDigit(
		height,
		width,
		4,
		0.7,
		80,
	)
	CaptchaInstance = base64Captcha.NewCaptcha(driver, base64Captcha.DefaultMemStore)
}

// Verify 校验验证码答案
// @param token 验证码令牌
// @param code 用户输入的验证码
// @return bool 是否验证通过
// @author sunct
func Verify(token, code string) bool {
	return base64Captcha.DefaultMemStore.Verify(token, code, true)
}

// ServeCaptcha 生成并输出验证码图片（PNG格式）
// @param w HTTP 响应写入器
// @author sunct
func ServeCaptcha(w http.ResponseWriter) {
	if CaptchaInstance == nil {
		InitCaptcha(160, 50)
	}

	id, b64s, _ := CaptchaInstance.Generate()

	w.Header().Set("Content-Type", "image/png")
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	w.Header().Set("Pragma", "no-cache")
	w.Header().Set("Expires", "0")
	w.Header().Set("X-Captcha-Token", id)

	prefix := "data:image/png;base64,"
	if strings.HasPrefix(b64s, prefix) {
		b64s = strings.TrimPrefix(b64s, prefix)
	}

	data, _ := base64.StdEncoding.DecodeString(b64s)
	w.Write(data)
}