package main

// The activator mesh (issue #550 Part 4 v2, Phase 3): multiple router boxes
// cooperate so a hibernating instance wakes on exactly one box, a request
// landing on the wrong box is reverse-proxied to the owner, and a dead box's
// instances wake elsewhere. Four mechanisms, deliberately this simple:
//
//  1. Membership — boxes learn who's alive from S3 heartbeat objects
//     (internal/activator/heartbeat.go), not gossip: the bucket is already the
//     shared substrate, so no new ports or protocols between boxes (which may
//     not even share a network). Each box PUTs mesh/members/<box> every
//     meshHeartbeatInterval and LISTs the prefix; a record older than
//     meshStaleAfter is a dead box. A box's identity is LCHAT_MESH_SELF — a
//     base URL its peers can reach it at; the env's presence turns the mesh
//     on, its absence keeps today's single-box behavior byte-identical.
//  2. Rendezvous (HRW) hashing — owner(host) is a pure function of the live
//     member set (hrwOwner): no routing table, no directory service, and a
//     membership change moves only ~1/N of hosts (only onto a joiner, only
//     off a leaver).
//  3. The lease stays the truth; hashing is only a hint. The owner claims the
//     per-host S3 lease exactly as on one box; a box that wrongly believes
//     it's the owner loses the CAS race loudly (503 + Retry-After) and the
//     client's retry lands right. Transient membership disagreement therefore
//     cannot split-brain — single-writer was never the mesh's job.
//  4. Proxy-to-owner — a request for a sleeping host whose HRW owner is
//     another live box is reverse-proxied to that box's peer URL, instance
//     hostname preserved (the peer routes by Host like any front door). A
//     dead/stale owner drops out of the member set, so the hash reassigns and
//     the request falls through to a local wake — failover IS the routine
//     wake path.
//  5. Domain scoping (LCHAT_MESH_DOMAINS) — HRW placement only ever applies to
//     hosts under a configured mesh domain; everything else always resolves
//     locally, on whichever box is asked. This isn't an optimization: a
//     registry row only ever exists on the boxes it was explicitly replicated
//     to (see registerInstance in control.go), so a host outside the
//     configured domains was never given a row anywhere but its original box,
//     and letting the hash "place" it elsewhere would just 404 there. It also
//     keeps placement bounded to hosts every mesh member can actually get a
//     matching TLS cert for (typically the same domains as
//     LCHAT_WILDCARD_DOMAINS's ACME DNS-01 cert, which — unlike per-host
//     HTTP-01 — any box can obtain independently of which box DNS currently
//     points at).
//  6. Region placement (LCHAT_MESH_REGION + a space's home_region) — the
//     multi-edge foundation (docs/multi-edge.md, Phase 1). Within the in-scope
//     set, HRW picks the owner from only the live boxes whose region matches
//     the space's home_region (candidatesFor), so a space lives on a box in its
//     region and relocating it is just a home_region flip (the hash re-homes it
//     onto the destination box, the lease still guaranteeing one writer). A
//     space with no home_region, OR whose region currently has no live box,
//     falls back to hashing across the WHOLE live set — byte-identical to
//     pre-region placement, which is what keeps today's single-region fleet
//     (and every self-hoster) unchanged. Boxes learn each other's region from
//     the heartbeat body (heartbeat.go's Regions); the home_region itself is
//     replicated into every box's registry alongside the route (meshreplicate.go).
//
// Drain (activatorSvc.shutdown) deletes the heartbeat FIRST, then checkpoints
// and releases as today — peers take over immediately, and deploys, box
// failure, and scale-down are the same code path.

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"lchat/internal/activator"
)

// meshForwardedHeader marks a request one box already proxied to another:
// "serve where you stand, never re-proxy". One hop max — if two boxes
// transiently disagree about ownership, the worst case is a wrong-box wake
// the lease arbitrates, never a proxy loop. An outsider spoofing it merely
// pins its own request to the box it hit, which is exactly what every request
// could do before the mesh existed; the lease still guarantees one writer.
// Stripped before any backend sees it (unixProxy), like the sidecar marker.
const meshForwardedHeader = "X-Lchat-Mesh-Forwarded"

