package localruntime

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

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/api/auth"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/uisession"
)

type wsTestSubscriptionProvider struct {
	snapshot uisession.SubscriptionSnapshot
	changes  chan uisession.Change
	releases atomic.Int32
}

func (p *wsTestSubscriptionProvider) Subscribe(string, time.Time) (uisession.Subscription, error) {
	var once sync.Once
	return uisession.Subscription{
		Snapshot: p.snapshot,
		Changes:  p.changes,
		Release: func() {
			once.Do(func() {
				p.releases.Add(1)
			})
		},
	}, nil
}

type wsTestTimer struct {
	mu sync.Mutex
	cb func()
}

func (t *wsTestTimer) AfterFunc(_ time.Duration, cb func()) func() {
	t.mu.Lock()
	t.cb = cb
	t.mu.Unlock()
	return func() {}
}

func (t *wsTestTimer) Fire() {
	t.mu.Lock()
	cb := t.cb
	t.mu.Unlock()
	if cb != nil {
		cb()
	}
}

func newWSStreamServer(
	t *testing.T,
	provider SubscriptionProvider,
	registry ServerRegistry,
	timer LeaseTimerRegistry,
) (string, func()) {
	t.Helper()
	router := gin.New()
	router.GET("/api/internal/servers/active/stream", NewActiveSelectionStreamHandler(StreamHandlerDeps{
		Provider: provider,
		Servers:  registry,
		Timer:    timer,
	}))
	server := httptest.NewServer(router)
	return "ws" + strings.TrimPrefix(server.URL, "http") +
		"/api/internal/servers/active/stream", server.Close
}

func dialAuthenticatedWS(t *testing.T, wsURL string) *websocket.Conn {
	t.Helper()
	conn, _, err := websocket.DefaultDialer.DialContext(
		context.Background(),
		wsURL+"?accessToken="+token.CreateSession(),
		nil,
	)
	require.NoError(t, err)
	require.NoError(t, conn.SetReadDeadline(time.Now().Add(2*time.Second)))
	return conn
}

// TestActiveSelectionStream_WebSocket_HelloAndSnapshot covers the
// happy-path handshake: a client that opens the WS with a valid session
// token via the accessToken query parameter receives a hello frame
// followed by a snapshot frame reflecting the current uisession state.
func TestActiveSelectionStream_WebSocket_HelloAndSnapshot(t *testing.T) {
	gin.SetMode(gin.TestMode)
	store := uisession.NewStore(uisession.Options{})
	registry := newTestServerRegistry(profile.AgentTeamServer{
		ID:          "server-a",
		Enabled:     true,
		Name:        "Production",
		WorkspaceID: "wsp-a",
		MachineID:   "mac-a",
	})

	router := gin.New()
	internal := router.Group("/api/internal")
	internal.GET("/servers/active/stream", NewActiveSelectionStreamHandler(StreamHandlerDeps{
		Provider: store,
		Servers:  registry,
	}))

	server := httptest.NewServer(router)
	defer server.Close()
	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/internal/servers/active/stream"

	sessionToken := token.CreateSession()
	conn, _, err := websocket.DefaultDialer.DialContext(
		context.Background(),
		wsURL+"?accessToken="+sessionToken,
		nil,
	)
	require.NoError(t, err)
	defer conn.Close()
	require.NoError(t, conn.SetReadDeadline(time.Now().Add(2*time.Second)))

	_, hello, err := conn.ReadMessage()
	require.NoError(t, err)
	var helloFrame activeSelectionStreamFrame
	require.NoError(t, json.Unmarshal(hello, &helloFrame))
	require.Equal(t, "hello", helloFrame.Type)

	_, snapshot, err := conn.ReadMessage()
	require.NoError(t, err)
	var snapshotFrame activeSelectionStreamFrame
	require.NoError(t, json.Unmarshal(snapshot, &snapshotFrame))
	require.Equal(t, "snapshot", snapshotFrame.Type)
	require.Equal(t, "empty", snapshotFrame.State)
}

