package server

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"database/sql"
	"encoding/base32"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/jmoiron/sqlx"
	"golang.org/x/crypto/bcrypt"
	"vdi-cloud/platform/internal/store"
)

type mfaRecord struct {
	AdminID        string     `db:"admin_id"`
	Secret         string     `db:"encrypted_secret"`
	LastCounter    int64      `db:"last_counter"`
	Pending        string     `db:"pending_secret"`
	PendingSession string     `db:"pending_session"`
	PendingExpires *time.Time `db:"pending_expires"`
	EnabledAt      *time.Time `db:"enabled_at"`
}

func (s *Server) cryptMFA(adminID, value string, seal bool) (string, error) {
	key := sha256.Sum256([]byte("vdi-admin-mfa-v1:" + s.Config.JWTSecret))
	block, _ := aes.NewCipher(key[:])
	gcm, _ := cipher.NewGCM(block)
	aad := []byte("admin_mfa:v1:" + adminID)
	if seal {
		nonce := make([]byte, gcm.NonceSize())
		if _, err := rand.Read(nonce); err != nil {
			return "", err
		}
		return base64.StdEncoding.EncodeToString(gcm.Seal(nonce, nonce, []byte(value), aad)), nil
	}
	data, err := base64.StdEncoding.DecodeString(value)
	if err != nil || len(data) < gcm.NonceSize() {
		return "", errors.New("invalid encrypted MFA secret")
	}
	plain, err := gcm.Open(nil, data[:gcm.NonceSize()], data[gcm.NonceSize():], aad)
	return string(plain), err
}

// RFC 6238 / RFC 4226, SHA-1 with dynamic truncation. Production uses 6 digits and 30 seconds.
func totpCode(secret string, counter int64, digits int) (string, error) {
	key, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(secret)
	if err != nil || len(key) < 20 || counter < 0 || (digits != 6 && digits != 8) {
		return "", errors.New("invalid TOTP parameters")
	}
	var input [8]byte
	binary.BigEndian.PutUint64(input[:], uint64(counter))
	mac := hmac.New(sha1.New, key)
	mac.Write(input[:])
	digest := mac.Sum(nil)
	offset := int(digest[len(digest)-1] & 15)
	value := binary.BigEndian.Uint32(digest[offset:offset+4]) & 0x7fffffff
	mod := uint32(1000000)
	if digits == 8 {
		mod = 100000000
	}
	return fmt.Sprintf("%0*d", digits, value%mod), nil
}

func matchTOTPCounter(secret, code string, now time.Time, last int64) (int64, bool) {
	if len(code) != 6 {
		return 0, false
	}
	for _, r := range code {
		if r < '0' || r > '9' {
			return 0, false
		}
	}
	for delta := int64(1); delta >= -1; delta-- {
		counter := now.Unix()/30 + delta
		if counter <= last {
			continue
		}
		value, err := totpCode(secret, counter, 6)
		if err == nil && hmac.Equal([]byte(value), []byte(code)) {
			return counter, true
		}
	}
	return 0, false
}

func recoveryHash(code string) string {
	code = strings.ToLower(strings.NewReplacer("-", "", " ", "").Replace(strings.TrimSpace(code)))
	if len(code) != 32 {
		return ""
	}
	if _, err := hex.DecodeString(code); err != nil {
		return ""
	}
	return store.Hash("vdi-mfa-recovery-v1:" + code)
}

func newRecoveryCodes(ctx context.Context, tx *sqlx.Tx, id string) ([]string, error) {
	if _, err := tx.ExecContext(ctx, "DELETE FROM admin_mfa_recovery WHERE admin_id=$1", id); err != nil {
		return nil, err
	}
	codes := make([]string, 0, 8)
	for i := 0; i < 8; i++ {
		var raw [16]byte
		if _, err := rand.Read(raw[:]); err != nil {
			return nil, err
		}
		value := strings.ToUpper(hex.EncodeToString(raw[:]))
		code := value[:8] + "-" + value[8:16] + "-" + value[16:24] + "-" + value[24:]
		if _, err := tx.ExecContext(ctx, "INSERT INTO admin_mfa_recovery(admin_id,code_hash) VALUES($1,$2)", id, recoveryHash(code)); err != nil {
			return nil, err
		}
		codes = append(codes, code)
	}
	return codes, nil
}

