package main

// Item 25: agent session tokens. After one scrypt-verified key auth the
// panel issues a random session token (only its SHA-256 hash is stored);
// subsequent requests authenticate by hash compare. These tests drive
// the real /api/agent/heartbeat route through the smoke harness.

import (
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"
	"time"
)

const testAgentKey = "test-agent-key-0123456789abcdef-XYZ"

func seedAgentApp(t *testing.T, key string) (*App, string) {
	t.Helper()
	cookieSecure = false
	store := NewStore(filepath.Join(t.TempDir(), "nexusgate.json"))
	if err := store.Init(); err != nil {
		t.Fatalf("store init: %v", err)
	}
	fp := TokenFingerprint(key)
	agentID := ""
	_, err := store.Transaction(func(d *Data) (any, error) {
		d.Servers = append(d.Servers, &Server{
			ID: "srv_tok", Name: "tok", Status: "online",
			CreatedAt: NowISO(), UpdatedAt: NowISO(),
		})
		agent := &Agent{
			ID:             NewID("agt"),
			ServerID:       "srv_tok",
			KeyHash:        HashSecret(key),
			KeyFingerprint: &fp,
			Status:         "active",
			CreatedAt:      NowISO(),
			LastSeenAt:     "2000-01-01T00:00:00.000Z",
		}
		agentID = agent.ID
		d.Agents = append(d.Agents, agent)
		return nil, nil
	})
	if err != nil {
		t.Fatalf("seed agent: %v", err)
	}
	return NewApp(store, NewSessionManager(12), NewOrchestrator(store)), agentID
}

func agentHeartbeat(t *testing.T, app *App, bearer string) *httptest.ResponseRecorder {
	t.Helper()
	return doRequest(app, http.MethodPost, "/api/agent/heartbeat", "{}", map[string]string{
		"Authorization": "Bearer " + bearer,
	})
}

func storedAgent(t *testing.T, app *App, agentID string) *Agent {
	t.Helper()
	var out *Agent
	app.store.View(func(d *Data) {
		for _, agent := range d.Agents {
			if agent.ID == agentID {
				copied := *agent
				out = &copied
			}
		}
	})
	if out == nil {
		t.Fatalf("agent %s not in store", agentID)
	}
	return out
}

func mutateAgent(t *testing.T, app *App, agentID string, fn func(*Agent)) {
	t.Helper()
	_, err := app.store.Transaction(func(d *Data) (any, error) {
		for _, agent := range d.Agents {
			if agent.ID == agentID {
				fn(agent)
			}
		}
		return nil, nil
	})
	if err != nil {
		t.Fatalf("mutate agent: %v", err)
	}
}

// Key auth succeeds and issues a session token whose hash (and only the
// hash) lands on the agent record.
func TestAgentSessionTokenIssuedOnKeyAuth(t *testing.T) {
	app, agentID := seedAgentApp(t, testAgentKey)
	rec := agentHeartbeat(t, app, testAgentKey)
	if rec.Code != 200 {
		t.Fatalf("key heartbeat: got %d (%s)", rec.Code, rec.Body.String())
	}
	token := rec.Header().Get(agentSessionTokenHeader)
	expires := rec.Header().Get(agentSessionTokenExpiresHeader)
	if token == "" || expires == "" {
		t.Fatalf("no session token issued: token=%q expires=%q", token, expires)
	}
	if token == testAgentKey {
		t.Fatal("session token must differ from the agent key")
	}
	stored := storedAgent(t, app, agentID)
	if stored.SessionTokenHash == nil || *stored.SessionTokenHash != TokenFingerprint(token) {
		t.Fatalf("stored hash does not match issued token: %+v", stored.SessionTokenHash)
	}
	if stored.SessionTokenExpiresAt == nil || *stored.SessionTokenExpiresAt != expires {
		t.Fatalf("stored expiry %v != header expiry %q", stored.SessionTokenExpiresAt, expires)
	}
	expMs, ok := ParseISOMs(expires)
	if !ok || expMs <= time.Now().UnixMilli() {
		t.Fatalf("issued expiry not in the future: %q", expires)
	}
}

// Token auth succeeds even when the key hash can no longer verify —
// proving the token path never touches scrypt — and does not re-issue.
func TestAgentSessionTokenAuthSkipsScrypt(t *testing.T) {
	app, agentID := seedAgentApp(t, testAgentKey)
	token := agentHeartbeat(t, app, testAgentKey).Header().Get(agentSessionTokenHeader)
	if token == "" {
		t.Fatal("setup: no token issued")
	}
	// Sabotage the key hash: any scrypt verification now fails.
	mutateAgent(t, app, agentID, func(agent *Agent) {
		agent.KeyHash = "scrypt:deadbeef:abcd"
	})
	rec := agentHeartbeat(t, app, token)
	if rec.Code != 200 {
		t.Fatalf("token heartbeat with sabotaged key hash: got %d", rec.Code)
	}
	if again := rec.Header().Get(agentSessionTokenHeader); again != "" {
		t.Fatalf("token path must not re-issue, got %q", again)
	}
	// And the (now unverifiable) key is rejected.
	if rec := agentHeartbeat(t, app, testAgentKey); rec.Code != 401 {
		t.Fatalf("sabotaged key: got %d, want 401", rec.Code)
	}
}

