package wire

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

func TestServerCapabilitiesDoNotMakeLegacyWakeStatusTransientBeforeStream(t *testing.T) {
	protocol := &Protocol{
		id:     "agentteam-test",
		sendCh: make(chan []byte, 1),
	}

	protocol.dispatch(context.Background(), []byte(`{
		"type":"server:capabilities",
		"capabilities":[" transient_legacy_wake_status ","","transient_legacy_wake_status"]
	}`))
	require.True(t, protocol.hasServerCapability(transientLegacyWakeStatusServerCapability))
	protocol.Emit(types.WakeStatusEvent{
		AgentID:   "agt_1",
		MessageID: "msg_1",
		WakeToken: "wake_1",
		Status:    consts.WakeStatusReplySent,
	})

	var frame map[string]json.RawMessage
	require.NoError(t, json.Unmarshal(<-protocol.sendCh, &frame))
	require.NotContains(t, frame, "transient")
}

func TestReadyAdvertisesTransientWakeStatusCapability(t *testing.T) {
	t.Setenv("PATH", t.TempDir())
	protocol := &Protocol{
		agents:    make(map[string]*agentState),
		dedupRing: newDedupRing(10),
	}

	message, _ := protocol.readyMessage()
	capabilities, ok := message["capabilities"].([]string)
	require.True(t, ok)
	require.Contains(t, capabilities, "wake:status-transient")
}

func TestConnectClearsServerCapabilitiesBeforePumpsStart(t *testing.T) {
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	release := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		require.NoError(t, err)
		defer conn.Close()
		<-release
	}))
	t.Cleanup(func() {
		close(release)
		server.Close()
	})

	protocol := &Protocol{
		id:          "agentteam-test",
		serverURL:   server.URL,
		apiKey:      "machine-key",
		workspaceID: "wsp_1",
	}
	protocol.replaceServerCapabilities([]string{"stale_connection_capability"})

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	require.NoError(t, protocol.connect(ctx))
	require.False(t, protocol.hasServerCapability("stale_connection_capability"))

	protocol.wsMu.Lock()
	conn := protocol.conn
	protocol.wsMu.Unlock()
	require.NotNil(t, conn)
	require.NoError(t, conn.Close())
}