// Caller locks the administrator first, serializing login, account changes and MFA writes.
func (s *Server) consumeMFA(ctx context.Context, tx *sqlx.Tx, row mfaRecord, code string) (bool, error) {
	if hash := recoveryHash(code); hash != "" {
		result, err := tx.ExecContext(ctx, "DELETE FROM admin_mfa_recovery WHERE admin_id=$1 AND code_hash=$2", row.AdminID, hash)
		if err != nil {
			return false, err
		}
		count, err := result.RowsAffected()
		return count == 1, err
	}
	secret, err := s.cryptMFA(row.AdminID, row.Secret, false)
	if err != nil {
		return false, err
	}
	counter, ok := matchTOTPCounter(secret, code, time.Now(), row.LastCounter)
	if !ok {
		return false, nil
	}
	_, err = tx.ExecContext(ctx, "UPDATE admin_mfa SET last_counter=$1 WHERE admin_id=$2", counter, row.AdminID)
	return err == nil, err
}

func (s *Server) mfaStatus(c *gin.Context) {
	var status struct {
		Enabled           bool `json:"enabled" db:"enabled"`
		RecoveryRemaining int  `json:"recoveryRemaining" db:"recovery_remaining"`
	}
	err := s.DB.GetContext(c.Request.Context(), &status, `SELECT EXISTS(SELECT 1 FROM admin_mfa WHERE admin_id=$1 AND encrypted_secret!='') enabled,
 (SELECT count(*) FROM admin_mfa_recovery WHERE admin_id=$1) recovery_remaining`, s.admin(c).ID)
	if err != nil {
		s.dbError(c, err)
		return
	}
	c.JSON(200, status)
}

func (s *Server) beginMFA(c *gin.Context) {
	var in struct {
		Password string `json:"password" binding:"required,max=200"`
	}
	if !s.decodeConfiguration(c, &in) || !s.allowAuthAttempt(c, s.admin(c).Username, "mfa") {
		return
	}
	tx, actor := s.accountTransaction(c, false)
	if tx == nil {
		return
	}
	defer tx.Rollback()
	if bcrypt.CompareHashAndPassword([]byte(actor.PasswordHash), []byte(in.Password)) != nil {
		s.fail(c, 400, "current_password_invalid", "当前密码不正确")
		return
	}
	var enabled bool
	if err := tx.GetContext(c.Request.Context(), &enabled, "SELECT EXISTS(SELECT 1 FROM admin_mfa WHERE admin_id=$1 AND encrypted_secret!='')", actor.ID); err != nil {
		s.dbError(c, err)
		return
	}
	if enabled {
		s.fail(c, 409, "mfa_enabled", "已开启双重验证，请先关闭后重新绑定")
		return
	}
	var raw [20]byte
	if _, err := rand.Read(raw[:]); err != nil {
		s.dbError(c, err)
		return
	}
	secret := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(raw[:])
	encrypted, err := s.cryptMFA(actor.ID, secret, true)
	expires := time.Now().Add(10 * time.Minute)
	if err == nil {
		_, err = tx.ExecContext(c.Request.Context(), `INSERT INTO admin_mfa(admin_id,pending_secret,pending_session,pending_expires) VALUES($1,$2,$3,$4)
 ON CONFLICT(admin_id) DO UPDATE SET pending_secret=$2,pending_session=$3,pending_expires=$4`, actor.ID, encrypted, c.GetString("sessionId"), expires)
	}
	if err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	uri := url.URL{Scheme: "otpauth", Host: "totp", Path: "/VDI Cloud:" + actor.Username}
	uri.RawQuery = url.Values{"secret": {secret}, "issuer": {"VDI Cloud"}, "algorithm": {"SHA1"}, "digits": {"6"}, "period": {"30"}}.Encode()
	c.JSON(200, gin.H{"secret": secret, "provisioningUri": uri.String(), "expiresAt": expires})
}

