package agentsock

import (
	"context"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

// fakeWorkers records every effect so a test can assert the ABSENCE of one.
// "the call was refused" is only half the property; "and nothing happened" is
// the half that catches an auth check placed after the action.
type fakeWorkers struct {
	mu       sync.Mutex
	tokens   map[string]string
	contacts []string
	due      map[string]int64
	dueErr   error
	slept    []string
	ready    map[string]string
}

func (f *fakeWorkers) WorkerToken(host string) string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.tokens[host]
}

func (f *fakeWorkers) NoteWorkerContact(host string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.contacts = append(f.contacts, host)
}

func (f *fakeWorkers) SetWorkerNextDue(host string, at int64) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.dueErr != nil {
		return f.dueErr
	}
	if f.due == nil {
		f.due = map[string]int64{}
	}
	f.due[host] = at
	return nil
}

func (f *fakeWorkers) SleepWorker(host string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.slept = append(f.slept, host)
}

func (f *fakeWorkers) NoteWorkerReady(host, version string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.ready == nil {
		f.ready = map[string]string{}
	}
	f.ready[host] = version
}

// effects renders every recorded side effect, so a test can assert a refused
// call changed NOTHING rather than checking one slice and missing the rest.
func (f *fakeWorkers) effects() (contacts, slept int, due map[string]int64, ready map[string]string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	due = map[string]int64{}
	for k, v := range f.due {
		due[k] = v
	}
	ready = map[string]string{}
	for k, v := range f.ready {
		ready[k] = v
	}
	return len(f.contacts), len(f.slept), due, ready
}

func (f *fakeWorkers) assertNoEffects(t *testing.T, why string) {
	t.Helper()
	contacts, slept, due, ready := f.effects()
	if contacts != 0 || slept != 0 || len(due) != 0 || len(ready) != 0 {
		t.Fatalf("%s: a refused call must have no effect at all; contacts=%d slept=%d due=%v ready=%v",
			why, contacts, slept, due, ready)
	}
}