func TestActiveSelectionStream_WebSocket_QueryAuthThroughProductionRouteShape(t *testing.T) {
	gin.SetMode(gin.TestMode)
	store := uisession.NewStore(uisession.Options{})
	registry := newTestServerRegistry()
	router := gin.New()
	RegisterActiveSelectionStreamRoute(router, Dependencies{
		Sessions: store,
		Servers:  registry,
	})
	protected := router.Group("/api/internal", auth.Middleware())
	protected.GET("/protected", func(c *gin.Context) { c.Status(http.StatusNoContent) })

	server := httptest.NewServer(router)
	defer server.Close()
	sessionToken := token.CreateSession()
	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") +
		"/api/internal/servers/active/stream?accessToken=" + sessionToken
	conn, _, err := websocket.DefaultDialer.DialContext(context.Background(), wsURL, nil)
	require.NoError(t, err)
	defer conn.Close()
	_, hello, err := conn.ReadMessage()
	require.NoError(t, err)
	require.Contains(t, string(hello), `"type":"hello"`)

	response, err := http.Get(server.URL + "/api/internal/protected?accessToken=" + sessionToken)
	require.NoError(t, err)
	defer response.Body.Close()
	require.Equal(t, http.StatusUnauthorized, response.StatusCode,
		"query auth must remain scoped to the active-selection WebSocket")
}

func TestActiveSelectionStream_WebSocket_ProductionRouteWrongTokenCloses4401(t *testing.T) {
	router := gin.New()
	RegisterActiveSelectionStreamRoute(router, Dependencies{
		Sessions: uisession.NewStore(uisession.Options{}),
		Servers:  newTestServerRegistry(),
	})
	server := httptest.NewServer(router)
	defer server.Close()
	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") +
		"/api/internal/servers/active/stream?accessToken=wrong"
	conn, _, err := websocket.DefaultDialer.DialContext(context.Background(), wsURL, nil)
	require.NoError(t, err)
	defer conn.Close()
	_, _, err = conn.ReadMessage()
	var closeErr *websocket.CloseError
	require.ErrorAs(t, err, &closeErr)
	require.Equal(t, activeStreamAuthRequiredClose, closeErr.Code)
}

func TestActiveSelectionStream_WebSocket_CrossOriginRejectedBeforeSubscribe(t *testing.T) {
	provider := &wsTestSubscriptionProvider{
		snapshot: uisession.SubscriptionSnapshot{State: uisession.SubscriptionStateEmpty},
		changes:  make(chan uisession.Change, 1),
	}
	wsURL, closeServer := newWSStreamServer(t, provider, newTestServerRegistry(), nil)
	defer closeServer()
	sessionToken := token.CreateSession()
	_, response, err := websocket.DefaultDialer.DialContext(
		context.Background(),
		wsURL+"?accessToken="+sessionToken,
		http.Header{"Origin": []string{"https://evil.example"}},
	)
	require.Error(t, err)
	require.NotNil(t, response)
	defer response.Body.Close()
	require.Equal(t, http.StatusForbidden, response.StatusCode)
	require.Equal(t, int32(0), provider.releases.Load(),
		"cross-origin rejection must happen before Subscribe")
}

func TestActiveSelectionStream_WebSocket_LeaseExpiryClosesAndReleasesOnce(t *testing.T) {
	provider := &wsTestSubscriptionProvider{
		snapshot: uisession.SubscriptionSnapshot{State: uisession.SubscriptionStateEmpty},
		changes:  make(chan uisession.Change, 1),
	}
	timer := &wsTestTimer{}
	wsURL, closeServer := newWSStreamServer(t, provider, newTestServerRegistry(), timer)
	defer closeServer()
	conn := dialAuthenticatedWS(t, wsURL)
	defer conn.Close()
	_, _, err := conn.ReadMessage() // hello
	require.NoError(t, err)
	_, _, err = conn.ReadMessage() // snapshot
	require.NoError(t, err)

	timer.Fire()
	_, _, err = conn.ReadMessage()
	require.Error(t, err)
	require.Eventually(t, func() bool { return provider.releases.Load() == 1 },
		time.Second, 10*time.Millisecond)
}

