// 🤖 Package idalloc is the shared, global, monotonic id allocator behind
// region-neutral space addressing (multi-edge Phase 1, docs/multi-edge.md).
//
// The problem it solves: once a space's hostname is <sqids(globalID)>.dev.11ch.at
// — ONE namespace, no country label — the id feeding sqids must be unique across
// the WHOLE fleet, or two region clouds minting concurrently would produce the
// same hostname. Home is the primary caller, but demo and direct/operator
// creates bypass home (they go straight to a box's control daemon), so every
// place that mints an id must draw from the SAME sequence.
//
// The source of truth is a single small object in the shared S3 bucket (the
// mesh's substrate — LCHAT_HIBERNATE_BUCKET) holding the high-water mark: the
// count of ids reserved so far. Advancing it is a compare-and-swap on S3's ETag,
// exactly the conditional-write style the hibernation lease (internal/activator)
// and the home vault (home/s3store.go) already ride — no new dependency, no
// second S3 client.
//
// Block allocation is the scaling trick. Id allocation happens once per space
// create, and Paul wants this to reach billions of spaces; one S3 round-trip per
// create would put the counter object on the hot path of every create fleet-wide.
// Instead an allocator reserves a RANGE of N ids in one CAS, then hands them out
// from memory, refilling with another CAS only when the block is spent. That
// bounds S3 ops to ~(creates / N) — with the default N=100, a rounding error on
// the bill even at scale (#836's lesson: request count IS the S3 cost).
//
// Guarantees: ids are unique and strictly increasing across all allocators (the
// CAS makes overlap impossible — a conflict just retries against the new
// high-water). The high-water lives in S3, so a restarted process continues past
// it, skipping the unused tail of its last block; ids need only be unique and
// increasing, never gapless, so skipped tails are fine. In-process hand-outs are
// mutex-guarded.
//
// Where it plugs in (the NEXT slice — this package does not wire it): cloud's
// createSpaceHandler (cloud.go) currently takes the id from the cloud.db spaces
// PK via cloudDB.allocateSpace; that id source becomes Allocate(ctx). The demo
// path (demo.go) and controld's direct add-instance create the same way. Home,
// cloud, and controld each construct an Allocator from the shared bucket and use
// it identically.
package idalloc

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	awss3 "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/smithy-go"
)

// DefaultBlockSize is the number of ids reserved per S3 CAS when a caller names
// none. 100 keeps S3 ops to ~1% of creates while wasting at most 99 ids on a
// restart — ids are cheap (uint64) and gaps are allowed, so a larger block only
// trades a slightly bigger restart gap for even fewer S3 ops.
const DefaultBlockSize uint64 = 100

// counterKey is the object under the bucket prefix that holds the high-water
// mark. It sits beside the mesh's membership heartbeats (<prefix>/mesh/members),
// under the same "mesh/" namespace — the bare label never collides with a
// per-instance prefix, which is always a dotted FQDN.
const counterKey = "mesh/id-sequence"

// maxReserveAttempts caps the CAS-conflict retry loop. Block allocation makes a
// conflict a 1-in-N event and each conflict means a competitor made progress
// (the high-water strictly rose), so the loser wins on a later attempt; the
// number of real allocator instances is tiny (home + a handful of region clouds
// + controld), so contention never approaches this. It exists only so a
// pathological store (e.g. a swap that always reports conflict) fails loudly
// instead of spinning forever. The context deadline is the primary bound.
const maxReserveAttempts = 100

// s3CallTimeout bounds each individual S3 call so a stalled counter object can't
// hang a space create indefinitely (the caller's ctx is the outer bound; this is
// the per-request floor, matching s3blobstore.go's per-call timeout).
const s3CallTimeout = 30 * time.Second

// counterStore is the minimal compare-and-swap surface the allocator needs over
// the one shared counter object. Splitting it out keeps the block/retry logic
// pure and unit-testable against an in-memory fake, while the S3-backed
// implementation stays a thin adapter.
//
// The version token is opaque (S3's ETag in production). By convention an empty
// token means "the object does not exist yet" — load returns "" when absent, and
// swap treats "" as a create-only precondition. S3 never returns an empty ETag
// for an existing object, so the convention is unambiguous there.
type counterStore interface {
	// load reads the current high-water value and its version token. A missing
	// object reads as value 0 with an empty token (the first writer's state).
	load(ctx context.Context) (value uint64, token string, err error)
	// swap writes newValue only if the object still matches token (or is still
	// absent, when token is ""). ok=false means a concurrent writer won the race
	// and the caller must reload and retry; a non-nil err is a real failure.
	swap(ctx context.Context, token string, newValue uint64) (ok bool, err error)
}