func post(h http.Handler, path, host, tok, body string) *httptest.ResponseRecorder {
	req := httptest.NewRequest("POST", path, strings.NewReader(body))
	if host != "" {
		req.Header.Set("Carlos-Sidecar-Host", host)
	}
	if tok != "" {
		req.Header.Set("Authorization", "Bearer "+tok)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func waitFor(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(2 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

var allVerbs = []string{"/v1/ready", "/v1/next-due", "/v1/idle", "/v1/heartbeat"}

func TestAgentSockRejectsWrongToken(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	h := New(f).Handler()

	for _, tc := range []struct{ name, host, tok string }{
		{"wrong token", "bot.x.internal", "bad"},
		{"no token", "bot.x.internal", ""},
		{"unknown host", "nope.x.internal", "good"},
		{"token for another host", "other.x.internal", "good"},
		{"no host header", "", "good"},
		{"empty-string token against an asleep host", "asleep.x.internal", ""},
	} {
		rec := post(h, "/v1/heartbeat", tc.host, tc.tok, "")
		if rec.Code != http.StatusUnauthorized {
			t.Fatalf("%s: code = %d, want 401", tc.name, rec.Code)
		}
	}
	f.assertNoEffects(t, "wrong-token sweep")
}

// The inversion this guards is total: WorkerToken("") means "unknown, asleep,
// or stopped", never "no token required". A caller who presents nothing must
// not match a host that holds nothing.
func TestAgentSockEmptyStoredTokenIsNotABypass(t *testing.T) {
	// "asleep" is present in the map with the empty token the real activator
	// returns for a host that is not currently awake or waking.
	f := &fakeWorkers{tokens: map[string]string{"asleep.x.internal": ""}}
	h := New(f).Handler()

	for _, tok := range []string{"", " ", "Bearer", "anything"} {
		for _, path := range allVerbs {
			body := ""
			switch path {
			case "/v1/ready":
				body = `{"version":"v1"}`
			case "/v1/next-due":
				body = `{"at":1756}`
			}
			rec := post(h, path, "asleep.x.internal", tok, body)
			if rec.Code != http.StatusUnauthorized {
				t.Fatalf("%s with token %q: code = %d, want 401", path, tok, rec.Code)
			}
		}
	}
	// The bare header form, which no post() helper shape produces: an
	// Authorization header whose credential is literally empty.
	req := httptest.NewRequest("POST", "/v1/heartbeat", nil)
	req.Header.Set("Carlos-Sidecar-Host", "asleep.x.internal")
	req.Header.Set("Authorization", "Bearer ")
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("empty bearer credential: code = %d, want 401", rec.Code)
	}
	f.assertNoEffects(t, "empty stored token")
}

// A caller that names no host is refused HERE, not by the accident that the
// activator's map happens to hold no "" key. This fake holds one deliberately:
// if the refusal were delegated to WorkerToken, an anonymous caller presenting
// this token would be authorized — and then act on host "", whatever that
// turns out to mean downstream.
func TestAgentSockAnonymousCallerIsRefusedHere(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"": "sekrit", "bot.x.internal": "good"}}
	h := New(f).Handler()
	for _, path := range allVerbs {
		body := ""
		switch path {
		case "/v1/ready":
			body = `{"version":"v1"}`
		case "/v1/next-due":
			body = `{"at":1756}`
		}
		if rec := post(h, path, "", "sekrit", body); rec.Code != http.StatusUnauthorized {
			t.Fatalf("%s with no host header: code = %d, want 401 — an unnamed caller has nothing to authorize", path, rec.Code)
		}
	}
	f.assertNoEffects(t, "anonymous caller")
}

func TestAgentSockTokenIsScopedToOneHost(t *testing.T) {
	// The security property: possession of A's token says nothing about B.
	f := &fakeWorkers{tokens: map[string]string{"a.x.internal": "tok-a", "b.x.internal": "tok-b"}}
	h := New(f).Handler()
	if rec := post(h, "/v1/idle", "b.x.internal", "tok-a", ""); rec.Code != http.StatusUnauthorized {
		t.Fatalf("A's token spoke about B: code = %d", rec.Code)
	}
	if rec := post(h, "/v1/ready", "b.x.internal", "tok-a", `{"version":"pwned"}`); rec.Code != http.StatusUnauthorized {
		t.Fatalf("A's token reported B's version: code = %d", rec.Code)
	}
	if rec := post(h, "/v1/next-due", "b.x.internal", "tok-a", `{"at":1}`); rec.Code != http.StatusUnauthorized {
		t.Fatalf("A's token rescheduled B: code = %d", rec.Code)
	}
	if rec := post(h, "/v1/heartbeat", "b.x.internal", "tok-a", ""); rec.Code != http.StatusUnauthorized {
		t.Fatalf("A's token kept B alive: code = %d", rec.Code)
	}
	f.assertNoEffects(t, "cross-host token")
}

// Every refusal reason must be indistinguishable on the wire: status, headers
// worth reading, and body. A caller must not be able to probe "does host X
// exist" or "is host X awake" by reading the difference.
func TestAgentSockRefusalsAreIndistinguishable(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{
		"awake.x.internal":  "good",
		"asleep.x.internal": "",
	}}
	h := New(f).Handler()

	type shape struct {
		code        int
		ctype, body string
	}
	var seen []shape
	for _, tc := range []struct{ name, host, tok string }{
		{"wrong token", "awake.x.internal", "bad"},
		{"wrong token, same length", "awake.x.internal", "goo1"},
		{"absent token", "awake.x.internal", ""},
		{"unknown host", "ghost.x.internal", "good"},
		{"asleep host", "asleep.x.internal", "good"},
		{"absent host header", "", "good"},
	} {
		rec := post(h, "/v1/heartbeat", tc.host, tc.tok, "")
		seen = append(seen, shape{rec.Code, rec.Header().Get("Content-Type"), rec.Body.String()})
		if rec.Code != http.StatusUnauthorized {
			t.Fatalf("%s: code = %d, want 401", tc.name, rec.Code)
		}
	}
	for i := 1; i < len(seen); i++ {
		if seen[i] != seen[0] {
			t.Fatalf("refusals are distinguishable: %+v vs %+v", seen[0], seen[i])
		}
	}
	f.assertNoEffects(t, "indistinguishable refusals")
}