func TestActiveSelectionStream_WebSocket_ClientCloseReleasesOnce(t *testing.T) {
	provider := &wsTestSubscriptionProvider{
		snapshot: uisession.SubscriptionSnapshot{State: uisession.SubscriptionStateEmpty},
		changes:  make(chan uisession.Change, 1),
	}
	wsURL, closeServer := newWSStreamServer(t, provider, newTestServerRegistry(), nil)
	defer closeServer()
	conn := dialAuthenticatedWS(t, wsURL)
	_, _, err := conn.ReadMessage()
	require.NoError(t, err)
	_, _, err = conn.ReadMessage()
	require.NoError(t, err)
	require.NoError(t, conn.Close())
	require.Eventually(t, func() bool { return provider.releases.Load() == 1 },
		time.Second, 10*time.Millisecond)
}

func TestActiveSelectionStream_WebSocket_SlowClientWriteDeadlineReleasesOnce(t *testing.T) {
	previousWriteWait := activeStreamDataWriteWait
	activeStreamDataWriteWait = 50 * time.Millisecond
	t.Cleanup(func() { activeStreamDataWriteWait = previousWriteWait })

	registry := newTestServerRegistry(profile.AgentTeamServer{
		ID: "server-a", Enabled: true, Name: strings.Repeat("x", 2<<20),
	})
	provider := &wsTestSubscriptionProvider{
		snapshot: uisession.SubscriptionSnapshot{State: uisession.SubscriptionStateEmpty},
		changes:  make(chan uisession.Change, 32),
	}
	wsURL, closeServer := newWSStreamServer(t, provider, registry, nil)
	defer closeServer()
	conn := dialAuthenticatedWS(t, wsURL)
	defer conn.Close()
	_, _, err := conn.ReadMessage()
	require.NoError(t, err)
	_, _, err = conn.ReadMessage()
	require.NoError(t, err)

	// The producer owns writes to changes; Release only records subscription
	// teardown and never closes a channel it does not own.
	for i := 0; i < cap(provider.changes); i++ {
		change := uisession.Change{
			Kind: uisession.ChangeKindSelected, State: uisession.SubscriptionStateSelected,
			ActiveServerID: "server-a", Revision: uint64(i + 1),
		}
		select {
		case provider.changes <- change:
		default:
			// The handler has already drained enough frames to reach its write
			// deadline. Stop producing rather than racing Release's close.
			goto filled
		}
	}
filled:
	require.Eventually(t, func() bool { return provider.releases.Load() == 1 },
		3*time.Second, 10*time.Millisecond,
		"blocked data writes must expire and release the subscription")
}