// Allocator hands out globally-unique, strictly-increasing ids drawn from blocks
// it reserves in the shared S3 counter. Safe for concurrent use.
type Allocator struct {
	store     counterStore
	blockSize uint64

	// mu guards the in-memory block AND serializes refills: holding it across the
	// refill CAS means two goroutines can never both reserve a block (which would
	// skip a whole block's worth of ids — still correct, but wasteful). Refills
	// are 1-in-blockSize, so the serialization is negligible.
	mu   sync.Mutex
	next uint64 // next id to hand out; ids are 1-based
	end  uint64 // one past the last id in the reserved block; next>=end ⇒ exhausted
}

// Config constructs an S3-backed Allocator. Bucket/Prefix/Region mirror the
// hibernation config (LCHAT_HIBERNATE_*) so a caller that already reads that env
// can build one with no new plumbing; the counter object lands at
// <Prefix>/mesh/id-sequence in Bucket.
type Config struct {
	Bucket    string // shared S3 bucket (the mesh substrate); required
	Prefix    string // key prefix inside the bucket (may be empty)
	Region    string // optional; the AWS default chain resolves "" from env/profile
	Endpoint  string // optional S3-compatible endpoint (R2, a test fake); implies path-style
	BlockSize uint64 // ids per CAS; 0 → DefaultBlockSize
}

// New builds an S3-backed Allocator. Credentials ride the standard AWS chain
// (the same as the hibernation lease and blob store); Endpoint lets an
// S3-compatible store or a test fake stand in.
func New(ctx context.Context, cfg Config) (*Allocator, error) {
	if strings.TrimSpace(cfg.Bucket) == "" {
		return nil, errors.New("idalloc: bucket is required")
	}
	var opts []func(*awsconfig.LoadOptions) error
	if cfg.Region != "" {
		opts = append(opts, awsconfig.WithRegion(cfg.Region))
	}
	awsCfg, err := awsconfig.LoadDefaultConfig(ctx, opts...)
	if err != nil {
		return nil, fmt.Errorf("idalloc: load aws config: %w", err)
	}
	client := awss3.NewFromConfig(awsCfg, func(o *awss3.Options) {
		if cfg.Endpoint != "" {
			o.BaseEndpoint = &cfg.Endpoint
			o.UsePathStyle = true // custom endpoints (R2, fakes) address by path
		}
		// Match s3blobstore.go: the SDK's default-on flexible checksums wrap
		// PutObject bodies in aws-chunked framing that non-AWS stores reject; the
		// counter body is tiny with a known length, so plain bodies are correct.
		o.RequestChecksumCalculation = aws.RequestChecksumCalculationWhenRequired
		o.ResponseChecksumValidation = aws.ResponseChecksumValidationWhenRequired
	})
	store := &s3Counter{
		client: client,
		bucket: cfg.Bucket,
		key:    path.Join(cfg.Prefix, counterKey),
	}
	return newWithStore(store, cfg.BlockSize), nil
}

// newWithStore builds an Allocator over any counterStore — the seam the tests
// drive with an in-memory fake, and the shared constructor path for New.
func newWithStore(store counterStore, blockSize uint64) *Allocator {
	if blockSize == 0 {
		blockSize = DefaultBlockSize
	}
	return &Allocator{store: store, blockSize: blockSize}
}

// Allocate returns the next globally-unique id, reserving a fresh block from S3
// when the current one is spent.
func (a *Allocator) Allocate(ctx context.Context) (uint64, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.next >= a.end {
		if err := a.reserveLocked(ctx); err != nil {
			return 0, err
		}
	}
	id := a.next
	a.next++
	return id, nil
}

// reserveLocked reserves the next block via a CAS on the shared counter,
// retrying on conflict. Called with a.mu held.
//
// The CAS: read the high-water H (the count reserved so far) and its version
// token, then conditional-write H+blockSize keyed on that token. Success means
// no other allocator advanced the counter in between, so this process now owns
// ids [H+1, H+blockSize]. A conflict means a competitor reserved a block first;
// its write bumped the token, so ours is stale and the store reports ok=false —
// we reload the now-higher H and try again, never reusing the ids the competitor
// took.
func (a *Allocator) reserveLocked(ctx context.Context) error {
	for attempt := 0; attempt < maxReserveAttempts; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}
		high, token, err := a.store.load(ctx)
		if err != nil {
			return fmt.Errorf("idalloc: load counter: %w", err)
		}
		newHigh := high + a.blockSize
		ok, err := a.store.swap(ctx, token, newHigh)
		if err != nil {
			return fmt.Errorf("idalloc: reserve block: %w", err)
		}
		if ok {
			a.next = high + 1  // ids are 1-based, so the block starts at H+1
			a.end = newHigh + 1 // exclusive: hand out H+1 .. H+blockSize
			return nil
		}
		// Lost the race — back off a touch before reloading, so a burst of
		// allocators contending on a fresh counter doesn't livelock hammering it.
		if err := sleepCtx(ctx, backoff(attempt)); err != nil {
			return err
		}
	}
	return fmt.Errorf("idalloc: reserve block: gave up after %d CAS conflicts", maxReserveAttempts)
}

