package edge

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/carlosframework/platform/internal/provision"
)

// waitFor polls cond until it holds or within elapses. Every prober assertion
// is time-based (a probe tick has to land before state can move), so the
// windows here are deliberately generous — an order of magnitude above the
// tick they wait on — rather than exact: a loaded CI box must not turn a
// correct prober into a red build.
func waitFor(t *testing.T, within time.Duration, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(within)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	if !cond() {
		t.Fatalf("timed out after %s waiting for %s", within, what)
	}
}

// switchableBackend is an httptest server whose answer flips between the
// configured status and 200 under an atomic flag, plus a hit counter.
//
// The flag exists because the obvious alternative — reassigning
// srv.Config.Handler once the server is live — is a data race: net/http reads
// Server.Handler on every connection it serves, so mutating it while the
// prober is mid-tick is exactly the kind of race `-race` exists to catch.
func switchableBackend(t *testing.T, downStatus int) (*httptest.Server, *atomic.Bool, *atomic.Int64) {
	t.Helper()
	ok := new(atomic.Bool)
	hits := new(atomic.Int64)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		if ok.Load() {
			w.WriteHeader(http.StatusOK)
			return
		}
		w.WriteHeader(downStatus)
	}))
	t.Cleanup(ts.Close)
	return ts, ok, hits
}

func TestProberFailsAwayAndBack(t *testing.T) {
	healthy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(healthy.Close)
	flaky, flakyOK, _ := switchableBackend(t, http.StatusInternalServerError)

	healthyAddr := healthy.Listener.Addr().String()
	flakyAddr := flaky.Listener.Addr().String()
	spec := probeSpec{
		Host:  "m.example.com",
		Addrs: []string{healthyAddr, flakyAddr},
		Health: provision.Health{
			Path: "/", Status: 200, IntervalMS: 20, FailAfter: 3, RiseAfter: 2,
		}.WithDefaults(),
	}

	p := newProberSet()
	defer p.Stop()
	// Optimism: a fresh Sync must not 503 a healthy fleet while the first
	// probes run, so both upstreams are live before any probe has landed.
	p.Sync(map[string]probeSpec{"m.example.com": spec})
	if live := p.Live("m.example.com"); len(live) != 2 {
		t.Fatalf("live immediately after Sync = %v, want both upstreams (all start live)", live)
	}

	waitFor(t, 2*time.Second, "the 500ing upstream to fail away", func() bool {
		return len(p.Live("m.example.com")) == 1
	})
	if live := p.Live("m.example.com"); live[0] != healthyAddr {
		t.Fatalf("live = %v, want only the healthy upstream %s", live, healthyAddr)
	}
	st := p.State()["m.example.com"]
	if len(st) != 2 {
		t.Fatalf("State() = %v, want one entry per upstream", st)
	}
	down := st[1]
	if down.Addr != flakyAddr || down.Up {
		t.Fatalf("State()[1] = %+v, want %s down", down, flakyAddr)
	}
	if down.Fails < 3 {
		t.Errorf("down upstream Fails = %d, want >= FailAfter (3)", down.Fails)
	}
	if down.LastError == "" {
		t.Error("down upstream LastError is empty, want the wrong-status reason")
	}
	if down.Since.IsZero() {
		t.Error("down upstream Since is zero, want the transition time")
	}

	// The flaky upstream recovers → rejoins after RiseAfter passes.
	flakyOK.Store(true)
	waitFor(t, 2*time.Second, "the recovered upstream to rejoin", func() bool {
		return len(p.Live("m.example.com")) == 2
	})
	if got := p.State()["m.example.com"][1]; !got.Up || got.Fails != 0 || got.LastError != "" {
		t.Errorf("recovered upstream = %+v, want up with cleared fails/error", got)
	}
}

