// Package hashcash implements SHA-256-based hashcash
// proof-of-work validation for abuse prevention.
//
// Stamp format: 1:bits:YYMMDD:resource::counter.
//
// The SHA-256 hash of the entire stamp string must have
// at least `bits` leading zero bits.
package hashcash

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

const (
	// stampVersion is the only supported hashcash version.
	stampVersion = "1"
	// stampFields is the number of fields in a stamp.
	stampFields = 6
	// maxStampAge is how old a stamp can be before
	// rejection.
	maxStampAge = 48 * time.Hour
	// maxFutureSkew allows stamps slightly in the future.
	maxFutureSkew = 1 * time.Hour
	// pruneInterval controls how often expired stamps are
	// removed from the spent set.
	pruneInterval = 10 * time.Minute
	// dateFormatShort is the YYMMDD date layout.
	dateFormatShort = "060102"
	// dateFormatLong is the YYMMDDHHMMSS date layout.
	dateFormatLong = "060102150405"
	// dateShortLen is the length of YYMMDD.
	dateShortLen = 6
	// dateLongLen is the length of YYMMDDHHMMSS.
	dateLongLen = 12
	// bitsPerByte is the number of bits in a byte.
	bitsPerByte = 8
	// fullByteMask is 0xFF, a mask for all bits in a byte.
	fullByteMask = 0xFF
)

var (
	errInvalidFields    = errors.New("invalid stamp field count")
	errBadVersion       = errors.New("unsupported stamp version")
	errInsufficientBits = errors.New("insufficient difficulty")
	errWrongResource    = errors.New("wrong resource")
	errStampExpired     = errors.New("stamp expired")
	errStampFuture      = errors.New("stamp date in future")
	errProofFailed      = errors.New("proof-of-work failed")
	errStampReused      = errors.New("stamp already used")
	errBadDateFormat    = errors.New("unrecognized date format")
)

// Validator checks hashcash stamps for validity and
// prevents replay attacks via an in-memory spent set.
type Validator struct {
	resource string
	mu       sync.Mutex
	spent    map[string]time.Time
}

// NewValidator creates a Validator for the given resource.
func NewValidator(resource string) *Validator {
	validator := &Validator{
		resource: resource,
		mu:       sync.Mutex{},
		spent:    make(map[string]time.Time),
	}

	go validator.pruneLoop()

	return validator
}

// Validate checks a hashcash stamp. It returns nil if the
// stamp is valid and has not been seen before.
func (v *Validator) Validate(
	stamp string,
	requiredBits int,
) error {
	if requiredBits <= 0 {
		return nil
	}

	parts := strings.Split(stamp, ":")
	if len(parts) != stampFields {
		return fmt.Errorf(
			"%w: expected %d, got %d",
			errInvalidFields, stampFields, len(parts),
		)
	}

	version := parts[0]
	bitsStr := parts[1]
	dateStr := parts[2]
	resource := parts[3]

	if err := v.validateHeader(
		version, bitsStr, resource, requiredBits,
	); err != nil {
		return err
	}

	stampTime, err := parseStampDate(dateStr)
	if err != nil {
		return err
	}

	if err := validateTime(stampTime); err != nil {
		return err
	}

	if err := validateProof(
		stamp, requiredBits,
	); err != nil {
		return err
	}

	return v.checkAndRecordStamp(stamp, stampTime)
}

func (v *Validator) validateHeader(
	version, bitsStr, resource string,
	requiredBits int,
) error {
	if version != stampVersion {
		return fmt.Errorf(
			"%w: %s", errBadVersion, version,
		)
	}

	claimedBits, err := strconv.Atoi(bitsStr)
	if err != nil || claimedBits < requiredBits {
		return fmt.Errorf(
			"%w: need %d bits",
			errInsufficientBits, requiredBits,
		)
	}

	if resource != v.resource {
		return fmt.Errorf(
			"%w: got %q, want %q",
			errWrongResource, resource, v.resource,
		)
	}

	return nil
}

func validateTime(stampTime time.Time) error {
	now := time.Now()

	if now.Sub(stampTime) > maxStampAge {
		return errStampExpired
	}

	if stampTime.Sub(now) > maxFutureSkew {
		return errStampFuture
	}

	return nil
}

func validateProof(stamp string, requiredBits int) error {
	hash := sha256.Sum256([]byte(stamp))
	if !hasLeadingZeroBits(hash[:], requiredBits) {
		return fmt.Errorf(
			"%w: need %d leading zero bits",
			errProofFailed, requiredBits,
		)
	}

	return nil
}

func (v *Validator) checkAndRecordStamp(
	stamp string,
	stampTime time.Time,
) error {
	v.mu.Lock()
	defer v.mu.Unlock()

	if _, ok := v.spent[stamp]; ok {
		return errStampReused
	}

	v.spent[stamp] = stampTime

	return nil
}

// hasLeadingZeroBits checks if the hash has at least n
// leading zero bits.
func hasLeadingZeroBits(hash []byte, numBits int) bool {
	fullBytes := numBits / bitsPerByte
	remainBits := numBits % bitsPerByte

	for idx := range fullBytes {
		if hash[idx] != 0 {
			return false
		}
	}

	if remainBits > 0 && fullBytes < len(hash) {
		mask := byte(
			fullByteMask << (bitsPerByte - remainBits),
		)

		if hash[fullBytes]&mask != 0 {
			return false
		}
	}

	return true
}

// parseStampDate parses a hashcash date stamp.
// Supports YYMMDD and YYMMDDHHMMSS formats.
func parseStampDate(dateStr string) (time.Time, error) {
	switch len(dateStr) {
	case dateShortLen:
		parsed, err := time.Parse(
			dateFormatShort, dateStr,
		)
		if err != nil {
			return time.Time{}, fmt.Errorf(
				"parse date: %w", err,
			)
		}

		return parsed, nil
	case dateLongLen:
		parsed, err := time.Parse(
			dateFormatLong, dateStr,
		)
		if err != nil {
			return time.Time{}, fmt.Errorf(
				"parse date: %w", err,
			)
		}

		return parsed, nil
	default:
		return time.Time{}, fmt.Errorf(
			"%w: %q", errBadDateFormat, dateStr,
		)
	}
}

// pruneLoop periodically removes expired stamps from the
// spent set.
func (v *Validator) pruneLoop() {
	ticker := time.NewTicker(pruneInterval)
	defer ticker.Stop()

	for range ticker.C {
		v.prune()
	}
}

func (v *Validator) prune() {
	cutoff := time.Now().Add(-maxStampAge)

	v.mu.Lock()
	defer v.mu.Unlock()

	for stamp, stampTime := range v.spent {
		if stampTime.Before(cutoff) {
			delete(v.spent, stamp)
		}
	}
}