Control-Center/go-backend/internal/gateway/wsclient_test.go

package gateway

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"code.cubecraftcreations.com/CubeCraft-Creations/Control-Center/go-backend/internal/handler"
	"code.cubecraftcreations.com/CubeCraft-Creations/Control-Center/go-backend/internal/models"

	"github.com/gorilla/websocket"
)

// ── Mock WebSocket server helper ─────────────────────────────────────────

// newTestWSServer creates an httptest.Server that upgrades to WebSocket and
// delegates each connection to handler. The server URL can be converted to
// a ws:// URL by replacing "http" with "ws".
func newTestWSServer(t *testing.T, handler func(conn *websocket.Conn)) *httptest.Server {
	t.Helper()
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		handler(conn)
	}))
	return srv
}

// wsURL converts an httptest.Server http URL to a ws URL.
func wsURL(srv *httptest.Server) string {
	return "ws" + strings.TrimPrefix(srv.URL, "http")
}

// ── Handshake helper for mock server ─────────────────────────────────────

// handleHandshake performs the server side of the v3 handshake:
//  1. Send connect.challenge
//  2. Read connect request
//  3. Send hello-ok response
//
// Returns the connect request frame for inspection.
func handleHandshake(t *testing.T, conn *websocket.Conn) map[string]any {
	t.Helper()

	// 1. Send connect.challenge
	challenge := map[string]any{
		"type":   "event",
		"event":  "connect.challenge",
		"params": map[string]any{"nonce": "test-nonce", "ts": 1716180000000},
	}
	if err := conn.WriteJSON(challenge); err != nil {
		t.Fatalf("server: write challenge: %v", err)
	}

	// 2. Read connect request
	var req map[string]any
	if err := conn.ReadJSON(&req); err != nil {
		t.Fatalf("server: read connect request: %v", err)
	}

	if req["method"] != "connect" {
		t.Fatalf("server: expected method=connect, got %v", req["method"])
	}

	// 3. Send hello-ok response
	// Note: helloOKResponse expects ConnID at the top level of the result,
	// matching the WSClient's JSON struct tags.
	result := map[string]any{
		"type":     "hello-ok",
		"protocol": 3,
		"connId":   "test-conn-123",
		"features": map[string]any{"methods": []string{}, "events": []string{}},
		"auth":     map[string]any{"role": "operator", "scopes": []string{"operator.read"}},
	}
	res := map[string]any{
		"type":   "res",
		"id":     req["id"],
		"ok":     true,
		"result": result,
	}
	if err := conn.WriteJSON(res); err != nil {
		t.Fatalf("server: write hello-ok: %v", err)
	}

	return req
}

// keepAlive reads frames from the connection until an error occurs
// (e.g., the client disconnects). Used as the default "do nothing"
// server loop after handshake.
func keepAlive(conn *websocket.Conn) {
	for {
		var m map[string]any
		if err := conn.ReadJSON(&m); err != nil {
			break
		}
	}
}

// ── 1. Test: Full handshake ──────────────────────────────────────────────

func TestWSClient_Handshake(t *testing.T) {
	srv := newTestWSServer(t, func(conn *websocket.Conn) {
		handleHandshake(t, conn)
		keepAlive(conn)
	})
	defer srv.Close()

	client := NewWSClient(WSConfig{URL: wsURL(srv), AuthToken: "test-token"}, nil, nil, slog.Default())

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	done := make(chan struct{})
	go func() {
		client.Start(ctx)
		close(done)
	}()

	// Wait briefly for handshake to complete
	time.Sleep(200 * time.Millisecond)

	// Verify connId was set
	client.connMu.Lock()
	connID := client.connId
	client.connMu.Unlock()

	if connID != "test-conn-123" {
		t.Errorf("expected connId 'test-conn-123', got %q", connID)
	}

	cancel()
	select {
	case <-done:
		// Client exited cleanly
	case <-time.After(3 * time.Second):
		t.Fatal("WSClient did not shut down after context cancellation")
	}
}