func TestProberSendsConfiguredHostHeader(t *testing.T) {
	var mu sync.Mutex
	var seenHost, seenPath string
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		seenHost, seenPath = r.Host, r.URL.Path
		mu.Unlock()
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(ts.Close)
	saw := func() (string, string) {
		mu.Lock()
		defer mu.Unlock()
		return seenHost, seenPath
	}

	// HostHeader set: it wins over the route host, so a shared upstream that
	// vhosts on Host can be probed on the name it actually answers to.
	p := newProberSet()
	p.Sync(map[string]probeSpec{"m.example.com": {
		Host:  "m.example.com",
		Addrs: []string{ts.Listener.Addr().String()},
		Health: provision.Health{
			Path: "/up", Status: 200, HostHeader: "probe.internal", IntervalMS: 20,
		}.WithDefaults(),
	}})
	waitFor(t, 2*time.Second, "the first probe with an explicit host header", func() bool {
		h, _ := saw()
		return h != ""
	})
	if h, path := saw(); h != "probe.internal" || path != "/up" {
		t.Fatalf("probe hit Host=%q path=%q, want probe.internal /up", h, path)
	}
	p.Stop()

	mu.Lock()
	seenHost, seenPath = "", ""
	mu.Unlock()

	// HostHeader unset: the route host is substituted — WithDefaults
	// deliberately leaves HostHeader empty so the prober can see "unset".
	p2 := newProberSet()
	defer p2.Stop()
	p2.Sync(map[string]probeSpec{"m.example.com": {
		Host:   "m.example.com",
		Addrs:  []string{ts.Listener.Addr().String()},
		Health: provision.Health{IntervalMS: 20}.WithDefaults(),
	}})
	waitFor(t, 2*time.Second, "the first probe with a defaulted host header", func() bool {
		h, _ := saw()
		return h != ""
	})
	if h, path := saw(); h != "m.example.com" || path != "/" {
		t.Fatalf("probe hit Host=%q path=%q, want the route host m.example.com and /", h, path)
	}
}

func TestSyncStopsVanishedHosts(t *testing.T) {
	backend, ok, hits := switchableBackend(t, http.StatusInternalServerError)
	ok.Store(true)

	p := newProberSet()
	defer p.Stop()
	spec := probeSpec{
		Host:   "m.example.com",
		Addrs:  []string{backend.Listener.Addr().String()},
		Health: provision.Health{IntervalMS: 20}.WithDefaults(),
	}
	p.Sync(map[string]probeSpec{"m.example.com": spec})
	waitFor(t, 2*time.Second, "the host to be probed at least twice", func() bool {
		return hits.Load() >= 2
	})

	// An unchanged Sync must NOT restart the prober (Reload runs on every
	// fingerprint move; restarting would reset every counter and re-probe
	// the fleet from scratch each time an unrelated route changed).
	before := p.State()["m.example.com"][0].Since
	p.Sync(map[string]probeSpec{"m.example.com": spec})
	if after := p.State()["m.example.com"][0].Since; !after.Equal(before) {
		t.Errorf("unchanged Sync restarted the prober (Since %v -> %v)", before, after)
	}

	// The host vanishes: its goroutine is stopped before Sync returns, so
	// the hit count is frozen from here on — no settle window needed.
	p.Sync(map[string]probeSpec{})
	if st := p.State(); len(st) != 0 {
		t.Fatalf("State() = %v, want empty after the host vanished", st)
	}
	if live := p.Live("m.example.com"); len(live) != 0 {
		t.Fatalf("Live() = %v, want empty after the host vanished", live)
	}
	frozen := hits.Load()
	time.Sleep(150 * time.Millisecond) // >> the 20ms interval it was probing at
	if now := hits.Load(); now != frozen {
		t.Fatalf("backend hit %d more time(s) after Sync dropped the host — goroutine leaked", now-frozen)
	}
}