const (
	// meshHeartbeatInterval paces the PUT + LIST beat. Membership changes are
	// rare events; this is liveness, not coordination — and each beat is a PUT
	// plus a membership LIST, and now a region refresh (heartbeat.go's Regions:
	// a second LIST plus one small GET per live peer — a handful of requests at
	// fleet scale), so it stays cheap for nothing membership-shaped in return.
	// Slower still would work, but
	// staleness (3 beats) rides on this, and 3m dead-box detection already
	// sits well inside the lease expiry that actually gates takeover.
	meshHeartbeatInterval = 60 * time.Second
	// meshStaleAfter is how old a heartbeat may be before its box is presumed
	// dead — three missed beats. It's judged against the reader's clock but
	// the timestamp is S3's LastModified (the bucket's clock), so only
	// reader-vs-S3 skew matters, tiny against 3m. Together with the lease
	// TTL (activator.LeaseTTL, 15m) this bounds failover: a hard-killed box's
	// hosts wake elsewhere once its heartbeat is stale AND its leases expire —
	// lease-dominated, so ~15m.
	meshStaleAfter = 3 * meshHeartbeatInterval
)

// meshStore is the S3 heartbeat surface (activator.HeartbeatStore in
// production), an interface so the mesh's unit tests share an in-memory
// "bucket" between boxes.
type meshStore interface {
	Put(ctx context.Context, peer, region string) error
	List(ctx context.Context) ([]activator.Heartbeat, error)
	// Regions returns each peer's advertised placement region (peer URL ->
	// region); a peer absent from the map advertises none (agnostic). Read off
	// the hot path, on the beat — see heartbeat.go's Regions.
	Regions(ctx context.Context) (map[string]string, error)
	Delete(ctx context.Context, peer string) error
}

type meshSvc struct {
	self       string // this box's peer base URL — its mesh identity
	store      meshStore
	domains    []string  // LCHAT_MESH_DOMAINS, lowercased — see inScope. Empty = unrestricted (test convenience only; meshFromEnv always supplies a non-empty list)
	token      string    // LCHAT_MESH_TOKEN — shared secret authenticating box-to-box route replication (meshreplicate.go); empty = unrestricted (test convenience only)
	region     string    // LCHAT_MESH_REGION — this box's placement region (e.g. "ie"); "" = region-agnostic (today's behaviour). Advertised in the heartbeat body; peers read it via Regions. See candidatesFor.
	reg        *Registry // local registry handle, for reading a host's replicated home_region (candidatesFor). nil = no lookup (direct-newMesh unit tests) = every space region-agnostic.
	interval   time.Duration
	staleAfter time.Duration
	now        func() time.Time // a seam for the unit tests' clock

	mu      sync.Mutex
	members []activator.Heartbeat // last LIST; staleness applied at read time
	regions map[string]string     // last Regions view: peer URL -> advertised region (self excluded — known locally via m.region)
	proxies map[string]*httputil.ReverseProxy
	started bool

	stopCh chan struct{}
	done   chan struct{}
}

func newMesh(self string, store meshStore, domains []string) *meshSvc {
	return &meshSvc{
		self:       self,
		store:      store,
		domains:    domains,
		interval:   meshHeartbeatInterval,
		staleAfter: meshStaleAfter,
		now:        time.Now,
		regions:    make(map[string]string),
		proxies:    make(map[string]*httputil.ReverseProxy),
		stopCh:     make(chan struct{}),
		done:       make(chan struct{}),
	}
}