// ── 2. Test: Send() with response matching ───────────────────────────────

func TestWSClient_Send(t *testing.T) {
	srv := newTestWSServer(t, func(conn *websocket.Conn) {
		handleHandshake(t, conn)

		// Read RPC requests and respond to each
		for {
			var req map[string]any
			if err := conn.ReadJSON(&req); err != nil {
				break
			}
			reqID, _ := req["id"].(string)
			method, _ := req["method"].(string)

			var result any
			switch method {
			case "agents.list":
				result = map[string]any{
					"agents": []map[string]any{
						{"id": "otto", "name": "Otto"},
					},
				}
			default:
				result = map[string]any{}
			}

			res := map[string]any{
				"type":   "res",
				"id":     reqID,
				"ok":     true,
				"result": result,
			}
			if err := conn.WriteJSON(res); err != nil {
				break
			}
		}
	})
	defer srv.Close()

	client := NewWSClient(WSConfig{URL: wsURL(srv), AuthToken: "test-token"}, nil, nil, slog.Default())

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	go client.Start(ctx)

	// Give the client time to complete handshake
	time.Sleep(300 * time.Millisecond)

	resp, err := client.Send("agents.list", nil)
	if err != nil {
		t.Fatalf("Send() returned error: %v", err)
	}

	// Verify the response payload
	var result map[string]any
	if err := json.Unmarshal(resp, &result); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}

	agents, ok := result["agents"].([]any)
	if !ok || len(agents) != 1 {
		t.Errorf("expected 1 agent in response, got %v", result)
	}

	cancel()
}

// ── 3. Test: Event handler routing ───────────────────────────────────────