// TestSyncReplacesChangedSpecWithoutAGap covers the replace path — the one
// Sync branch where a host is removed AND re-added in the same call. Two
// invariants ride on how that is ordered, and both are here so a refactor that
// simplifies the ordering fails loudly:
//
//  1. Live(host) is NEVER empty across the swap. The replacement is built
//     (all-live, per the optimism rule) inside the same critical section that
//     removes the outgoing prober, so a request racing a Reload sees the old
//     table or the new one, never a gap. A gap would be a 503 for a route whose
//     upstreams are perfectly healthy — a self-inflicted outage on every
//     addr-set edit.
//  2. The dropped backend is not probed again once Sync RETURNS. Retired
//     probers are stop()ed (which waits for the goroutine to exit) before any
//     replacement goroutine is launched, so there is no straggler tick against
//     an upstream the route no longer names.
func TestSyncReplacesChangedSpecWithoutAGap(t *testing.T) {
	const host = "m.example.com"
	newOK := func() (*httptest.Server, *atomic.Int64) {
		srv, ok, hits := switchableBackend(t, http.StatusInternalServerError)
		ok.Store(true)
		return srv, hits
	}
	a, _ := newOK()
	b, bHits := newOK()
	c, _ := newOK()
	aAddr, bAddr, cAddr := a.Listener.Addr().String(), b.Listener.Addr().String(), c.Listener.Addr().String()

	spec := func(addrs ...string) probeSpec {
		return probeSpec{
			Host:  host,
			Addrs: addrs,
			Health: provision.Health{
				Path: "/", Status: 200, IntervalMS: 20, FailAfter: 3, RiseAfter: 2,
			}.WithDefaults(),
		}
	}

	p := newProberSet()
	defer p.Stop()
	p.Sync(map[string]probeSpec{host: spec(aAddr, bAddr)})
	waitFor(t, 2*time.Second, "the outgoing addr set to be probed", func() bool {
		return bHits.Load() >= 2
	})

	// Hammer Live from a second goroutine for the whole duration of the Sync.
	// No sleep in the loop: the window being guarded is a few instructions
	// wide, so the poller has to be as tight as the race it is hunting. The
	// polling channel is closed after its first reading, so the swap cannot
	// start before the poller is actually running — otherwise a scheduler that
	// gets round to the goroutine late turns the guard into a no-op.
	stop := make(chan struct{})
	done := make(chan struct{})
	var sawGap bool
	polls := new(atomic.Int64)
	go func() {
		defer close(done)
		for {
			polls.Add(1)
			if live := p.Live(host); len(live) == 0 {
				sawGap = true
				return
			}
			select {
			case <-stop:
				return
			default:
			}
		}
	}()
	// Do not swap until the poller has been scheduled and taken several
	// readings — otherwise a late-starting goroutine would see only the
	// finished state and the guard would silently assert nothing.
	waitFor(t, 2*time.Second, "the Live poller to be running concurrently", func() bool {
		return polls.Load() > 2
	})

	p.Sync(map[string]probeSpec{host: spec(aAddr, cAddr)})

	// Freeze the retired backend's counter the instant Sync returns: stop() is
	// synchronous, so no settle window is allowed here.
	frozen := bHits.Load()
	close(stop)
	<-done
	if sawGap {
		t.Fatalf("Live(%s) was empty during the spec swap — a healthy route would have 503ed", host)
	}

	time.Sleep(150 * time.Millisecond) // >> the 20ms interval it was probing at
	if now := bHits.Load(); now != frozen {
		t.Fatalf("dropped upstream probed %d more time(s) after Sync returned — straggler tick from a retired prober", now-frozen)
	}

	// The new spec is the one running: the replacement's addrs, all live.
	live := p.Live(host)
	if len(live) != 2 || live[0] != aAddr || live[1] != cAddr {
		t.Fatalf("live after the swap = %v, want the new addr set [%s %s]", live, aAddr, cAddr)
	}
}

// TestSyncNormalizesHealthBeforeComparing pins finding #2 of the fix round: the
// same probing job expressed with a raw Health{} and with a WithDefaults()ed one
// must compare EQUAL, or a caller that takes a mixed path (the registry parse on
// one Reload, a defaulted spec on the next) restarts every prober on the box —
// resetting every fail/rise run and re-probing the whole fleet — on each Reload.
func TestSyncNormalizesHealthBeforeComparing(t *testing.T) {
	const host = "m.example.com"
	backend, ok, _ := switchableBackend(t, http.StatusInternalServerError)
	ok.Store(true)

	p := newProberSet()
	defer p.Stop()
	proberFor := func() *hostProber {
		p.mu.Lock()
		defer p.mu.Unlock()
		return p.hosts[host]
	}

	raw := probeSpec{
		Host:   host,
		Addrs:  []string{backend.Listener.Addr().String()},
		Health: provision.Health{IntervalMS: 20},
	}
	defaulted := raw
	defaulted.Health = raw.Health.WithDefaults()
	if raw.Health == defaulted.Health {
		t.Fatal("the two spellings are identical; this test is asserting nothing")
	}

	p.Sync(map[string]probeSpec{host: raw})
	first := proberFor()
	if first == nil {
		t.Fatal("no prober after the first Sync")
	}
	// Identity, not just Since: the exact goroutine must survive both spellings.
	p.Sync(map[string]probeSpec{host: defaulted})
	if got := proberFor(); got != first {
		t.Error("a defaulted spec restarted the prober started from the raw one")
	}
	p.Sync(map[string]probeSpec{host: raw})
	if got := proberFor(); got != first {
		t.Error("a raw spec restarted the prober started from the defaulted one")
	}
}