func TestAgentSockVerbs(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	h := New(f).Handler()

	if rec := post(h, "/v1/ready", "bot.x.internal", "good", `{"version":"abc123"}`); rec.Code != http.StatusNoContent {
		t.Fatalf("ready: %d", rec.Code)
	}
	if f.ready["bot.x.internal"] != "abc123" {
		t.Fatal("ready must record the running version — it is a sidecar's only version signal")
	}
	if rec := post(h, "/v1/next-due", "bot.x.internal", "good", `{"at":1756}`); rec.Code != http.StatusNoContent {
		t.Fatalf("next-due: %d", rec.Code)
	}
	if f.due["bot.x.internal"] != 1756 {
		t.Fatalf("due = %d", f.due["bot.x.internal"])
	}
	if rec := post(h, "/v1/idle", "bot.x.internal", "good", ""); rec.Code != http.StatusNoContent {
		t.Fatalf("idle: %d", rec.Code)
	}
	if len(f.slept) != 1 {
		t.Fatal("idle must ask for a stop")
	}
	if rec := post(h, "/v1/heartbeat", "bot.x.internal", "good", ""); rec.Code != http.StatusNoContent {
		t.Fatalf("heartbeat: %d", rec.Code)
	}
	if len(f.contacts) != 1 || f.contacts[0] != "bot.x.internal" {
		t.Fatalf("heartbeat must stamp contact for the CALLER's host; contacts = %v", f.contacts)
	}
}

// The verb acts on the host the TOKEN authorized, never on a host named
// anywhere in the body — there is no host field in any body, and adding one
// would be the confused-deputy bug this pins.
func TestAgentSockActsOnlyOnTheAuthorizedHost(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"a.x.internal": "tok-a", "b.x.internal": "tok-b"}}
	h := New(f).Handler()

	if rec := post(h, "/v1/next-due", "a.x.internal", "tok-a", `{"at":99,"host":"b.x.internal"}`); rec.Code != http.StatusNoContent {
		t.Fatalf("next-due: %d", rec.Code)
	}
	if _, ok := f.due["b.x.internal"]; ok {
		t.Fatal("a body field must never redirect an effect onto another host")
	}
	if f.due["a.x.internal"] != 99 {
		t.Fatalf("due for the authorized host = %d, want 99", f.due["a.x.internal"])
	}
}

func TestAgentSockRefusesGarbage(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	h := New(f).Handler()
	for _, body := range []string{"not json", `{"at":"soon"}`, strings.Repeat("x", 1<<20), `{"at":-1}`} {
		rec := post(h, "/v1/next-due", "bot.x.internal", "good", body)
		if rec.Code != http.StatusBadRequest {
			t.Fatalf("body %.20q: code = %d, want 400", body, rec.Code)
		}
	}
	// An oversized body on a JSON-bearing verb must be refused, not truncated
	// into something that happens to parse.
	big := `{"version":"` + strings.Repeat("v", 1<<20) + `"}`
	if rec := post(h, "/v1/ready", "bot.x.internal", "good", big); rec.Code != http.StatusBadRequest {
		t.Fatalf("oversized ready: code = %d, want 400", rec.Code)
	}
	// The SIZE CAP itself, not the field bounds behind it: a body that is
	// perfectly well-formed and whose fields are all in range, but megabytes
	// long. Without a bounded reader this parses and acts.
	pad := strings.Repeat("x", 1<<20)
	if rec := post(h, "/v1/next-due", "bot.x.internal", "good", `{"at":1756,"pad":"`+pad+`"}`); rec.Code != http.StatusBadRequest {
		t.Fatalf("padded next-due: code = %d, want 400 — the body must be read under a cap", rec.Code)
	}
	if rec := post(h, "/v1/ready", "bot.x.internal", "good", `{"version":"abc","pad":"`+pad+`"}`); rec.Code != http.StatusBadRequest {
		t.Fatalf("padded ready: code = %d, want 400 — the body must be read under a cap", rec.Code)
	}
	for _, body := range []string{"not json", `{"version":42}`, "{\"version\":\"a\\u0000b\"}"} {
		if rec := post(h, "/v1/ready", "bot.x.internal", "good", body); rec.Code != http.StatusBadRequest {
			t.Fatalf("ready body %.20q: code = %d, want 400", body, rec.Code)
		}
	}
	if _, ok := f.due["bot.x.internal"]; ok {
		t.Fatal("a malformed body must leave no effect")
	}
	if len(f.ready) != 0 {
		t.Fatalf("a malformed ready must leave no effect: %v", f.ready)
	}

	if rec := post(h, "/v1/nope", "bot.x.internal", "good", ""); rec.Code != http.StatusNotFound {
		t.Fatalf("unknown verb: %d", rec.Code)
	}
	req := httptest.NewRequest("GET", "/v1/idle", nil)
	// GET on a POST-only verb
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusMethodNotAllowed && rec.Code != http.StatusUnauthorized {
		t.Fatalf("GET: %d", rec.Code)
	}
}