func TestWSClient_EventRouting(t *testing.T) {
	eventReceived := make(chan json.RawMessage, 1)

	srv := newTestWSServer(t, func(conn *websocket.Conn) {
		handleHandshake(t, conn)

		// After handshake, send a test event
		evt := map[string]any{
			"type":   "event",
			"event":  "test.event",
			"params": map[string]any{"greeting": "hello from server"},
		}
		if err := conn.WriteJSON(evt); err != nil {
			t.Logf("server: write event: %v", err)
			return
		}

		keepAlive(conn)
	})
	defer srv.Close()

	client := NewWSClient(WSConfig{URL: wsURL(srv), AuthToken: "test-token"}, nil, nil, slog.Default())

	// Register event handler BEFORE starting the client
	client.OnEvent("test.event", func(payload json.RawMessage) {
		eventReceived <- payload
	})

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	go client.Start(ctx)

	// Wait for the event handler to fire
	select {
	case payload := <-eventReceived:
		var data map[string]any
		if err := json.Unmarshal(payload, &data); err != nil {
			t.Fatalf("unmarshal event payload: %v", err)
		}
		if greeting, _ := data["greeting"].(string); greeting != "hello from server" {
			t.Errorf("expected greeting 'hello from server', got %q", greeting)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for event handler to fire")
	}

	cancel()
}

// ── 4. Test: Concurrent Send ─────────────────────────────────────────────

func TestWSClient_ConcurrentSend(t *testing.T) {
	var reqCount atomic.Int32

	srv := newTestWSServer(t, func(conn *websocket.Conn) {
		handleHandshake(t, conn)

		// Read RPC requests and respond to each
		for {
			var req map[string]any
			if err := conn.ReadJSON(&req); err != nil {
				break
			}
			reqID, _ := req["id"].(string)
			n := reqCount.Add(1)

			res := map[string]any{
				"type":   "res",
				"id":     reqID,
				"ok":     true,
				"result": map[string]any{"index": n, "method": req["method"]},
			}
			if err := conn.WriteJSON(res); err != nil {
				break
			}
		}
	})
	defer srv.Close()

	client := NewWSClient(WSConfig{URL: wsURL(srv), AuthToken: "test-token"}, nil, nil, slog.Default())

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	go client.Start(ctx)

	// Give the client time to complete handshake
	time.Sleep(300 * time.Millisecond)

	// Fire 3 concurrent Send() calls
	type sendResult struct {
		method  string
		payload json.RawMessage
		err     error
	}
	results := make(chan sendResult, 3)

	methods := []string{"agents.list", "sessions.list", "agents.config"}
	for _, method := range methods {
		go func(m string) {
			resp, err := client.Send(m, nil)
			results <- sendResult{method: m, payload: resp, err: err}
		}(method)
	}

	// Collect all results
	for i := 0; i < 3; i++ {
		select {
		case r := <-results:
			if r.err != nil {
				t.Errorf("Send(%q) returned error: %v", r.method, r.err)
				continue
			}
			var result map[string]any
			if err := json.Unmarshal(r.payload, &result); err != nil {
				t.Errorf("Send(%q) unmarshal error: %v", r.method, err)
				continue
			}
			gotMethod, _ := result["method"].(string)
			if gotMethod != r.method {
				t.Errorf("Send(%q) got response for %q (mismatched)", r.method, gotMethod)
			}
		case <-time.After(5 * time.Second):
			t.Fatal("timed out waiting for concurrent Send results")
		}
	}

	cancel()
}

// ── 5. Test: Clean shutdown ──────────────────────────────────────────────

func TestWSClient_CleanShutdown(t *testing.T) {
	srv := newTestWSServer(t, func(conn *websocket.Conn) {
		handleHandshake(t, conn)
		keepAlive(conn)
	})
	defer srv.Close()

	client := NewWSClient(WSConfig{URL: wsURL(srv), AuthToken: "test-token"}, nil, nil, slog.Default())

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

	done := make(chan struct{})
	go func() {
		client.Start(ctx)
		close(done)
	}()

	// Let the client connect and complete handshake
	time.Sleep(200 * time.Millisecond)

	// Cancel context — should trigger clean shutdown
	cancel()

	select {
	case <-done:
		// Client exited cleanly — pass
	case <-time.After(3 * time.Second):
		t.Fatal("WSClient did not shut down cleanly within timeout")
	}
}

// ── Pure utility tests (from CUB-205) ─────────────────────────────────────

func TestMapSessionStatus(t *testing.T) {
	tests := []struct {
		input    string
		expected models.AgentStatus
	}{
		{"running", models.AgentStatusActive},
		{"streaming", models.AgentStatusActive},
		{"done", models.AgentStatusIdle},
		{"error", models.AgentStatusError},
		{"", models.AgentStatusIdle},
		{"garbage", models.AgentStatusIdle},
	}
	for _, tt := range tests {
		result := mapSessionStatus(tt.input)
		if result != tt.expected {
			t.Errorf("mapSessionStatus(%q) = %q, want %q", tt.input, result, tt.expected)
		}
	}
}

func TestAgentItemToCard(t *testing.T) {
	t.Run("full fields", func(t *testing.T) {
		item := agentListItem{
			ID:      "dex",
			Name:    "Dex",
			Role:    "backend",
			Channel: "telegram",
		}
		card := agentItemToCard(item)
		if card.ID != "dex" {
			t.Errorf("ID = %q, want %q", card.ID, "dex")
		}
		if card.DisplayName != "Dex" {
			t.Errorf("DisplayName = %q, want %q", card.DisplayName, "Dex")
		}
		if card.Role != "backend" {
			t.Errorf("Role = %q, want %q", card.Role, "backend")
		}
		if card.Channel != "telegram" {
			t.Errorf("Channel = %q, want %q", card.Channel, "telegram")
		}
		if card.Status != models.AgentStatusIdle {
			t.Errorf("Status = %q, want %q", card.Status, models.AgentStatusIdle)
		}
	})

	t.Run("empty fields use defaults", func(t *testing.T) {
		item := agentListItem{
			ID: "otto",
		}
		card := agentItemToCard(item)
		if card.ID != "otto" {
			t.Errorf("ID = %q, want %q", card.ID, "otto")
		}
		if card.DisplayName != "otto" {
			t.Errorf("DisplayName = %q, want %q (should fallback to ID)", card.DisplayName, "otto")
		}
		if card.Role != "agent" {
			t.Errorf("Role = %q, want %q (default)", card.Role, "agent")
		}
		if card.Channel != "unknown" {
			t.Errorf("Channel = %q, want %q (per Grimm requirement)", card.Channel, "unknown")
		}
		if card.Status != models.AgentStatusIdle {
			t.Errorf("Status = %q, want %q", card.Status, models.AgentStatusIdle)
		}
	})

	t.Run("empty name falls back to ID", func(t *testing.T) {
		item := agentListItem{
			ID:   "hex",
			Name: "",
			Role: "database",
		}
		card := agentItemToCard(item)
		if card.DisplayName != "hex" {
			t.Errorf("DisplayName = %q, want %q (ID fallback)", card.DisplayName, "hex")
		}
	})
}

// ── 6. Test: Initial sync ordering (readLoop active before Send) ──────────

// TestConnectAndRun_InitialSyncOrdering verifies that the WS client
// completes initial sync successfully. This test would hang/timeout if
// readLoop were NOT started before initialSync, because Send() relies on
// readLoop→routeFrame→handleResponse to deliver RPC responses.
func TestConnectAndRun_InitialSyncOrdering(t *testing.T) {
	repo := &mockAgentRepo{agents: make(map[string]models.AgentCardData)}
	broker := handler.NewBroker()
	capture := newBroadcastCapture(broker)
	defer capture.close()

	srv := newTestWSServer(t, func(conn *websocket.Conn) {
		// Handshake
		handleHandshake(t, conn)

		// After handshake, respond to RPCs
		for {
			var req map[string]any
			if err := conn.ReadJSON(&req); err != nil {
				break
			}
			reqID, _ := req["id"].(string)
			method, _ := req["method"].(string)

			var result any
			switch method {
			case "agents.list":
				result = []map[string]any{
					{"id": "otto", "name": "Otto", "role": "Orchestrator", "channel": "discord"},
					{"id": "dex", "name": "Dex", "role": "Backend Dev", "channel": "telegram"},
				}
			case "sessions.list":
				result = []map[string]any{
					{"sessionKey": "s1", "agentId": "otto", "status": "running", "totalTokens": 500, "lastActivityAt": "2025-05-20T12:00:00Z"},
				}
			default:
				result = map[string]any{}
			}

			res := map[string]any{
				"type":   "res",
				"id":     reqID,
				"ok":     true,
				"result": result,
			}
			if err := conn.WriteJSON(res); err != nil {
				break
			}
		}
	})
	defer srv.Close()

	client := NewWSClient(WSConfig{URL: wsURL(srv), AuthToken: "test-token"}, repo, broker, slog.Default())

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	done := make(chan struct{})
	go func() {
		client.Start(ctx)
		close(done)
	}()

	// Wait for initial sync to complete by checking repo state.
	// The agents should be persisted from the RPC responses.
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		repo.mu.Lock()
		_, ottoOK := repo.agents["otto"]
		_, dexOK := repo.agents["dex"]
		repo.mu.Unlock()
		if ottoOK && dexOK {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}

	repo.mu.Lock()
	_, ottoOK := repo.agents["otto"]
	_, dexOK := repo.agents["dex"]
	repo.mu.Unlock()

	if !ottoOK {
		t.Error("otto not found in repo after initial sync — readLoop may not have been active before Send()")
	}
	if !dexOK {
		t.Error("dex not found in repo after initial sync — readLoop may not have been active before Send()")
	}

	cancel()
	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("WSClient did not shut down cleanly")
	}
}