// An expired token is rejected; the key still works and re-issues a
// fresh token (rotation).
func TestAgentSessionTokenExpiryFallsBackToKey(t *testing.T) {
	app, agentID := seedAgentApp(t, testAgentKey)
	token := agentHeartbeat(t, app, testAgentKey).Header().Get(agentSessionTokenHeader)
	mutateAgent(t, app, agentID, func(agent *Agent) {
		past := ISOFromTime(time.Now().Add(-time.Hour))
		agent.SessionTokenExpiresAt = &past
	})
	if rec := agentHeartbeat(t, app, token); rec.Code != 401 {
		t.Fatalf("expired token: got %d, want 401", rec.Code)
	}
	rec := agentHeartbeat(t, app, testAgentKey)
	if rec.Code != 200 {
		t.Fatalf("key after expiry: got %d", rec.Code)
	}
	fresh := rec.Header().Get(agentSessionTokenHeader)
	if fresh == "" || fresh == token {
		t.Fatalf("expected a fresh token after expiry, got %q", fresh)
	}
	if rec := agentHeartbeat(t, app, fresh); rec.Code != 200 {
		t.Fatalf("fresh token: got %d", rec.Code)
	}
}

// A token whose stored hash no longer matches (tampered / rotated
// server-side) is rejected; the key re-issues.
func TestAgentSessionTokenHashMismatchFallsBack(t *testing.T) {
	app, agentID := seedAgentApp(t, testAgentKey)
	token := agentHeartbeat(t, app, testAgentKey).Header().Get(agentSessionTokenHeader)
	mutateAgent(t, app, agentID, func(agent *Agent) {
		wrong := TokenFingerprint("some-other-token")
		agent.SessionTokenHash = &wrong
	})
	if rec := agentHeartbeat(t, app, token); rec.Code != 401 {
		t.Fatalf("mismatched token: got %d, want 401", rec.Code)
	}
	rec := agentHeartbeat(t, app, testAgentKey)
	if rec.Code != 200 || rec.Header().Get(agentSessionTokenHeader) == "" {
		t.Fatalf("key re-issue after mismatch: code=%d token=%q", rec.Code, rec.Header().Get(agentSessionTokenHeader))
	}
}

// Revoking the agent kills both the token and the key immediately, and
// Revoke() clears the stored token material.
func TestAgentSessionTokenRevokedAgentRejected(t *testing.T) {
	app, agentID := seedAgentApp(t, testAgentKey)
	token := agentHeartbeat(t, app, testAgentKey).Header().Get(agentSessionTokenHeader)
	mutateAgent(t, app, agentID, func(agent *Agent) { agent.Revoke() })
	stored := storedAgent(t, app, agentID)
	if stored.SessionTokenHash != nil || stored.SessionTokenExpiresAt != nil {
		t.Fatal("Revoke() must clear session token fields")
	}
	if rec := agentHeartbeat(t, app, token); rec.Code != 401 {
		t.Fatalf("revoked agent token: got %d, want 401", rec.Code)
	}
	if rec := agentHeartbeat(t, app, testAgentKey); rec.Code != 401 {
		t.Fatalf("revoked agent key: got %d, want 401", rec.Code)
	}
}

// The token is bound to its own agent record: it must not authenticate
// a request as a different agent (observed via which record's
// LastSeenAt the heartbeat advances).
func TestAgentSessionTokenBoundToRecord(t *testing.T) {
	app, agentID := seedAgentApp(t, testAgentKey)
	otherKey := "other-agent-key-9876543210"
	otherFP := TokenFingerprint(otherKey)
	otherID := ""
	_, err := app.store.Transaction(func(d *Data) (any, error) {
		other := &Agent{
			ID: NewID("agt"), ServerID: "srv_tok",
			KeyHash: HashSecret(otherKey), KeyFingerprint: &otherFP,
			Status: "active", CreatedAt: NowISO(),
			LastSeenAt: "2000-01-01T00:00:00.000Z",
		}
		otherID = other.ID
		d.Agents = append(d.Agents, other)
		return nil, nil
	})
	if err != nil {
		t.Fatalf("seed second agent: %v", err)
	}
	token := agentHeartbeat(t, app, testAgentKey).Header().Get(agentSessionTokenHeader)
	if rec := agentHeartbeat(t, app, token); rec.Code != 200 {
		t.Fatalf("token heartbeat: got %d", rec.Code)
	}
	if got := storedAgent(t, app, agentID).LastSeenAt; got == "2000-01-01T00:00:00.000Z" {
		t.Fatal("token heartbeat did not advance its own agent's LastSeenAt")
	}
	if got := storedAgent(t, app, otherID).LastSeenAt; got != "2000-01-01T00:00:00.000Z" {
		t.Fatalf("token heartbeat advanced the WRONG agent (LastSeenAt=%s)", got)
	}
}

// Key-only agents (pre-token clients) keep working request after
// request; each key auth rotates the stored token.
func TestAgentKeyOnlyRequestsKeepWorking(t *testing.T) {
	app, agentID := seedAgentApp(t, testAgentKey)
	var prevHash string
	for i := 0; i < 3; i++ {
		rec := agentHeartbeat(t, app, testAgentKey)
		if rec.Code != 200 {
			t.Fatalf("key heartbeat %d: got %d", i, rec.Code)
		}
		stored := storedAgent(t, app, agentID)
		if stored.SessionTokenHash == nil || *stored.SessionTokenHash == prevHash {
			t.Fatalf("key heartbeat %d did not rotate the token", i)
		}
		prevHash = *stored.SessionTokenHash
	}
}