func TestActiveSelectionStream_WebSocket_ResolveMissFailsClosed(t *testing.T) {
	t.Run("initial snapshot", func(t *testing.T) {
		provider := &wsTestSubscriptionProvider{
			snapshot: uisession.SubscriptionSnapshot{
				State:          uisession.SubscriptionStateSelected,
				ActiveServerID: "missing",
			},
			changes: make(chan uisession.Change, 1),
		}
		wsURL, closeServer := newWSStreamServer(t, provider, newTestServerRegistry(), nil)
		defer closeServer()
		conn := dialAuthenticatedWS(t, wsURL)
		defer conn.Close()
		_, hello, err := conn.ReadMessage()
		require.NoError(t, err)
		require.Contains(t, string(hello), `"type":"hello"`)
		_, _, err = conn.ReadMessage()
		require.Error(t, err, "missing server must close before a snapshot frame")
		require.Equal(t, int32(1), provider.releases.Load())
	})

	t.Run("post snapshot change", func(t *testing.T) {
		registry := newTestServerRegistry(profile.AgentTeamServer{
			ID: "server-a", Enabled: true, Name: "Server A",
		})
		provider := &wsTestSubscriptionProvider{
			snapshot: uisession.SubscriptionSnapshot{
				State: uisession.SubscriptionStateEmpty,
			},
			changes: make(chan uisession.Change, 1),
		}
		wsURL, closeServer := newWSStreamServer(t, provider, registry, nil)
		defer closeServer()
		conn := dialAuthenticatedWS(t, wsURL)
		defer conn.Close()
		_, _, err := conn.ReadMessage()
		require.NoError(t, err)
		_, snapshot, err := conn.ReadMessage()
		require.NoError(t, err)
		require.Contains(t, string(snapshot), `"type":"snapshot"`)
		provider.changes <- uisession.Change{
			Kind: uisession.ChangeKindSelected, State: uisession.SubscriptionStateSelected,
			ActiveServerID: "missing", Revision: 2,
		}
		_, _, err = conn.ReadMessage()
		require.Error(t, err, "missing changed server must close without a partial frame")
		require.Equal(t, int32(1), provider.releases.Load())
	})
}

// TestActiveSelectionStream_WebSocket_AuthFailureClosesWith4401 covers
// the auth-failure close path: a client that opens the WS without a
// valid session token receives a 4401 close frame so the reconnect
// agentteam can map it to onTerminalAuth.
func TestActiveSelectionStream_WebSocket_AuthFailureClosesWith4401(t *testing.T) {
	gin.SetMode(gin.TestMode)
	store := uisession.NewStore(uisession.Options{})
	registry := newTestServerRegistry()

	router := gin.New()
	internal := router.Group("/api/internal")
	internal.GET("/servers/active/stream", NewActiveSelectionStreamHandler(StreamHandlerDeps{
		Provider: store,
		Servers:  registry,
	}))

	server := httptest.NewServer(router)
	defer server.Close()
	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/internal/servers/active/stream"

	conn, _, err := websocket.DefaultDialer.DialContext(context.Background(), wsURL, nil)
	require.NoError(t, err)
	defer conn.Close()
	require.NoError(t, conn.SetReadDeadline(time.Now().Add(2*time.Second)))

	_, _, err = conn.ReadMessage()
	var closeErr *websocket.CloseError
	require.ErrorAs(t, err, &closeErr)
	require.Equal(t, activeStreamAuthRequiredClose, closeErr.Code)
}

// TestActiveSelectionStream_WebSocket_SessionIdentityFromHeader covers
// non-browser clients that authenticate via the Authorization Bearer
// header rather than the accessToken query parameter.
func TestActiveSelectionStream_WebSocket_SessionIdentityFromHeader(t *testing.T) {
	gin.SetMode(gin.TestMode)
	store := uisession.NewStore(uisession.Options{})
	registry := newTestServerRegistry()

	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	internal.GET("/servers/active/stream", NewActiveSelectionStreamHandler(StreamHandlerDeps{
		Provider: store,
		Servers:  registry,
	}))

	server := httptest.NewServer(router)
	defer server.Close()
	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/internal/servers/active/stream"

	sessionToken := token.CreateSession()
	conn, _, err := websocket.DefaultDialer.DialContext(
		context.Background(),
		wsURL,
		http.Header{"Authorization": []string{"Bearer " + sessionToken}},
	)
	require.NoError(t, err)
	defer conn.Close()
	require.NoError(t, conn.SetReadDeadline(time.Now().Add(2*time.Second)))

	_, hello, err := conn.ReadMessage()
	require.NoError(t, err)
	var helloFrame activeSelectionStreamFrame
	require.NoError(t, json.Unmarshal(hello, &helloFrame))
	require.Equal(t, "hello", helloFrame.Type)
}