func (s *Server) confirmMFA(c *gin.Context) {
	var in struct {
		Code string `json:"code" binding:"required,len=6"`
	}
	if !s.decodeConfiguration(c, &in) || !s.allowAuthAttempt(c, s.admin(c).Username, "mfa") {
		return
	}
	tx, actor := s.accountTransaction(c, false)
	if tx == nil {
		return
	}
	defer tx.Rollback()
	ctx := c.Request.Context()
	var row mfaRecord
	err := tx.GetContext(ctx, &row, "SELECT * FROM admin_mfa WHERE admin_id=$1", actor.ID)
	if errors.Is(err, sql.ErrNoRows) {
		s.fail(c, 409, "mfa_setup_expired", "请重新开始绑定")
		return
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	if row.Secret != "" {
		s.fail(c, 409, "mfa_enabled", "已开启双重验证")
		return
	}
	if row.Pending == "" || row.PendingExpires == nil || !row.PendingExpires.After(time.Now()) || row.PendingSession != c.GetString("sessionId") {
		s.fail(c, 409, "mfa_setup_expired", "绑定已过期或来自其他登录，请重新开始")
		return
	}
	secret, err := s.cryptMFA(actor.ID, row.Pending, false)
	if err != nil {
		s.dbError(c, err)
		return
	}
	counter, ok := matchTOTPCounter(secret, in.Code, time.Now(), -1)
	if !ok {
		s.fail(c, 400, "mfa_invalid", "验证码错误或已过期，请核对设备时间")
		return
	}
	_, err = tx.ExecContext(ctx, "UPDATE admin_mfa SET encrypted_secret=pending_secret,last_counter=$1,enabled_at=now(),pending_secret='',pending_session='',pending_expires=NULL WHERE admin_id=$2", counter, actor.ID)
	var codes []string
	if err == nil {
		codes, err = newRecoveryCodes(ctx, tx, actor.ID)
	}
	if err == nil {
		_, err = tx.ExecContext(ctx, "DELETE FROM admin_sessions WHERE admin_id=$1 AND id!=$2", actor.ID, c.GetString("sessionId"))
	}
	if err == nil {
		err = accountAudit(ctx, tx, actor.Username, "enable_mfa", actor.Username, c.GetString("requestId"))
	}
	if err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	c.JSON(200, gin.H{"recoveryCodes": codes})
}

func (s *Server) changeMFA(c *gin.Context, disable bool) {
	var in struct {
		Password string `json:"password" binding:"required,max=200"`
		Code     string `json:"code" binding:"required,max=80"`
	}
	if !s.decodeConfiguration(c, &in) || !s.allowAuthAttempt(c, s.admin(c).Username, "mfa") {
		return
	}
	tx, actor := s.accountTransaction(c, false)
	if tx == nil {
		return
	}
	defer tx.Rollback()
	ctx := c.Request.Context()
	if bcrypt.CompareHashAndPassword([]byte(actor.PasswordHash), []byte(in.Password)) != nil {
		s.fail(c, 400, "current_password_invalid", "当前密码不正确")
		return
	}
	var row mfaRecord
	err := tx.GetContext(ctx, &row, "SELECT * FROM admin_mfa WHERE admin_id=$1 AND encrypted_secret!=''", actor.ID)
	if errors.Is(err, sql.ErrNoRows) {
		s.fail(c, 409, "mfa_disabled", "尚未开启双重验证")
		return
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	ok, err := s.consumeMFA(ctx, tx, row, in.Code)
	if err != nil {
		s.dbError(c, err)
		return
	}
	if !ok {
		s.fail(c, 400, "mfa_invalid", "验证码错误、已使用或已过期")
		return
	}
	action := "rotate_mfa_recovery"
	var codes []string
	if disable {
		action = "disable_mfa"
		_, err = tx.ExecContext(ctx, "DELETE FROM admin_mfa WHERE admin_id=$1", actor.ID)
	} else {
		codes, err = newRecoveryCodes(ctx, tx, actor.ID)
	}
	if err == nil {
		_, err = tx.ExecContext(ctx, "DELETE FROM admin_sessions WHERE admin_id=$1 AND id!=$2", actor.ID, c.GetString("sessionId"))
	}
	if err == nil {
		err = accountAudit(ctx, tx, actor.Username, action, actor.Username, c.GetString("requestId"))
	}
	if err == nil {
		err = tx.Commit()
	}
	if err != nil {
		s.dbError(c, err)
		return
	}
	if disable {
		c.Status(204)
	} else {
		c.JSON(200, gin.H{"recoveryCodes": codes})
	}
}

func (s *Server) loginMFA(c *gin.Context, tx *sqlx.Tx, adminID, code string) bool {
	var row mfaRecord
	err := tx.GetContext(c.Request.Context(), &row, "SELECT * FROM admin_mfa WHERE admin_id=$1 AND encrypted_secret!=''", adminID)
	if errors.Is(err, sql.ErrNoRows) {
		return true
	}
	if err != nil {
		s.dbError(c, err)
		return false
	}
	if code == "" {
		s.fail(c, 401, "mfa_required", "请输入身份验证器动态码或一次性恢复码")
		return false
	}
	ok, err := s.consumeMFA(c.Request.Context(), tx, row, code)
	if err != nil {
		s.dbError(c, err)
		return false
	}
	if !ok {
		s.fail(c, 401, "mfa_invalid", "验证码错误、已使用或已过期")
		return false
	}
	return true
}