// 🤖 Mesh membership heartbeats (issue #550 Part 4 v2, Phase 3 — "the
// activator mesh"). Boxes learn who's alive from the bucket itself: each live
// box PUTs one object under a shared members prefix on an interval, LISTs the
// prefix to see its peers, and DELETEs its object on drain. S3 was chosen over
// gossip (memberlist) deliberately: the bucket is already the mesh's shared
// substrate — the leases and replicas live there — so membership needs no new
// ports or protocols between boxes, which may not even share a network (the
// current prod box plus a future AWS pool). Membership changes are rare
// events; a ~15s poll is plenty.
//
// The KEY is the record: an object key is base64url(peer URL), so one LIST —
// no per-member GETs — yields the whole membership, and LastModified (the
// bucket's clock, not the writer's) is the liveness timestamp, so writer
// clock skew never enters the staleness judgement. The body duplicates the
// URL as JSON purely for operators poking at the bucket. On a versioned
// bucket the beats churn noncurrent versions; the standard keep-N lifecycle
// prunes them.

package activator

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"path"
	"strings"
	"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"
)

// Heartbeat is one box's liveness record in the mesh membership prefix.
type Heartbeat struct {
	Peer     string    // the box's peer base URL — its mesh identity
	LastSeen time.Time // the object's S3 LastModified — the bucket's clock
}

// HeartbeatStore reads and writes box heartbeats under s3://bucket/prefix/.
type HeartbeatStore struct {
	client *awss3.Client
	bucket string
	prefix string
}

// NewHeartbeatStore builds a store on the default AWS credential chain, like
// AcquireLease. Region may be empty; the chain resolves it.
func NewHeartbeatStore(ctx context.Context, bucket, prefix, region string) (*HeartbeatStore, error) {
	var opts []func(*awsconfig.LoadOptions) error
	if region != "" {
		opts = append(opts, awsconfig.WithRegion(region))
	}
	cfg, err := awsconfig.LoadDefaultConfig(ctx, opts...)
	if err != nil {
		return nil, fmt.Errorf("load aws config: %w", err)
	}
	return &HeartbeatStore{
		client: awss3.NewFromConfig(cfg),
		bucket: bucket,
		prefix: strings.Trim(prefix, "/"),
	}, nil
}

// encodeHeartbeatKey / decodeHeartbeatKey turn a peer URL into an S3-safe key
// leaf and back (URLs carry '/' and ':', which would fake folder structure).
func encodeHeartbeatKey(peer string) string {
	return base64.RawURLEncoding.EncodeToString([]byte(peer))
}

func decodeHeartbeatKey(leaf string) (string, error) {
	b, err := base64.RawURLEncoding.DecodeString(leaf)
	if err != nil {
		return "", err
	}
	return string(b), nil
}

func (s *HeartbeatStore) key(peer string) string {
	return path.Join(s.prefix, encodeHeartbeatKey(peer))
}

// statsKey is a SIBLING object next to the membership heartbeat, under a
// separate "<prefix>-stats/" path — never mixed into the heartbeat body
// itself, so the hot membership path (List, polled every meshHeartbeatInterval
// for HRW placement) never pays to write or parse a body it doesn't need.
// Stats are read far less often (deploy/status-publish.sh's 5-minute Fleet
// publish), by a different reader (bash, via the AWS CLI, not this package).
func (s *HeartbeatStore) statsKey(peer string) string {
	return path.Join(s.prefix+"-stats", encodeHeartbeatKey(peer))
}

// BoxStats is one box's self-published capacity snapshot: band words and
// counts only, matching the Fleet dashboard's public disclosure rule
// (status.md — "never remaining slot counts... percentages, or cap values").
// Awake/Asleep are nil when unknown (never guessed, same convention as
// fleet.json's own awake/asleep). Load is a band word (comfortable/busy/
// critical) or "" when unknown — never a raw memory ratio.
type BoxStats struct {
	// Label is this box's own whimsical, hash-derived name ("ada-stargazing",
	// meshstats.go's boxLabel — the exact same algorithm and word lists as
	// deploy/status-publish.sh's box_label, kept in sync by hand: neither
	// side can compute the OTHER box's label, since it hashes a machine id
	// only that box has, so each box must self-report the label, never the
	// id it was derived from.
	Label  string `json:"label"`
	Awake  *int   `json:"awake"`
	Asleep *int   `json:"asleep"`
	Load   string `json:"load,omitempty"`
	// Region is this box's own coarse continent+country string ("EU ·
	// Finland" — see status.md's PRIVACY rule), set only via the operator's
	// explicit LCHAT_BOX_REGION env (meshstats.go's boxRegion) — unlike the
	// local box's own status-publish.sh reading, a peer has no cloud-metadata
	// auto-detect here, so absent env means omitted, never guessed.
	Region string `json:"region,omitempty"`
}

