package mfa

import (
	"bytes"
	"image/png"
	"log"
	"strings"
	"time"

	"gitee.com/carlmax_my/console-core-go/pkg/core"
	"gitee.com/carlmax_my/console-core-go/pkg/jwt"

	"github.com/pquerna/otp"
	"github.com/pquerna/otp/totp"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

type MfaPatchData struct {
	UpdatedAt   time.Time
	MfaEnabled  bool
	MfaKey      string
	MfaRecCodes string
	MfaSecret   string
}

type MfaLoginCallbackFunc func(core.Context, *jwt.CustomClaims)

type SaveMfaKeyToDbCb[T string | primitive.ObjectID] func(uid T, data *MfaPatchData) error

type MfaLoginParam struct {
	MfaToken string `json:"mfa_token,omitempty"`
	MfaCode  string `json:"mfa_code,omitempty"`
	MfaBind  bool   `json:"mfa_bind,omitempty"` // bind mfa when login, for enforcement mode
}

type LoginParam struct {
	MfaLoginParam
	Username  string `json:"username,omitempty"`
	Password  string `json:"password,omitempty"`
	GrantType string `json:"grant_type,omitempty"` // 'password' | 'mfa' | 'captcha'; // default is password
}

func GenerateMfaKey(account, issuer string) (*otp.Key, error) {
	return totp.Generate(totp.GenerateOpts{
		Issuer:      issuer,
		AccountName: account, // TODO, EMAIL?
	})
}

func SaveMfaKeyToDb[T string | primitive.ObjectID](key *otp.Key, uid T, recCodes []string, cb SaveMfaKeyToDbCb[T]) error {
	mfaKey := key.String()
	mfaSecret := key.Secret()
	mfaRecCodes := strings.Join(recCodes, ",")
	// log.Println("SaveMfaKeyToDb====+", mfaKey+";"+mfaRecCodes)
	doc := MfaPatchData{
		UpdatedAt:   time.Now(),
		MfaEnabled:  true,
		MfaKey:      mfaKey,
		MfaRecCodes: mfaRecCodes,
		MfaSecret:   mfaSecret,
	}

	return cb(uid, &doc)
}

func ClearUserMfaKeyInDb[T string | primitive.ObjectID](uid T, cb SaveMfaKeyToDbCb[T]) error {
	// log.Println("ClearUserMfaKeyInDb====+", uid.Hex())
	doc := MfaPatchData{
		UpdatedAt:   time.Now(),
		MfaEnabled:  false,
		MfaKey:      "",
		MfaRecCodes: "",
		MfaSecret:   "",
	}
	return cb(uid, &doc)
}

func GenerateRecCode(key *otp.Key) (string, error) {
	// 生成TOTP恢复码
	// 这里的0是UTC时间的时间偏移量（以秒为单位）
	// 如果你需要为过去的时间生成一个TOTP，可以改变这个偏移量
	return totp.GenerateCode(key.Secret(), time.Now())
}
func GenerateRecCodes(key *otp.Key, count int) []string {
	codes := make([]string, 0)
	for i := 1; i <= count; i++ {
		code, err := GenerateRecCode(key)
		if err != nil {
			log.Println("generateRecCodes error: + ", err.Error())
		} else if code == "" {
			log.Println("generateRecCodes error(empty code)")
		} else {
			codes = append(codes, code)
		}
	}
	return codes
}

func GetMfaKeyPngBytes(key *otp.Key) ([]byte, error) {
	var buf bytes.Buffer
	img, err := key.Image(200, 200)
	if err != nil {
		return nil, err
	}
	png.Encode(&buf, img)

	return buf.Bytes(), nil
}