// Package pow implements a self-contained hashcash-style proof of work for
// the /submit endpoint. The client computes "<unix-ts>:<nonce>" locally
// (no challenge round-trip) such that SHA-256 of the stamp has a required
// number of leading zero bits. The server only has to verify, not compute.
package pow

import (
	"crypto/sha256"
	"errors"
	"math/bits"
	"strconv"
	"strings"
	"sync"
	"time"
)

// MaxAge bounds how long a stamp remains acceptable, limiting replay window
// and precomputation of stamps far in advance.
const MaxAge = 5 * time.Minute

// maxClockSkew tolerates a stamp timestamped slightly ahead of the server's clock.
const maxClockSkew = time.Minute

// Verifier checks proof-of-work stamps against a fixed difficulty and keeps
// a bounded in-memory cache of seen stamps for replay protection.
type Verifier struct {
	difficulty int

	mu   sync.Mutex
	seen map[string]time.Time
}

func NewVerifier(difficultyBits int) *Verifier {
	return &Verifier{
		difficulty: difficultyBits,
		seen:       make(map[string]time.Time),
	}
}

// Verify validates stamp "<unix-ts>:<nonce>": well-formed, not expired,
// meets the required leading-zero-bit difficulty, and not replayed.
func (v *Verifier) Verify(stamp string) error {
	tsPart, _, ok := strings.Cut(stamp, ":")
	if !ok {
		return errors.New("malformed stamp")
	}
	tsSec, err := strconv.ParseInt(tsPart, 10, 64)
	if err != nil {
		return errors.New("malformed timestamp")
	}
	ts := time.Unix(tsSec, 0)
	now := time.Now()
	if now.Sub(ts) > MaxAge || ts.Sub(now) > maxClockSkew {
		return errors.New("stamp expired or in the future")
	}

	sum := sha256.Sum256([]byte(stamp))
	if leadingZeroBits(sum[:]) < v.difficulty {
		return errors.New("insufficient proof of work")
	}

	v.mu.Lock()
	defer v.mu.Unlock()
	v.cleanupLocked(now)
	if _, dup := v.seen[stamp]; dup {
		return errors.New("replayed stamp")
	}
	v.seen[stamp] = now
	return nil
}

func (v *Verifier) cleanupLocked(now time.Time) {
	for stamp, seenAt := range v.seen {
		if now.Sub(seenAt) > MaxAge {
			delete(v.seen, stamp)
		}
	}
}

func leadingZeroBits(b []byte) int {
	count := 0
	for _, by := range b {
		if by == 0 {
			count += 8
			continue
		}
		count += bits.LeadingZeros8(by)
		break
	}
	return count
}