package files

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

	"github.com/golang-jwt/jwt/v5"
)

const captchaTTL = 10 * time.Minute

type captchaClaims struct {
	Expected int `json:"ans"`
	jwt.RegisteredClaims
}

// jtiStore tracks used token IDs to make each captcha token single-use.
type jtiStore struct {
	mu   sync.Mutex
	used map[string]time.Time // jti → token expiry
}

var usedJTIs = &jtiStore{used: make(map[string]time.Time)}

func init() {
	// Periodically remove expired JTIs so the map doesn't grow forever.
	go func() {
		ticker := time.NewTicker(captchaTTL)
		defer ticker.Stop()
		for range ticker.C {
			usedJTIs.mu.Lock()
			now := time.Now()
			for jti, exp := range usedJTIs.used {
				if now.After(exp) {
					delete(usedJTIs.used, jti)
				}
			}
			usedJTIs.mu.Unlock()
		}
	}()
}

// claim marks a JTI as used. Returns false if it was already used.
func (s *jtiStore) claim(jti string, exp time.Time) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.used[jti]; exists {
		return false
	}
	s.used[jti] = exp
	return true
}

func randInt(min, max int) int {
	n, _ := rand.Int(rand.Reader, big.NewInt(int64(max-min+1)))
	return int(n.Int64()) + min
}

func randJTI() string {
	b := make([]byte, 16)
	rand.Read(b) //nolint:errcheck
	return hex.EncodeToString(b)
}

// CaptchaHandler generates a math question and returns it with a signed token
// that embeds the expected answer. The token is single-use and expires in 10 minutes.
func CaptchaHandler(jwtSecret string) http.HandlerFunc {
	type response struct {
		Question string `json:"question"`
		Token    string `json:"token"`
	}

	return func(w http.ResponseWriter, r *http.Request) {
		ops := []string{"+", "-", "×"}
		op := ops[randInt(0, 2)]

		var a, b, expected int
		switch op {
		case "+":
			a, b = randInt(2, 20), randInt(1, 15)
			expected = a + b
		case "-":
			a, b = randInt(5, 20), randInt(1, 10)
			expected = a - b
		case "×":
			a, b = randInt(1, 10), randInt(1, 10)
			expected = a * b
		}

		exp := time.Now().Add(captchaTTL)
		token := jwt.NewWithClaims(jwt.SigningMethodHS256, captchaClaims{
			Expected: expected,
			RegisteredClaims: jwt.RegisteredClaims{
				ID:        randJTI(),
				ExpiresAt: jwt.NewNumericDate(exp),
			},
		})
		signed, err := token.SignedString([]byte(jwtSecret))
		if err != nil {
			http.Error(w, "internal error", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(response{ //nolint:errcheck
			Question: fmt.Sprintf("%d %s %d", a, op, b),
			Token:    signed,
		})
	}
}

// verifyCaptcha validates the signed token and checks that answerStr matches
// the embedded expected answer. Each token can only be used once.
func verifyCaptcha(tokenStr, answerStr, jwtSecret string) error {
	answer, err := strconv.Atoi(answerStr)
	if err != nil {
		return errors.New("invalid answer")
	}

	parsed, err := jwt.ParseWithClaims(tokenStr, &captchaClaims{}, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, errors.New("unexpected signing method")
		}
		return []byte(jwtSecret), nil
	})
	if err != nil || !parsed.Valid {
		return errors.New("invalid or expired token")
	}

	claims, ok := parsed.Claims.(*captchaClaims)
	if !ok || claims.ID == "" {
		return errors.New("malformed token")
	}

	if claims.Expected != answer {
		return errors.New("wrong answer")
	}

	if !usedJTIs.claim(claims.ID, claims.ExpiresAt.Time) {
		return errors.New("token already used")
	}

	return nil
}