// Copyright 2024 TamaBox. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package captcha

import (
	"crypto/rand"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"math/big"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/flamego/flamego"
	"github.com/flamego/recaptcha"
	"github.com/flamego/session"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"

	"github.com/tamakyi/TamaBox/internal/conf"
	"github.com/tamakyi/TamaBox/internal/i18n"
)

const (
	// Session key prefix for storing math captcha answers
	sessionKeyPrefix = "math_captcha_"
	// TTL for each captcha in session
	captchaTTL = 5 * time.Minute
	// maxCaptchaAttempts 同一条验证码被尝试（含模态框预校验）的最大次数，
	// 超过即销毁该验证码，迫使刷新，防止对答案做无限轮询爆破。
	maxCaptchaAttempts = 5
)

// MathProblem represents an arithmetic captcha problem.
type MathProblem struct {
	ID       string `json:"id"`
	Question string `json:"question"`
}

// storedAnswer holds the expected answer and creation time.
type StoredAnswer struct {
	Answer    int
	CreatedAt time.Time
	// Attempts 累计校验尝试次数（含模态框预校验），用于达到上限后销毁验证码。
	Attempts int
}

// GenerateMath creates a new arithmetic captcha problem and stores the answer in session.
func GenerateMath(s session.Session) (*MathProblem, error) {
	// Generate random operands
	op1, err := rand.Int(rand.Reader, big.NewInt(90))
	if err != nil {
		return nil, errors.Wrap(err, "generate operand 1")
	}
	op2, err := rand.Int(rand.Reader, big.NewInt(90))
	if err != nil {
		return nil, errors.Wrap(err, "generate operand 2")
	}
	a := int(op1.Int64()) + 1 // 1-90
	b := int(op2.Int64()) + 1 // 1-90

	// Pick a random operator (+, -, *)
	opIdx, err := rand.Int(rand.Reader, big.NewInt(3))
	if err != nil {
		return nil, errors.Wrap(err, "generate operator")
	}
	var operator string
	var answer int
	switch opIdx.Int64() {
	case 0:
		operator = "+"
		answer = a + b
	case 1:
		operator = "-"
		// Ensure non-negative result
		if a < b {
			a, b = b, a
		}
		answer = a - b
	default:
		operator = "×"
		// Keep numbers small for multiplication
		a = a%10 + 1
		b = b%10 + 1
		answer = a * b
	}

	// 生成唯一 ID
	idBytes := make([]byte, 8)
	if _, err := rand.Read(idBytes); err != nil {
		return nil, errors.Wrap(err, "generate captcha id")
	}
	id := hex.EncodeToString(idBytes)

	// 清理策略说明：flamego session 接口不支持遍历，无法主动清理过期条目。
	// 泄漏可控：验证码一次性消费（VerifyMath 校验后立即删除），
	// 且每个用户同一时刻只生成 1 个新 key；session 本身过期后存储层整体清理。

	// Store answer in session
	key := sessionKeyPrefix + id
	s.Set(key, StoredAnswer{
		Answer:    answer,
		CreatedAt: time.Now(),
	})

	question := fmt.Sprintf("%d %s %d =", a, operator, b)
	return &MathProblem{
		ID:       id,
		Question: question,
	}, nil
}

// VerifyMath checks if the given answer matches the stored captcha.
// Returns (true, nil) if correct, (false, nil) if incorrect, (false, error) on error.
func VerifyMath(s session.Session, token string, lang string) (bool, error) {
	if token == "" {
		return false, errors.New(i18n.T(lang, "captcha.errEmpty"))
	}

	// Token format: "captcha_id:answer"
	parts := strings.SplitN(token, ":", 2)
	if len(parts) != 2 {
		return false, errors.New(i18n.T(lang, "captcha.errFormat"))
	}

	id := parts[0]
	answerStr := parts[1]

	answer, err := strconv.Atoi(strings.TrimSpace(answerStr))
	if err != nil {
		return false, errors.New(i18n.T(lang, "captcha.errAnswerFormat"))
	}

	// Retrieve stored answer
	key := sessionKeyPrefix + id
	stored, ok := s.Get(key).(StoredAnswer)
	if !ok {
		return false, errors.New(i18n.T(lang, "captcha.errExpired"))
	}

	// Check expiry
	if time.Since(stored.CreatedAt) > captchaTTL {
		s.Delete(key)
		return false, errors.New(i18n.T(lang, "captcha.errExpired"))
	}

	// Delete after verification attempt (one-time use)
	s.Delete(key)

	// 常量时间比较，避免答案比对产生可观测的时序侧信道。
	if subtle.ConstantTimeCompare([]byte(strconv.Itoa(answer)), []byte(strconv.Itoa(stored.Answer))) != 1 {
		return false, errors.New(i18n.T(lang, "captcha.errAnswerWrong"))
	}

	return true, nil
}