func TestEjectIsImmediateAndProbeRejoins(t *testing.T) {
	a := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(a.Close)
	b := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(b.Close)
	aAddr, bAddr := a.Listener.Addr().String(), b.Listener.Addr().String()

	p := newProberSet()
	defer p.Stop()
	p.Sync(map[string]probeSpec{"m.example.com": {
		Host:  "m.example.com",
		Addrs: []string{aAddr, bAddr},
		Health: provision.Health{
			Path: "/", Status: 200, IntervalMS: 20, FailAfter: 3, RiseAfter: 2,
		}.WithDefaults(),
	}})
	waitFor(t, 2*time.Second, "both upstreams live", func() bool {
		return len(p.Live("m.example.com")) == 2
	})

	// Eject is the in-request dial-failure fast path: the drop is visible on
	// the very next Live(), with no probe tick in between.
	p.Eject("m.example.com", aAddr)
	live := p.Live("m.example.com")
	if len(live) != 1 || live[0] != bAddr {
		t.Fatalf("live right after Eject = %v, want only %s", live, bAddr)
	}
	got := p.State()["m.example.com"][0]
	if got.Addr != aAddr || got.Up {
		t.Fatalf("ejected upstream state = %+v, want %s down", got, aAddr)
	}
	if got.Fails != 3 {
		t.Errorf("ejected upstream Fails = %d, want FailAfter (3)", got.Fails)
	}
	if got.LastError == "" {
		t.Error("ejected upstream LastError is empty, want the ejection reason")
	}

	// The upstream was never actually unhealthy, so probe passes rejoin it.
	waitFor(t, 2*time.Second, "the ejected upstream to rejoin on probe passes", func() bool {
		return len(p.Live("m.example.com")) == 2
	})

	// Ejecting an addr this host does not have (or an unknown host) is a
	// no-op, not a panic: the proxy can race a Reload that dropped either.
	p.Eject("m.example.com", "203.0.113.9:1")
	p.Eject("nobody.example.com", aAddr)
	if len(p.Live("m.example.com")) != 2 {
		t.Fatalf("stray Eject disturbed the live set: %v", p.Live("m.example.com"))
	}
}

// probeState reads one upstream's state off a hostProber the way State() would,
// for tests that drive record() directly instead of waiting on probe ticks.
func probeState(hp *hostProber, addr string) UpstreamState {
	hp.mu.Lock()
	defer hp.mu.Unlock()
	st := hp.state[addr]
	if st == nil {
		return UpstreamState{}
	}
	return UpstreamState{Addr: addr, Up: st.up, Since: st.since, Fails: st.fails, LastError: st.lastError}
}

// TestProberClampsNonPositiveInterval guards the clamp that keeps a bad health
// JSON from taking the whole edge process down: time.NewTicker PANICS on a
// non-positive duration, and the prober's ticker lives in the goroutine that
// serves every probed route on the box. Sync's WithDefaults normalization
// handles interval_ms 0; it passes NEGATIVES straight through, and
// SetUpstreams (plus the route PATCH endpoint that follows it) writes the
// health column directly, so the clamp is the only thing between a -1 and a
// crash.
func TestProberClampsNonPositiveInterval(t *testing.T) {
	for _, ms := range []int{0, -1, -60000} {
		hp := newHostProber(probeSpec{
			Host:   "m.example.com",
			Addrs:  []string{"10.0.0.1:80"},
			Health: provision.Health{IntervalMS: ms},
		}, nil)
		if hp.interval != fallbackProbeInterval {
			t.Errorf("interval_ms %d clamped to %v, want the %v fallback", ms, hp.interval, fallbackProbeInterval)
		}
	}

	// And end to end through Sync: a negative interval survives WithDefaults,
	// so this is the shape that would reach time.NewTicker. It must probe, not
	// panic.
	backend, ok, hits := switchableBackend(t, http.StatusInternalServerError)
	ok.Store(true)
	p := newProberSet()
	defer p.Stop()
	p.Sync(map[string]probeSpec{"m.example.com": {
		Host:   "m.example.com",
		Addrs:  []string{backend.Listener.Addr().String()},
		Health: provision.Health{Path: "/", Status: 200, IntervalMS: -1},
	}})
	waitFor(t, 2*time.Second, "the first probe of a negative-interval spec", func() bool {
		return hits.Load() >= 1
	})
	time.Sleep(300 * time.Millisecond)
	// At the clamped 2000ms cadence this window holds the immediate first
	// probe and possibly a second on a slow box; an unclamped interval would
	// be in the hundreds. A bound, not an exact count.
	if n := hits.Load(); n > 5 {
		t.Fatalf("backend probed %d times in ~300ms — the negative interval was not clamped to %v", n, fallbackProbeInterval)
	}
}