// meshFromEnv reads the box's mesh identity. LCHAT_MESH_SELF is a base URL
// peer boxes can reach this box at, scheme included — http on a private
// network, or https where the box presents a verifiable cert (see
// newPeerProxy). Presence turns the mesh on; absent = single box, zero mesh
// codepaths. The mesh rides on the hibernation bucket, so it requires
// LCHAT_HIBERNATE_BUCKET. It also requires LCHAT_MESH_DOMAINS (same
// comma-list shape as LCHAT_WILDCARD_DOMAINS, parsed by the same
// splitDomains): failing loudly here beats silently letting HRW "place" a
// host onto a box that has no registry row for it (a 404, not a graceful
// wake) or that could never get it a matching cert — see inScope. And
// LCHAT_MESH_TOKEN, the shared secret every mesh member authenticates
// box-to-box route replication with (meshreplicate.go) — without it any box
// could hand any other box arbitrary routes.
//
// LCHAT_MESH_REGION is optional: this box's placement region (e.g. "ie"). Set,
// it opts the box into region-scoped placement (candidatesFor); unset, the box
// is region-agnostic and everything behaves exactly as before regions existed.
func meshFromEnv(hc *hibernateConfig) (*meshSvc, error) {
	self := strings.TrimSpace(os.Getenv("LCHAT_MESH_SELF"))
	if self == "" {
		return nil, nil
	}
	if hc == nil {
		return nil, errors.New("LCHAT_MESH_SELF set but LCHAT_HIBERNATE_BUCKET missing — the mesh rides on the hibernation bucket and only places hibernating instances")
	}
	domains := splitDomains(os.Getenv("LCHAT_MESH_DOMAINS"))
	if len(domains) == 0 {
		return nil, errors.New("LCHAT_MESH_SELF set but LCHAT_MESH_DOMAINS missing — mesh placement needs an explicit domain allowlist so HRW never assigns a host to a box that has no registry row for it or can't obtain it a matching cert; set it to the domains every mesh member can serve (typically the same value as LCHAT_WILDCARD_DOMAINS)")
	}
	token := strings.TrimSpace(os.Getenv("LCHAT_MESH_TOKEN"))
	if token == "" {
		return nil, errors.New("LCHAT_MESH_SELF set but LCHAT_MESH_TOKEN missing — route replication between mesh members needs a shared secret; generate one and set it identically on every mesh box")
	}
	self = strings.TrimRight(self, "/")
	u, err := url.Parse(self)
	if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") || u.Path != "" {
		return nil, fmt.Errorf("LCHAT_MESH_SELF %q: want a base URL like https://box-a.internal or http://10.0.0.5:8080", self)
	}
	store, err := activator.NewHeartbeatStore(context.Background(), hc.Bucket, hc.meshPrefix(), hc.Region)
	if err != nil {
		return nil, fmt.Errorf("mesh heartbeat store: %w", err)
	}
	m := newMesh(self, store, domains)
	m.token = token
	m.region = strings.TrimSpace(os.Getenv("LCHAT_MESH_REGION"))
	return m, nil
}

// start registers this box (one synchronous beat, so the first sweep and the
// first request already see a membership view) and begins the heartbeat loop.
func (m *meshSvc) start() {
	m.beat()
	m.mu.Lock()
	m.started = true
	m.mu.Unlock()
	go func() {
		defer close(m.done)
		t := time.NewTicker(m.interval)
		defer t.Stop()
		for {
			select {
			case <-m.stopCh:
				return
			case <-t.C:
				m.beat()
			}
		}
	}()
}

