package wire

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

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

type connectionRunRoundTripper func(*http.Request) (*http.Response, error)

func (fn connectionRunRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	return fn(req)
}

func TestConnectionRunStrongJoinDrainsOldSubordinatesBeforeReplacement(t *testing.T) {
	channel := &Protocol{
		id:     "connection-run-test",
		sendCh: make(chan []byte, 16),
	}
	oldRun := newConnectionRun(context.Background())
	releaseOld := make(chan struct{})
	cancelled := make(chan string, 3)
	oldSendErrors := make(chan error, 3)

	// These model the three previously detached classes. They deliberately wait
	// past cancellation and enqueue a final frame while unwinding; the reconnect
	// boundary must join them and drain those frames before a replacement starts.
	for _, class := range []string{"handler", "reporter", "ticker"} {
		class := class
		require.True(t, oldRun.start(oldRun.context(), func(ctx context.Context) {
			<-ctx.Done()
			cancelled <- class
			<-releaseOld
			oldSendErrors <- channel.sendJSON(map[string]string{
				"type":       "old:" + class,
				"generation": "old",
			})
		}))
	}

	finished := make(chan struct{})
	go func() {
		channel.finishConnectionRun(oldRun, nil, "test_replacement")
		close(finished)
	}()
	for range 3 {
		select {
		case <-cancelled:
		case <-time.After(time.Second):
			t.Fatal("old subordinate did not observe run cancellation")
		}
	}

	// closeAdmission and start share one mutex. Once cancellation is observable,
	// the owner snapshot is stable and a late nested task cannot escape behind it.
	require.False(t, oldRun.start(oldRun.context(), func(context.Context) {
		t.Error("task admitted after terminal run snapshot")
	}))
	select {
	case <-finished:
		t.Fatal("strong run boundary returned before admitted owners unwound")
	case <-time.After(30 * time.Millisecond):
	}

	close(releaseOld)
	select {
	case <-finished:
	case <-time.After(time.Second):
		t.Fatal("strong run boundary did not join admitted owners")
	}
	for range 3 {
		require.NoError(t, <-oldSendErrors)
	}
	require.Empty(t, channel.sendCh, "old-generation frames survived the post-owner drain")

	newRun := newConnectionRun(context.Background())
	newSent := make(chan error, 1)
	require.True(t, newRun.start(newRun.context(), func(context.Context) {
		newSent <- channel.sendJSON(map[string]string{
			"type":       "new:ready",
			"generation": "new",
		})
	}))
	select {
	case err := <-newSent:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("replacement run did not send")
	}

	var frame map[string]string
	require.NoError(t, json.Unmarshal(<-channel.sendCh, &frame))
	require.Equal(t, "new", frame["generation"])
	owners := newRun.closeAdmission()
	require.NoError(t, waitForConnectionRunOwners(context.Background(), owners))
}

func TestConnectionRunBoundedWaitCanBeUpgradedToExactStrongJoin(t *testing.T) {
	run := newConnectionRun(context.Background())
	entered := make(chan struct{})
	release := make(chan struct{})
	require.True(t, run.start(run.context(), func(context.Context) {
		close(entered)
		<-release // Deliberately model a subordinate slow to honor cancellation.
	}))
	<-entered

	owners := run.closeAdmission()
	boundedCtx, cancelBounded := context.WithTimeout(context.Background(), 25*time.Millisecond)
	defer cancelBounded()
	require.ErrorIs(t, waitForConnectionRunOwners(boundedCtx, owners), context.DeadlineExceeded)

	strongDone := make(chan error, 1)
	go func() { strongDone <- waitForConnectionRunOwners(context.Background(), owners) }()
	select {
	case err := <-strongDone:
		t.Fatalf("strong waiter crossed the still-live owner boundary: %v", err)
	case <-time.After(25 * time.Millisecond):
	}
	close(release)
	select {
	case err := <-strongDone:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("strong waiter did not join the original owner snapshot")
	}
}

func TestProtocolBoundedStopThenStrongStopJoinsConnectionTokenOwner(t *testing.T) {
	t.Setenv("PATH", t.TempDir())
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	connected := make(chan struct{})
	var connectedOnce sync.Once
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		conn, err := upgrader.Upgrade(w, req, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		connectedOnce.Do(func() { close(connected) })
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}))
	t.Cleanup(server.Close)

	tokenEntered := make(chan struct{})
	releaseToken := make(chan struct{})
	var tokenEnteredOnce sync.Once
	channel := &Protocol{
		id:        "strong-token-owner",
		enabled:   true,
		serverURL: server.URL,
		apiKey:    "machine-key",
		authToken: "machine-key",
		httpClient: &http.Client{Transport: connectionRunRoundTripper(func(req *http.Request) (*http.Response, error) {
			if !strings.HasSuffix(req.URL.Path, "/daemon/session-token") {
				return &http.Response{
					StatusCode: http.StatusNotFound,
					Body:       io.NopCloser(strings.NewReader(`{}`)),
					Header:     make(http.Header),
					Request:    req,
				}, nil
			}
			tokenEnteredOnce.Do(func() { close(tokenEntered) })
			<-releaseToken // Deliberately ignore request cancellation.
			return &http.Response{
				StatusCode: http.StatusOK,
				Body:       io.NopCloser(strings.NewReader(`{"token":"user-token"}`)),
				Header:     make(http.Header),
				Request:    req,
			}, nil
		})},
		sendCh:                 make(chan []byte, sendBufSize),
		agents:                 make(map[string]*agentState),
		agentTokenRefresh:      make(map[string]*agentTokenRefreshState),
		dedupRing:              newDedupRing(dedupSetCap),
		workspaceListRequests:  make(map[string]chan workspaceListResult),
		remoteEndpointRegistry: defaultDesktopRemoteEndpointRegistry(),
		backoffFactory:         defaultReconnectBackoff,
		disp:                   &agentteamTestDispatcher{},
	}

	runCtx, cancelRun := context.WithCancel(context.Background())
	t.Cleanup(cancelRun)
	require.NoError(t, channel.Start(runCtx))
	select {
	case <-connected:
	case <-time.After(2 * time.Second):
		t.Fatal("protocol did not establish websocket connection")
	}
	// onConnect intentionally delays this best-effort exchange for two seconds.
	select {
	case <-tokenEntered:
	case <-time.After(4 * time.Second):
		t.Fatal("token prefetch owner did not start")
	}

	boundedCtx, cancelBounded := context.WithTimeout(context.Background(), 40*time.Millisecond)
	require.ErrorIs(t, channel.Stop(boundedCtx), context.DeadlineExceeded)
	cancelBounded()

	stopDone := make(chan error, 1)
	go func() { stopDone <- channel.Stop(context.Background()) }()
	select {
	case err := <-stopDone:
		t.Fatalf("Protocol.Stop crossed the still-live token owner: %v", err)
	case <-time.After(50 * time.Millisecond):
	}
	close(releaseToken)
	select {
	case err := <-stopDone:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("Protocol.Stop did not join the token owner")
	}
}