// Package captcha is a stateless, self-hosted proof-of-work CAPTCHA server
// implementing the Cap.js (trycap.dev) protocol in pure Go. No third-party
// service, no API keys, no database: all state is carried in HMAC-signed
// tokens plus an in-process nonce map. Wire-compatible with the cap-widget.
package captcha

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"strconv"
	"sync/atomic"
	"time"
)

const captchaTokenField = "cap-token"

type Challenge struct {
	C int `json:"c"`
	S int `json:"s"`
	D int `json:"d"`
}

type ChallengeResponse struct {
	Challenge Challenge `json:"challenge"`
	Token     string    `json:"token"`
	Expires   int64     `json:"expires"`
}

type RedeemResponse struct {
	Success bool   `json:"success"`
	Token   string `json:"token,omitempty"`
	Expires int64  `json:"expires,omitempty"`
}

type Server struct {
	secret         atomic.Pointer[[]byte]
	count, size, d int
	challengeTTL   time.Duration
	tokenTTL       time.Duration
	nonces         NonceStore
	usedTokens     NonceStore
}

type Option func(*Server)

func WithChallenge(c, s, d int) Option {
	return func(sv *Server) { sv.count, sv.size, sv.d = c, s, d }
}
func WithChallengeExpiry(d time.Duration) Option { return func(sv *Server) { sv.challengeTTL = d } }
func WithTokenExpiry(d time.Duration) Option     { return func(sv *Server) { sv.tokenTTL = d } }
func WithNonceStore(ns NonceStore) Option        { return func(sv *Server) { sv.nonces = ns } }

// WithUsedTokenStore injects the single-use store for redemption tokens
// (mirrors WithNonceStore; a durable cross-process store can implement it later).
func WithUsedTokenStore(ns NonceStore) Option { return func(sv *Server) { sv.usedTokens = ns } }

// New builds a Server. Defaults match the Cap.js reference: 50 sub-challenges,
// 32-char salts, difficulty 4, 10-min challenge / 5-min token expiry.
func New(secret []byte, opts ...Option) *Server {
	s := &Server{
		count:        50,
		size:         32,
		d:            4,
		challengeTTL: 10 * time.Minute,
		tokenTTL:     5 * time.Minute,
		nonces:       NewMemoryNonceStore(),
		usedTokens:   NewMemoryNonceStore(),
	}
	s.secret.Store(&secret)
	for _, o := range opts {
		o(s)
	}
	return s
}

func (s *Server) secretBytes() []byte { return *s.secret.Load() }

// SetSecret swaps the HMAC secret at runtime. Outstanding challenge and
// verification tokens signed with the old secret immediately fail
// verification, so callers can use this to invalidate all issued tokens.
func (s *Server) SetSecret(secret []byte) { s.secret.Store(&secret) }

func (s *Server) CreateChallenge() (ChallengeResponse, error) {
	buf := make([]byte, 25)
	if _, err := rand.Read(buf); err != nil {
		return ChallengeResponse{}, err
	}
	nonce := hex.EncodeToString(buf)
	expires := nowMs() + s.challengeTTL.Milliseconds()
	token := makeChallengeToken(s.secretBytes(), nonce, expires, s.count, s.size, s.d)
	return ChallengeResponse{
		Challenge: Challenge{C: s.count, S: s.size, D: s.d},
		Token:     token,
		Expires:   expires,
	}, nil
}

func (s *Server) Redeem(token string, solutions []string) RedeemResponse {
	claims := verifyChallengeToken(s.secretBytes(), token)
	if claims == nil {
		return RedeemResponse{Success: false}
	}
	ttl := max(time.Duration(claims.ExpiresMs-nowMs())*time.Millisecond, 0)
	if !s.nonces.MarkUsed(claims.Nonce, ttl) {
		return RedeemResponse{Success: false}
	}
	if len(solutions) != claims.C {
		return RedeemResponse{Success: false}
	}
	for i := 1; i <= claims.C; i++ {
		salt := prng(token+strconv.Itoa(i), claims.S)
		target := prng(token+strconv.Itoa(i)+"d", claims.D)
		sum := sha256.Sum256([]byte(salt + solutions[i-1]))
		h := hex.EncodeToString(sum[:])
		if len(h) < len(target) || h[:len(target)] != target {
			return RedeemResponse{Success: false}
		}
	}
	expires := nowMs() + s.tokenTTL.Milliseconds()
	vt, err := makeVerificationToken(s.secretBytes(), expires)
	if err != nil {
		return RedeemResponse{Success: false}
	}
	return RedeemResponse{Success: true, Token: vt, Expires: expires}
}

// VerifyToken validates a redemption token and consumes it: a token is
// SINGLE-USE and cannot be replayed within its TTL. The token is burned only on
// a successful verification (valid HMAC + unexpired + not previously used);
// forged/expired tokens never touch the store.
//
// Burn semantics — IMPORTANT for callers: consumption happens here, at the
// captcha layer, BEFORE any downstream business logic runs. So if a caller
// verifies the captcha and THEN rejects the request for a non-captcha reason
// (wrong password, invalid email, form validation error), the token is already
// spent. This is deliberate: it closes the replay window. The Cap widget
// re-issues a fresh token on its next solve, so a re-rendered form/error
// fragment that still carries the widget (with its reset flow) lets the visitor
// re-solve and retry — callers MUST keep the widget present on rejection paths.
func (s *Server) VerifyToken(token string) bool {
	claims := parseVerificationToken(s.secretBytes(), token)
	if claims == nil {
		return false
	}
	ttl := max(time.Duration(claims.ExpiresMs-nowMs())*time.Millisecond, 0)
	// MarkUsed returns false if this token id was already consumed → replay.
	return s.usedTokens.MarkUsed(claims.ID, ttl)
}

// VerifyRequest reads the cap-token form field and verifies it. Callers that
// have already parsed the form (typical for HTMX handlers) can rely on
// FormValue; it triggers ParseMultipartForm/ParseForm lazily otherwise.
func (s *Server) VerifyRequest(r *http.Request) bool {
	return s.VerifyToken(r.FormValue(captchaTokenField))
}