// beat PUTs our heartbeat (with our region) and refreshes the member view, plus
// the peer-region view WHEN this box has opted into region placement. Every
// failure is soft by design: a missed PUT makes peers see us stale (they take
// over; the lease arbitrates), a failed LIST keeps the last known membership (a
// stale view misroutes at worst — same arbitration), and a failed region
// refresh keeps the last regions (at worst a space scopes to the whole fleet
// for one beat, never wrong — candidatesFor's fallback IS all-boxes).
//
// Region placement is opt-in PER MESH: a box that declares no region
// (LCHAT_MESH_REGION unset — prod today) skips the peer-region fetch entirely,
// so it adds ZERO S3 requests over pre-region behaviour — an inert feature
// costs nothing. Such a box leaves m.regions empty and candidatesFor always
// falls back to all-boxes HRW. A mixed fleet (some boxes with regions, some
// without) is safe too: an agnostic box just can't region-scope, so it hashes
// across everyone, and any resulting placement disagreement is absorbed by the
// S3 single-writer lease exactly as ordinary membership disagreement is — a
// wrong-box wake, never corruption.
func (m *meshSvc) beat() {
	timeout := m.interval
	if timeout < 5*time.Second {
		timeout = 5 * time.Second
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	// Put stays unconditional — the heartbeat is membership (its key + the write
	// timestamp), region-or-not; an unset region just writes an empty-region body.
	if err := m.store.Put(ctx, m.self, m.region); err != nil {
		log.Printf("mesh: heartbeat put failed: %v — peers will see this box stale after %s", err, m.staleAfter)
	}
	members, err := m.store.List(ctx)
	if err != nil {
		log.Printf("mesh: membership list failed: %v — keeping the last known view", err)
		return
	}
	m.mu.Lock()
	m.members = members
	m.mu.Unlock()
	if m.region == "" {
		return // not region-scoped — no peer regions to fetch, zero extra S3
	}
	// Region-scoped: refresh peers' regions. A box's region is static, so this
	// per-beat re-fetch could later narrow to fetch-on-new-peer only (a
	// planet-scale optimisation — out of scope for this slice, where the fleet
	// is a handful of boxes and a per-beat GET each is noise).
	regions, err := m.store.Regions(ctx)
	if err != nil {
		log.Printf("mesh: region refresh failed: %v — keeping the last known regions", err)
		return
	}
	m.mu.Lock()
	m.regions = regions
	m.mu.Unlock()
}

// halt stops the heartbeat loop without touching the heartbeat object — the
// shape a crashed box leaves behind (its record rots until staleAfter). The
// drain path is stop; tests use halt to simulate hard death.
func (m *meshSvc) halt() {
	m.mu.Lock()
	started := m.started
	m.started = false
	m.mu.Unlock()
	if !started {
		return
	}
	close(m.stopCh)
	<-m.done
}

// stop is the drain path: halt the loop (a beat racing the delete would
// resurrect us), then delete our heartbeat so peers take over immediately.
// The caller (activatorSvc.shutdown) checkpoints and releases AFTER this —
// that ordering is the drain contract.
func (m *meshSvc) stop(ctx context.Context) {
	m.halt()
	if err := m.store.Delete(ctx, m.self); err != nil {
		log.Printf("mesh: heartbeat delete failed: %v — peers take over when it goes stale (%s)", err, m.staleAfter)
	}
}

// liveBoxes is the current membership: every fresh heartbeat with a sane peer
// URL, self always included — our own PUT may not have landed or listed yet,
// and a box always counts itself, so the set is never empty and hrwOwner is
// always defined. Staleness is applied at read time, so a box that dies
// between polls drops out the moment its record ages past staleAfter.
func (m *meshSvc) liveBoxes() []string {
	now := m.now()
	m.mu.Lock()
	members := m.members
	m.mu.Unlock()
	boxes := []string{m.self}
	for _, hb := range members {
		if hb.Peer == m.self || now.Sub(hb.LastSeen) > m.staleAfter {
			continue
		}
		// A garbage record (someone else's object decoding to nonsense) must
		// not enter the hash — it would attract ~1/N of hosts and blackhole
		// them behind an unproxyable "peer".
		if u, err := url.Parse(hb.Peer); err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") {
			continue
		}
		boxes = append(boxes, hb.Peer)
	}
	return boxes
}

// candidatesFor returns the live boxes eligible to own host — the member set
// hrwOwner picks from. Placement keys on the space's home_region (registry.go,
// replicated to every in-scope box by meshreplicate.go), so a space pinned to a
// region is owned by a live box IN that region; relocation is a home_region flip
// and the hash re-homes it onto the destination box (docs/multi-edge.md).
//
// The fallback to ALL live boxes is load-bearing: a space with no home_region
// (every space today), or one whose region has no live box right now, hashes
// across the whole fleet exactly as before regions existed. So with no
// LCHAT_MESH_REGION set anywhere and empty home_region everywhere, this returns
// liveBoxes() and owns/remoteOwner are byte-identical to pre-region placement.
func (m *meshSvc) candidatesFor(host string) []string {
	live := m.liveBoxes()
	home := m.homeRegionOf(host)
	if home == "" {
		return live // region-agnostic space — the whole fleet, as before regions
	}
	m.mu.Lock()
	regions := m.regions
	m.mu.Unlock()
	var matched []string
	for _, b := range live {
		region := regions[b]
		if b == m.self {
			region = m.region // our own region is known locally, never via S3
		}
		if region == home {
			matched = append(matched, b)
		}
	}
	if len(matched) == 0 {
		// No live box advertises the space's region (none configured for it yet,
		// or that region's box is down): fall back to the whole fleet so the
		// space is still served somewhere — availability, and exact back-compat.
		return live
	}
	return matched
}