// backoff is capped exponential with full jitter — small, since a conflict
// resolves in one or two retries in any realistic fleet.
func backoff(attempt int) time.Duration {
	const base = 5 * time.Millisecond
	const cap = 250 * time.Millisecond
	d := base << attempt
	if d > cap || d <= 0 { // <=0 guards the shift overflowing at high attempts
		d = cap
	}
	return time.Duration(rand.Int63n(int64(d) + 1))
}

// sleepCtx sleeps for d unless ctx is cancelled first.
func sleepCtx(ctx context.Context, d time.Duration) error {
	if d <= 0 {
		return ctx.Err()
	}
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-t.C:
		return nil
	}
}

// --- S3-backed counterStore --------------------------------------------------

// s3Counter is the production counterStore: the high-water lives as a decimal
// string in one S3 object, and CAS rides S3's ETag conditional writes — the same
// mechanism home/s3store.go uses for its user docs, over the same aws-sdk-go-v2
// client s3blobstore.go builds.
type s3Counter struct {
	client *awss3.Client
	bucket string
	key    string
}

func (c *s3Counter) load(ctx context.Context) (uint64, string, error) {
	ctx, cancel := context.WithTimeout(ctx, s3CallTimeout)
	defer cancel()
	out, err := c.client.GetObject(ctx, &awss3.GetObjectInput{Bucket: &c.bucket, Key: &c.key})
	if err != nil {
		if isNotFound(err) {
			return 0, "", nil // never allocated yet: value 0, empty (create-only) token
		}
		return 0, "", err
	}
	defer out.Body.Close()
	// The body is a short decimal; cap the read so a corrupt/huge object can't
	// balloon memory.
	b, err := io.ReadAll(io.LimitReader(out.Body, 64))
	if err != nil {
		return 0, "", err
	}
	v, err := strconv.ParseUint(strings.TrimSpace(string(b)), 10, 64)
	if err != nil {
		return 0, "", fmt.Errorf("idalloc: corrupt counter object %q: %w", strings.TrimSpace(string(b)), err)
	}
	token := ""
	if out.ETag != nil {
		token = *out.ETag // S3 ETags are quoted; kept verbatim for If-Match
	}
	return v, token, nil
}

func (c *s3Counter) swap(ctx context.Context, token string, newValue uint64) (bool, error) {
	ctx, cancel := context.WithTimeout(ctx, s3CallTimeout)
	defer cancel()
	body := []byte(strconv.FormatUint(newValue, 10))
	n := int64(len(body))
	in := &awss3.PutObjectInput{
		Bucket:        &c.bucket,
		Key:           &c.key,
		Body:          bytes.NewReader(body),
		ContentLength: &n,
	}
	if token == "" {
		star := "*"
		in.IfNoneMatch = &star // create-only: fails 412 if someone created it first
	} else {
		in.IfMatch = &token // replace-only: fails 412 if the object moved on
	}
	if _, err := c.client.PutObject(ctx, in); err != nil {
		if isCASConflict(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// isNotFound folds the shapes a missing object answers with (mirrors
// s3blobstore.go's s3NotFound — kept local so idalloc doesn't import package
// main).
func isNotFound(err error) bool {
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) {
		switch apiErr.ErrorCode() {
		case "NoSuchKey", "NotFound", "404":
			return true
		}
	}
	return false
}

// isCASConflict reports whether a PutObject error is a lost conditional write:
// 412 (If-Match / If-None-Match precondition failed) or 409 (S3's
// ConditionalRequestConflict when two conditional writes race). Matched by both
// the S3 error code and the raw HTTP status, so an S3-compatible store that only
// sets one of them still classifies correctly.
func isCASConflict(err error) bool {
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) {
		switch apiErr.ErrorCode() {
		case "PreconditionFailed", "ConditionalRequestConflict":
			return true
		}
	}
	var httpErr interface{ HTTPStatusCode() int }
	if errors.As(err, &httpErr) {
		switch httpErr.HTTPStatusCode() {
		case http.StatusPreconditionFailed, http.StatusConflict:
			return true
		}
	}
	return false
}