// VerifyOnly checks if the answer is correct WITHOUT consuming the captcha.
// This is used by the modal to validate before actually submitting the form.
// 同一验证码的答案预校验次数计入 Attempts，达到上限即销毁，防止脚本无限轮询爆破。
func VerifyOnly(s session.Session, id string, answer int) bool {
	key := sessionKeyPrefix + id
	stored, ok := s.Get(key).(StoredAnswer)
	if !ok {
		return false
	}
	if time.Since(stored.CreatedAt) > captchaTTL {
		s.Delete(key)
		return false
	}

	// 尝试次数限制：超过上限即销毁，迫使刷新。
	stored.Attempts++
	if stored.Attempts > maxCaptchaAttempts {
		s.Delete(key)
		return false
	}
	s.Set(key, stored)

	// 常量时间比较，避免答案比对产生可观测的时序侧信道。
	return subtle.ConstantTimeCompare([]byte(strconv.Itoa(answer)), []byte(strconv.Itoa(stored.Answer))) == 1
}

// sessionLang 从会话读取匿名用户显式选择的语言（与 Contexter 一致，语言偏好现托管于
// session 而非独立 lang cookie，避免被 session 中间件的 Set-Cookie 覆盖）；未设置时回退默认。
func sessionLang(s session.Session) string {
	if v := s.Get("lang"); v != nil {
		if str, ok := v.(string); ok && str != "" {
			return i18n.Resolve(str)
		}
	}
	return i18n.Default
}

// VerifyHandler is an HTTP handler that checks a captcha answer without consuming it.
// Used by the modal to validate before form submission.
func VerifyHandler(ctx flamego.Context, s session.Session) {
	lang := sessionLang(s)
	var req struct {
		ID     string `json:"id"`
		Answer int    `json:"answer"`
	}
	if err := json.NewDecoder(ctx.Request().Request.Body).Decode(&req); err != nil {
		ctx.ResponseWriter().WriteHeader(http.StatusBadRequest)
		_ = json.NewEncoder(ctx.ResponseWriter()).Encode(map[string]interface{}{
			"code":    40000,
			"message": i18n.T(lang, "captcha.json.badRequest"),
		})
		return
	}

	if req.ID == "" {
		ctx.ResponseWriter().WriteHeader(http.StatusBadRequest)
		_ = json.NewEncoder(ctx.ResponseWriter()).Encode(map[string]interface{}{
			"code":    40000,
			"message": i18n.T(lang, "captcha.json.idRequired"),
		})
		return
	}

	correct := VerifyOnly(s, req.ID, req.Answer)

	ctx.ResponseWriter().Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(ctx.ResponseWriter()).Encode(map[string]interface{}{
		"code":    0,
		"correct": correct,
	})
}

// GenerateHandler is an HTTP handler that generates a math captcha and returns it as JSON.
func GenerateHandler(ctx flamego.Context, s session.Session) {
	lang := sessionLang(s)
	problem, err := GenerateMath(s)
	if err != nil {
		ctx.ResponseWriter().WriteHeader(http.StatusInternalServerError)
		_ = json.NewEncoder(ctx.ResponseWriter()).Encode(map[string]interface{}{
			"code":    50000,
			"message": i18n.T(lang, "captcha.json.genFailed"),
		})
		return
	}

	ctx.ResponseWriter().Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(ctx.ResponseWriter()).Encode(map[string]interface{}{
		"code": 0,
		"data": problem,
	})
}

// Verify is a unified captcha verification function.
// It automatically chooses between recaptcha/turnstile and math captcha based on config.
// Returns (true, nil) if verification passes, (false, nil) if it fails, (false, error) on error.
func Verify(token string, remoteIP string, s session.Session, r recaptcha.RecaptchaV3, lang string) (bool, error) {
	if conf.Recaptcha.CaptchaType == "math" {
		return VerifyMath(s, token, lang)
	}

	// reCAPTCHA / Turnstile
	resp, err := r.Verify(token, remoteIP)
	if err != nil {
		logrus.WithError(err).Error("Failed to verify recaptcha/turnstile")
		return false, errors.New(i18n.T(lang, "captcha.errService"))
	}
	if !resp.Success {
		return false, errors.New(i18n.T(lang, "captcha.error"))
	}
	return true, nil
}