// homeRegionOf reads host's replicated home_region from the LOCAL registry
// (meshreplicate.go copies it onto every in-scope box, so no network). "" when
// the mesh has no registry handle (direct-newMesh unit tests), the host has no
// row, or its region is unset — all of which candidatesFor treats as agnostic.
func (m *meshSvc) homeRegionOf(host string) string {
	if m.reg == nil {
		return ""
	}
	r, err := m.reg.get(host)
	if err != nil || r == nil {
		return ""
	}
	return r.HomeRegion
}

// regionLive reports whether any live box currently advertises region — the
// precondition `lchat relocate` checks before flipping a space's home_region
// onto a destination region: there is no point re-homing a space to a region
// no box serves (candidatesFor would just fall back to the whole fleet and the
// space would never actually move). It reads the same cached membership +
// peer-region view HRW placement uses, so the caller must beat() first to
// refresh it (fanOutMeshRoute's pattern). This box's own region is known
// locally (m.region), never via S3.
func (m *meshSvc) regionLive(region string) bool {
	if region == "" {
		return false
	}
	if m.region == region {
		return true
	}
	m.mu.Lock()
	regions := m.regions
	m.mu.Unlock()
	for _, b := range m.liveBoxes() {
		if b == m.self {
			continue // handled by m.region above
		}
		if regions[b] == region {
			return true
		}
	}
	return false
}

// inScope reports whether host participates in HRW placement at all. Only
// hosts under a configured LCHAT_MESH_DOMAINS domain are ever hashed across
// boxes; everything else always resolves to whichever box is asked — correct
// because registerInstance only ever replicates a route to the boxes covered
// by those same domains (meshreplicate.go), so no other box has a row for it
// anyway, and this box (which was asked) is the only one that could possibly
// serve it. An empty domain list is unrestricted — a test-construction
// convenience (direct newMesh callers exercising bare HRW); meshFromEnv
// always supplies a non-empty list in production.
func (m *meshSvc) inScope(host string) bool {
	return hostInDomains(host, m.domains)
}

// hostInDomains is inScope's matching rule, factored out so `lchat mesh-seed`
// (meshreplicate.go) can filter routes to replicate without needing a full
// meshSvc: host equals a configured domain, or is a subdomain of one.
func hostInDomains(host string, domains []string) bool {
	if len(domains) == 0 {
		return true
	}
	host = strings.ToLower(host)
	for _, d := range domains {
		if host == d || strings.HasSuffix(host, "."+d) {
			return true
		}
	}
	return false
}

// owns reports whether this box is host's HRW owner in the live member set —
// the due-time sweeper's gate: only the owner due-wakes a host. Out-of-scope
// hosts always report true: this box is the only one ever asked about them
// (see inScope), so there is nothing to arbitrate.
func (m *meshSvc) owns(host string) bool {
	if !m.inScope(host) {
		return true
	}
	return hrwOwner(host, m.candidatesFor(host)) == m.self
}

// remoteOwner returns the peer to proxy host to when its HRW owner is another
// live box. ok=false means serve locally: host is out of mesh scope, we are
// the owner, or the previous owner is dead/stale (it's out of the member set,
// so the hash has already reassigned — possibly to us — and a local wake is
// the failover).
func (m *meshSvc) remoteOwner(host string) (peer string, ok bool) {
	if !m.inScope(host) {
		return "", false
	}
	if o := hrwOwner(host, m.candidatesFor(host)); o != m.self {
		return o, true
	}
	return "", false
}