// PutStats writes this box's stats snapshot. Called on its own slower cadence
// (meshstats.go), independent of the membership beat.
func (s *HeartbeatStore) PutStats(ctx context.Context, peer string, stats BoxStats) error {
	body, err := json.Marshal(stats)
	if err != nil {
		return err
	}
	_, err = s.client.PutObject(ctx, &awss3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(s.statsKey(peer)),
		Body:   bytes.NewReader(body),
	})
	return err
}

// Put writes (or refreshes) this box's heartbeat. The write itself is the
// timestamp — LastModified advances on every beat. The body also carries the
// box's placement region (LCHAT_MESH_REGION), so peers can learn it via Regions
// and scope placement to a region (mesh.go's candidatesFor). region is omitted
// when empty, so a region-agnostic box's body stays byte-identical to before
// regions existed. The KEY is unchanged either way (base64url(peer)), so
// membership/List and every operator tool that decodes the key
// (deploy/*-posture.sh) are untouched — region rides the body, never the key.
func (s *HeartbeatStore) Put(ctx context.Context, peer, region string) error {
	rec := map[string]string{"peer": peer}
	if region != "" {
		rec["region"] = region
	}
	body, err := json.Marshal(rec)
	if err != nil {
		return err
	}
	_, err = s.client.PutObject(ctx, &awss3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(s.key(peer)),
		Body:   bytes.NewReader(body),
	})
	return err
}

// List returns every heartbeat under the prefix. Objects whose key leaf isn't
// a decodable peer record are skipped — a foreign object in the prefix must
// never poison membership.
func (s *HeartbeatStore) List(ctx context.Context) ([]Heartbeat, error) {
	var out []Heartbeat
	p := awss3.NewListObjectsV2Paginator(s.client, &awss3.ListObjectsV2Input{
		Bucket: aws.String(s.bucket),
		Prefix: aws.String(s.prefix + "/"),
	})
	for p.HasMorePages() {
		page, err := p.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		for _, obj := range page.Contents {
			peer, err := decodeHeartbeatKey(path.Base(aws.ToString(obj.Key)))
			if err != nil {
				continue
			}
			var last time.Time
			if obj.LastModified != nil {
				last = *obj.LastModified
			}
			out = append(out, Heartbeat{Peer: peer, LastSeen: last})
		}
	}
	return out, nil
}

// Regions returns each peer's advertised placement region (the "region" field
// its heartbeat body carries — see Put), keyed by peer URL. Deliberately
// SEPARATE from List: List reads membership + liveness from keys + LastModified
// alone (no per-member GET — the property the whole membership design rests
// on), whereas region lives in the body, so this GETs each small object. The
// mesh calls it off the hot request path, on its beat, into an in-memory cache
// — region changes rarely (only on relocation) and a box only needs peers'
// regions to scope placement. Best-effort per object: a body that won't GET or
// decode, or that advertises no region (an older box, or a region-agnostic
// one), is simply absent from the map, which the caller treats as agnostic.
func (s *HeartbeatStore) Regions(ctx context.Context) (map[string]string, error) {
	out := map[string]string{}
	p := awss3.NewListObjectsV2Paginator(s.client, &awss3.ListObjectsV2Input{
		Bucket: aws.String(s.bucket),
		Prefix: aws.String(s.prefix + "/"),
	})
	for p.HasMorePages() {
		page, err := p.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		for _, obj := range page.Contents {
			key := aws.ToString(obj.Key)
			peer, err := decodeHeartbeatKey(path.Base(key))
			if err != nil {
				continue // a foreign object must never poison the region view
			}
			region := s.getRegion(ctx, key)
			if region == "" {
				continue // no region advertised — leave the peer agnostic
			}
			out[peer] = region
		}
	}
	return out, nil
}

// getRegion GETs one heartbeat object and pulls its advertised region from the
// body, "" on any error — a soft miss, never fatal to the region refresh (a
// stale region view at worst scopes a space to the whole fleet for one beat).
func (s *HeartbeatStore) getRegion(ctx context.Context, key string) string {
	obj, err := s.client.GetObject(ctx, &awss3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return ""
	}
	defer obj.Body.Close()
	var rec struct {
		Region string `json:"region"`
	}
	if err := json.NewDecoder(io.LimitReader(obj.Body, 1<<12)).Decode(&rec); err != nil {
		return ""
	}
	return rec.Region
}

// Delete removes this box's heartbeat — the FIRST step of a drain, before any
// checkpoint or lease release, so peers take over immediately.
func (s *HeartbeatStore) Delete(ctx context.Context, peer string) error {
	_, err := s.client.DeleteObject(ctx, &awss3.DeleteObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(s.key(peer)),
	})
	return err
}