// A GET must not act even when it carries a perfectly good credential — the
// method gate has to come before the effect, not merely before the 200.
func TestAgentSockGetNeverActs(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	h := New(f).Handler()
	for _, path := range allVerbs {
		for _, method := range []string{"GET", "HEAD", "PUT", "DELETE"} {
			req := httptest.NewRequest(method, path, nil)
			req.Header.Set("Carlos-Sidecar-Host", "bot.x.internal")
			req.Header.Set("Authorization", "Bearer good")
			rec := httptest.NewRecorder()
			h.ServeHTTP(rec, req)
			if rec.Code == http.StatusNoContent {
				t.Fatalf("%s %s succeeded; only POST may act", method, path)
			}
		}
	}
	f.assertNoEffects(t, "non-POST methods")
}

// An empty next-due body is malformed: "at" is the whole message, and a
// zero-valued default would silently reschedule the worker to the epoch.
func TestAgentSockNextDueRequiresAt(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	h := New(f).Handler()
	for _, body := range []string{"", "{}", "null"} {
		if rec := post(h, "/v1/next-due", "bot.x.internal", "good", body); rec.Code != http.StatusBadRequest {
			t.Fatalf("next-due body %q: code = %d, want 400", body, rec.Code)
		}
	}
	if len(f.due) != 0 {
		t.Fatalf("due must be untouched: %v", f.due)
	}
}

// The version field must be PRESENT — a typo'd body is a loud 400, never a
// silent "unknown" that also ends the settle window. It may be empty, which is
// a child honestly saying it does not know its own build.
func TestAgentSockReadyRequiresTheVersionField(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	h := New(f).Handler()
	for _, body := range []string{"", "{}", "null", `{"verison":"abc"}`} {
		if rec := post(h, "/v1/ready", "bot.x.internal", "good", body); rec.Code != http.StatusBadRequest {
			t.Fatalf("ready body %q: code = %d, want 400", body, rec.Code)
		}
	}
	if len(f.ready) != 0 {
		t.Fatalf("ready must be untouched: %v", f.ready)
	}
	if rec := post(h, "/v1/ready", "bot.x.internal", "good", `{"version":""}`); rec.Code != http.StatusNoContent {
		t.Fatalf("explicitly-unknown version: code = %d, want 204", rec.Code)
	}
	if v, ok := f.ready["bot.x.internal"]; !ok || v != "" {
		t.Fatalf("ready = %q, %v; an explicit empty version is honest, and still ready", v, ok)
	}
}

// Bodyless verbs stay bodyless: idle and heartbeat carry no arguments, so
// nothing a caller sends can steer them.
func TestAgentSockBodylessVerbsIgnoreTheirBody(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	h := New(f).Handler()
	if rec := post(h, "/v1/heartbeat", "bot.x.internal", "good", `{"host":"other.x.internal"}`); rec.Code != http.StatusNoContent {
		t.Fatalf("heartbeat with a body: %d", rec.Code)
	}
	if len(f.contacts) != 1 || f.contacts[0] != "bot.x.internal" {
		t.Fatalf("contacts = %v", f.contacts)
	}
}

func TestAgentSockNextDueStoreFailureIsNot204(t *testing.T) {
	f := &fakeWorkers{
		tokens: map[string]string{"bot.x.internal": "good"},
		dueErr: errors.New("index write failed"),
	}
	h := New(f).Handler()
	rec := post(h, "/v1/next-due", "bot.x.internal", "good", `{"at":1756}`)
	if rec.Code != http.StatusInternalServerError {
		t.Fatalf("code = %d, want 500 — a child must learn its schedule did not land", rec.Code)
	}
}