// TestProberClampsNonPositiveFailAndRiseCounts is the other half of the same
// amendment — a non-positive fail_after must not fail an upstream away on a
// single blip, nor a non-positive rise_after rejoin it on a single pass — and
// it also pins the recovery-state truth: a DOWN upstream keeps its Fails and
// LastError through the passing probes of a recovery, so the dashboard never
// shows `Up:false, Fails:0, LastError:""` (down for no stated reason) to an
// operator watching a failover. Driven through record() directly, so it asserts
// exact counts with no timing in play.
func TestProberClampsNonPositiveFailAndRiseCounts(t *testing.T) {
	const addr = "10.0.0.1:80"
	hp := newHostProber(probeSpec{
		Host:   "m.example.com",
		Addrs:  []string{addr},
		Health: provision.Health{Path: "/", Status: 200, IntervalMS: 20, FailAfter: -1, RiseAfter: 0},
	}, nil)
	if hp.failAt != fallbackFailAfter || hp.riseAt != fallbackRiseAfter {
		t.Fatalf("clamped counts = fail %d rise %d, want %d/%d", hp.failAt, hp.riseAt, fallbackFailAfter, fallbackRiseAfter)
	}

	boom := errors.New("dial: connection refused")
	for i := 1; i < fallbackFailAfter; i++ {
		hp.record(addr, boom)
		if st := probeState(hp, addr); !st.Up {
			t.Fatalf("upstream went down after %d failure(s), want the clamped fail_after (%d)", i, fallbackFailAfter)
		}
	}
	hp.record(addr, boom)
	down := probeState(hp, addr)
	if down.Up {
		t.Fatalf("upstream still up after %d failures, want down", fallbackFailAfter)
	}
	if down.Fails != fallbackFailAfter || down.LastError != boom.Error() {
		t.Fatalf("down state = %+v, want %d fails and the dial reason", down, fallbackFailAfter)
	}

	for i := 1; i < fallbackRiseAfter; i++ {
		hp.record(addr, nil)
		st := probeState(hp, addr)
		if st.Up {
			t.Fatalf("upstream rejoined after %d pass(es), want the clamped rise_after (%d)", i, fallbackRiseAfter)
		}
		// Still down → still owes the operator a reason.
		if st.Fails != fallbackFailAfter || st.LastError != boom.Error() {
			t.Fatalf("recovering (still down) state = %+v, want the failure count and reason held until it comes up", st)
		}
	}
	hp.record(addr, nil)
	up := probeState(hp, addr)
	if !up.Up || up.Fails != 0 || up.LastError != "" {
		t.Fatalf("recovered state = %+v, want up with fails and error cleared on the transition", up)
	}
}

// TestProberStateReturnsCopies pins the Task 6 contract: the status snapshot
// serializes what State() returns while probes keep running, so the snapshot
// must never share a slice with the live table.
func TestProberStateReturnsCopies(t *testing.T) {
	backend, ok, _ := switchableBackend(t, http.StatusInternalServerError)
	ok.Store(true)

	p := newProberSet()
	defer p.Stop()
	addr := backend.Listener.Addr().String()
	p.Sync(map[string]probeSpec{"m.example.com": {
		Host:   "m.example.com",
		Addrs:  []string{addr},
		Health: provision.Health{IntervalMS: 20}.WithDefaults(),
	}})

	st := p.State()
	st["m.example.com"][0].Addr = "clobbered"
	delete(st, "m.example.com")
	if again := p.State()["m.example.com"]; len(again) != 1 || again[0].Addr != addr {
		t.Fatalf("State() = %v after a caller mutated its copy, want the live table intact", again)
	}
	live := p.Live("m.example.com")
	if len(live) != 1 {
		t.Fatalf("Live() = %v, want the single upstream", live)
	}
	live[0] = "clobbered"
	if again := p.Live("m.example.com"); again[0] != addr {
		t.Fatalf("Live() = %v after a caller mutated its copy, want %s", again, addr)
	}
}