package daemon

import (
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/store/agent"
)

func TestServerCapabilitiesSentOnSessionRegistration(t *testing.T) {
	t.Parallel()

	serverConn, desktopConn := newServerCapabilitiesWebSocketPair(t)
	hub := New(
		slog.New(slog.NewTextHandler(io.Discard, nil)),
		&recordingStartStore{},
		"",
		nil,
	)
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	serveDone := make(chan error, 1)
	go func() {
		serveDone <- hub.ServeConnection(ctx, agent.MachineIdentity{
			MachineID:       "mac_capabilities",
			WorkspaceID:     "wsp_capabilities",
			CreatedByUserID: "usr_capabilities",
		}, serverConn)
	}()

	require.NoError(t, desktopConn.SetReadDeadline(time.Now().Add(time.Second)))
	var frame struct {
		Type         string   `json:"type"`
		Capabilities []string `json:"capabilities"`
	}
	require.NoError(t, desktopConn.ReadJSON(&frame))
	require.Equal(t, "server:capabilities", frame.Type)
	require.Contains(t, frame.Capabilities, "transient_legacy_wake_status")

	cancel()
	require.NoError(t, desktopConn.Close())
	select {
	case <-serveDone:
	case <-time.After(time.Second):
		t.Fatal("ServeConnection did not stop after cancellation")
	}
}

func newServerCapabilitiesWebSocketPair(t *testing.T) (*websocket.Conn, *websocket.Conn) {
	t.Helper()

	serverConnCh := make(chan *websocket.Conn, 1)
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		require.NoError(t, err)
		serverConnCh <- conn
	}))
	t.Cleanup(server.Close)

	desktopConn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
	require.NoError(t, err)
	t.Cleanup(func() { _ = desktopConn.Close() })

	select {
	case serverConn := <-serverConnCh:
		t.Cleanup(func() { _ = serverConn.Close() })
		return serverConn, desktopConn
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for server websocket")
		return nil, nil
	}
}