func TestAgentSockServeCleansUpItsSocket(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "agent.sock")
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- New(&fakeWorkers{}).Serve(ctx, path) }()
	waitFor(t, "socket to appear", func() bool { _, err := os.Stat(path); return err == nil })
	cancel()
	if err := <-done; err != nil {
		t.Fatalf("Serve returned %v on a clean shutdown", err)
	}
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		t.Fatal("a stale socket file must not survive shutdown — the next start would fail to bind")
	}
}

// A hard kill leaves the socket file behind. The next start must take it over,
// not die on "address already in use".
func TestAgentSockServeReplacesAStaleSocketFile(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "agent.sock")
	if err := os.WriteFile(path, []byte("stale"), 0o600); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- New(&fakeWorkers{}).Serve(ctx, path) }()
	waitFor(t, "socket to appear", func() bool {
		c, err := net.Dial("unix", path)
		if err != nil {
			return false
		}
		c.Close()
		return true
	})
	cancel()
	<-done
}

// The other half of stale-socket removal: a socket somebody is LISTENING on is
// not stale. Unlinking it would orphan the live agent — its children would go
// on talking to a process nothing else can reach — so refuse instead.
func TestAgentSockServeRefusesALiveSocket(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "agent.sock")
	ln, err := net.Listen("unix", path)
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			c.Close()
		}
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	if err := New(&fakeWorkers{}).Serve(ctx, path); err == nil {
		t.Fatal("Serve must refuse a socket another listener already owns")
	}
	if _, err := os.Stat(path); err != nil {
		t.Fatalf("the live socket was unlinked: %v", err)
	}
}

// End to end over a real unix socket: the shape a sidecar actually dials.
func TestAgentSockOverTheRealSocket(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "agent.sock")
	f := &fakeWorkers{tokens: map[string]string{"bot.x.internal": "good"}}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- New(f).Serve(ctx, path) }()
	defer func() { cancel(); <-done }()
	waitFor(t, "socket to appear", func() bool { _, err := os.Stat(path); return err == nil })

	cl := &http.Client{
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
				return (&net.Dialer{}).DialContext(ctx, "unix", path)
			},
		},
		Timeout: 5 * time.Second,
	}
	call := func(tok string) int {
		req, err := http.NewRequest("POST", "http://agent/v1/heartbeat", nil)
		if err != nil {
			t.Fatal(err)
		}
		req.Header.Set("Carlos-Sidecar-Host", "bot.x.internal")
		req.Header.Set("Authorization", "Bearer "+tok)
		resp, err := cl.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		resp.Body.Close()
		return resp.StatusCode
	}
	if code := call("good"); code != http.StatusNoContent {
		t.Fatalf("live heartbeat: %d", code)
	}
	if code := call("bad"); code != http.StatusUnauthorized {
		t.Fatalf("live bad token: %d", code)
	}

	// The mode is checked only AFTER a request has been answered, and the
	// ordering is the whole point: Serve chmods before it starts serving, so
	// a completed response is proof the chmod has happened. Polling on the
	// file merely EXISTING returns inside the listen→chmod window, where an
	// unlucky read sees whatever the umask left (0755 under the common 022)
	// and fails a correct implementation.
	if fi, err := os.Stat(path); err != nil {
		t.Fatal(err)
	} else if perm := fi.Mode().Perm(); perm&0o007 != 0 {
		t.Fatalf("socket mode %o is reachable by other — the token is the boundary, but the mode should not widen the audience", perm)
	}
	contacts, _, _, _ := f.effects()
	if contacts != 1 {
		t.Fatalf("contacts = %d, want exactly the authorized one", contacts)
	}
}

// Many children hammering one agent at once — the handler holds no per-request
// state, and -race proves it.
func TestAgentSockConcurrentCallers(t *testing.T) {
	f := &fakeWorkers{tokens: map[string]string{"a.x.internal": "tok-a", "b.x.internal": "tok-b"}}
	h := New(f).Handler()
	var wg sync.WaitGroup
	for i := range 50 {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			if i%2 == 0 {
				post(h, "/v1/heartbeat", "a.x.internal", "tok-a", "")
			} else {
				post(h, "/v1/heartbeat", "b.x.internal", "tok-a", "")
			}
		}(i)
	}
	wg.Wait()
	contacts, _, _, _ := f.effects()
	if contacts != 25 {
		t.Fatalf("contacts = %d, want 25 — only a.x.internal's own token may stamp it", contacts)
	}
}