// ── 7. Test: Event not lost during initial sync (regression) ───────────────

// TestConnectAndRun_EventNotLostDuringSync verifies that live gateway events
// arriving during initial sync are NOT dropped. This is a regression test
// for the race where readLoop started before registerEventHandlers(),
// causing events read during that window to be logged as "unhandled" and lost.
//
// The mock server sends a live event (sessions.changed) right after the
// handshake, interleaved with the RPC responses for agents.list and
// sessions.list. The test asserts the event is received by the handler.
func TestConnectAndRun_EventNotLostDuringSync(t *testing.T) {
	repo := &mockAgentRepo{agents: make(map[string]models.AgentCardData)}
	broker := handler.NewBroker()
	capture := newBroadcastCapture(broker)
	defer capture.close()

	// Pre-seed an agent so the event handler can update it.
	repo.agents["otto"] = models.AgentCardData{
		ID:          "otto",
		DisplayName: "Otto",
		Status:      models.AgentStatusIdle,
	}

	srv := newTestWSServer(t, func(conn *websocket.Conn) {
		// Handshake
		handleHandshake(t, conn)

		// After handshake, process RPCs and inject a live event.
		for {
			var req map[string]any
			if err := conn.ReadJSON(&req); err != nil {
				break
			}
			reqID, _ := req["id"].(string)
			method, _ := req["method"].(string)

			// Respond to agents.list RPC
			if method == "agents.list" {
				// Before responding, inject a live event — simulates
				// a gateway pushing a presence update during sync.
				evt := map[string]any{
					"type":   "event",
					"event":  "presence",
					"params": map[string]any{"agentId": "otto", "connected": true, "lastActivityAt": "2025-05-20T12:30:00Z"},
				}
				if err := conn.WriteJSON(evt); err != nil {
					break
				}

				// Now send the RPC response
				res := map[string]any{
					"type":   "res",
					"id":     reqID,
					"ok":     true,
					"result": []map[string]any{
						{"id": "otto", "name": "Otto", "role": "Orchestrator", "channel": "discord"},
					},
				}
				if err := conn.WriteJSON(res); err != nil {
					break
				}
				continue
			}

			// Respond to sessions.list RPC
			if method == "sessions.list" {
				res := map[string]any{
					"type":   "res",
					"id":     reqID,
					"ok":     true,
					"result": []map[string]any{},
				}
				if err := conn.WriteJSON(res); err != nil {
					break
				}
				continue
			}

			// Default response for other methods
			res := map[string]any{
				"type":   "res",
				"id":     reqID,
				"ok":     true,
				"result": map[string]any{},
			}
			if err := conn.WriteJSON(res); err != nil {
				break
			}
		}
	})
	defer srv.Close()

	client := NewWSClient(WSConfig{URL: wsURL(srv), AuthToken: "test-token"}, repo, broker, slog.Default())

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	done := make(chan struct{})
	go func() {
		client.Start(ctx)
		close(done)
	}()

	// Wait for the presence event to be processed by checking the repo.
	// The presence handler updates the agent, so we check for the
	// lastActivityAt change.
	deadline := time.Now().Add(5 * time.Second)
	var lastActivity string
	for time.Now().Before(deadline) {
		repo.mu.Lock()
		if a, ok := repo.agents["otto"]; ok {
			lastActivity = a.LastActivity
		}
		repo.mu.Unlock()
		if lastActivity == "2025-05-20T12:30:00Z" {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}

	if lastActivity != "2025-05-20T12:30:00Z" {
		t.Errorf("presence event during sync was lost: lastActivity = %q, want %q", lastActivity, "2025-05-20T12:30:00Z")
	}

	cancel()
	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("WSClient did not shut down cleanly")
	}
}

func TestStrPtr(t *testing.T) {
	s := "hello"
	p := strPtr(s)
	if p == nil {
		t.Fatal("strPtr returned nil")
	}
	if *p != s {
		t.Errorf("strPtr(%q) = %q, want %q", s, *p, s)
	}

	empty := ""
	ep := strPtr(empty)
	if *ep != empty {
		t.Errorf("strPtr(empty) = %q, want %q", *ep, empty)
	}
}