// proxyTo forwards one request to the owning peer, instance hostname
// preserved. Proxies are built lazily, one per peer, and live for the mesh's
// life — peer URLs are stable box identities.
func (m *meshSvc) proxyTo(w http.ResponseWriter, r *http.Request, peer string) {
	m.mu.Lock()
	p := m.proxies[peer]
	if p == nil {
		u, err := url.Parse(peer)
		if err != nil || u.Host == "" {
			// liveBoxes filters garbage, so this is belt and braces.
			m.mu.Unlock()
			log.Printf("mesh: unusable peer URL %q for %s", peer, r.Host)
			http.Error(w, "instance's owner box is unreachable; retry shortly", http.StatusServiceUnavailable)
			return
		}
		p = newPeerProxy(m.self, u)
		m.proxies[peer] = p
	}
	m.mu.Unlock()
	p.ServeHTTP(w, r)
}

// newPeerProxy builds the reverse proxy to one peer box. The peer routes by
// Host exactly like the public front door, so the instance hostname must
// survive the hop; the scheme comes from the peer URL and is honored:
//
//   - http peers (a private network between boxes): dial the peer address,
//     Host header = the instance host.
//   - https peers: SNI *and* certificate verification run against the
//     instance hostname — the peer presents its wildcard/per-host cert, which
//     proves it serves this host (it could never prove it IS 10.x.y.z). The
//     outgoing URL keeps the instance host and a custom dialer pins the TCP
//     connection to the peer's address — the same shape as the box's own
//     /etc/hosts loopback pin. InsecureSkipVerify is never used.
func newPeerProxy(self string, peer *url.URL) *httputil.ReverseProxy {
	dialAddr := peerDialAddr(peer)
	return &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			pr.SetXForwarded()
			pr.Out.URL.Scheme = peer.Scheme
			if peer.Scheme == "https" {
				pr.Out.URL.Host = pr.In.Host // SNI + cert verification target
			} else {
				pr.Out.URL.Host = peer.Host
			}
			pr.Out.Host = pr.In.Host // the peer routes by Host, like any client
			pr.Out.Header.Set(meshForwardedHeader, self)
		},
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
				return (&net.Dialer{Timeout: 10 * time.Second}).DialContext(ctx, network, dialAddr)
			},
			// Bounded like unixProxy's pool, for the same leak reasons.
			MaxIdleConns:    32,
			IdleConnTimeout: 90 * time.Second,
		},
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			// An owner that died but hasn't gone stale yet (or a box-to-box
			// network fault): 503 + retry. Within staleAfter the membership
			// view drops it and retries wake the instance locally instead —
			// staleness bounds this window.
			log.Printf("mesh: proxy %s to %s: %v", r.Host, peer, err)
			w.Header().Set("Retry-After", "5")
			http.Error(w, "instance's owner box is unreachable; retry shortly", http.StatusServiceUnavailable)
		},
	}
}

// peerDialAddr is the TCP address a peer base URL dials to, default ports
// applied — computed once per proxy, since the peer URL never changes.
func peerDialAddr(peer *url.URL) string {
	if peer.Port() != "" {
		return peer.Host
	}
	if peer.Scheme == "https" {
		return peer.Hostname() + ":443"
	}
	return peer.Hostname() + ":80"
}

// hrwOwner picks host's owner from boxes by rendezvous (highest-random-weight)
// hashing: the box with the highest hash(box||host) wins. A pure function of
// (host, member set) — every box computes the same answer from membership
// alone, and when a box joins or leaves only ~1/N of hosts change owner (a
// joiner only steals, a leaver's hosts scatter over the survivors; nothing
// else moves). The result is a placement HINT, not a lock: the S3 lease stays
// the single-writer truth, which is why this can be twenty lines.
func hrwOwner(host string, boxes []string) string {
	var best string
	var bestScore uint64
	for _, b := range boxes {
		h := fnv.New64a()
		io.WriteString(h, b)
		h.Write([]byte{0}) // separator: ("ab","c") must not collide with ("a","bc")
		io.WriteString(h, host)
		if s := h.Sum64(); best == "" || s > bestScore || (s == bestScore && b > best) {
			best, bestScore = b, s
		}
	}
	return best
}