package dispatcher

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

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/manager"
)

// TestMachineCatalogReconnectConvergesAndReplacesDesktopProcess exercises the
// production Desktop transport chain across an offline interval: initial
// machine Agent Catalog WebSocket payload → verified HTTP Catalog fetch →
// disconnect while the server revokes → reconnect authoritative payload →
// stale runtime lease revocation and real Process retirement.
func TestMachineCatalogReconnectConvergesAndReplacesDesktopProcess(t *testing.T) {
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	var catalogVersion atomic.Int64
	var connectionCount atomic.Int32
	catalogVersion.Store(7)
	disconnectForRevoke := make(chan struct{})
	websocketConnected := make(chan int32, 4)
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}

	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		switch {
		case strings.HasSuffix(request.URL.Path, "/daemon/connect"):
			connection, err := upgrader.Upgrade(writer, request, nil)
			require.NoError(t, err)
			generation := connectionCount.Add(1)
			websocketConnected <- generation
			defer connection.Close()
			if generation == 1 {
				go func() {
					select {
					case <-disconnectForRevoke:
						catalogVersion.Store(8)
						_ = connection.Close()
					case <-request.Context().Done():
					}
				}()
			}
			for {
				_, raw, err := connection.ReadMessage()
				if err != nil {
					return
				}
				var envelope struct {
					Type string `json:"type"`
				}
				if json.Unmarshal(raw, &envelope) != nil || envelope.Type != "machine:agent:list" {
					continue
				}
				if catalogVersion.Load() == 7 {
					require.NoError(t, connection.WriteJSON(machineCatalogResponse(7)))
				} else {
					require.NoError(t, connection.WriteJSON(machineCatalogResponse(8)))
				}
			}
		case strings.HasSuffix(request.URL.Path, "/runtime/ext"):
			writer.Header().Set("Content-Type", "application/json")
			if catalogVersion.Load() == 7 {
				writer.Header().Set("ETag", `"ext-v1-granted"`)
				if request.Header.Get("If-None-Match") == `"ext-v1-granted"` {
					writer.WriteHeader(http.StatusNotModified)
					return
				}
				fmt.Fprint(
					writer,
					`{"extensions":[{"extensionName":"github","displayName":"GitHub","actions":[{"name":"list","summary":"List","safety":{"class":"read"},"inputSchema":{"$schema":"https://json-schema.org/draft/2020-12/schema","type":"object"},"outputSchema":{"type":"object"}}]}]}`,
				)
				return
			}
			writer.Header().Set("ETag", `"ext-v1-empty"`)
			if request.Header.Get("If-None-Match") == `"ext-v1-empty"` {
				writer.WriteHeader(http.StatusNotModified)
				return
			}
			fmt.Fprint(writer, `{"extensions":[]}`)
		default:
			http.NotFound(writer, request)
		}
	}))
	defer server.Close()

	manager := manager.New(nil, nil)
	dispatcher, err := NewWithConnectionManager(server.URL, manager, Dependencies{})
	require.NoError(t, err)
	t.Cleanup(dispatcher.Close)
	connection, err := wire.New(profile.AgentTeamServer{
		ID: "agentteam-main", Enabled: true, ServerURL: server.URL,
		APIKey: "machine-key", WorkspaceID: "ws-1",
	}, dispatcher)
	require.NoError(t, err)
	manager.Register(connection)
	runCtx, cancelRun := context.WithCancel(context.Background())
	t.Cleanup(cancelRun)
	require.NoError(t, connection.Start(runCtx))
	select {
	case generation := <-websocketConnected:
		require.Equal(t, int32(1), generation)
	case <-time.After(time.Second):
		t.Fatal("Desktop Protocol did not connect to machine Catalog WebSocket")
	}

	key := extensionCatalogKey{serverID: "agentteam-main", agentID: "agent-1"}
	require.Eventually(t, func() bool {
		hint, snapshot, publishedRevision, published, ok := dispatcher.extensionCatalogState(key)
		return ok && hint.Version == 7 && published &&
			publishedRevision == 7 && snapshot.ETag() == `"ext-v1-granted"` &&
			!snapshot.Empty()
	}, 3*time.Second, 10*time.Millisecond)
	_, snapshot, _, _, _ := dispatcher.extensionCatalogState(key)
	api := connection.(*wire.Protocol).APIClient("agent-1")
	plc := execTestLaunchContext("agentteam-main")
	plc.startCfg.LaunchID = "process-generation-1"
	execConfig, err := dispatcher.activateExec(
		context.Background(),
		plc,
		"session-1",
		extensionExecResolved(t, api, snapshot),
	)
	require.NoError(t, err)
	tokenValue, err := os.ReadFile(execConfig.TokenFile)
	require.NoError(t, err)

	driver := newBlockingDispatcherManagedDriver()
	process, err := agent.NewProcess(
		context.Background(),
		driver,
		drivercontract.LaunchConfig{
			AgentID: "agent-1", WorkDir: t.TempDir(), Exec: execConfig,
		},
		agent.ProcessOptions{},
	)
	require.NoError(t, err)
	require.True(t, driver.waitStartedWithin(time.Second))
	t.Cleanup(func() {
		driver.allowExit()
		process.Stop()
		process.WaitForExit(time.Second)
	})
	session := &interactiveSession{
		sessionKey: "agent-1", agentID: "agent-1", process: process,
		binding: conversationBinding{channelID: "agentteam-main", serverAgentID: "agent-1"},
	}
	dispatcher.sessionManager.Set("agent-1", session)

	close(disconnectForRevoke)
	select {
	case generation := <-websocketConnected:
		require.GreaterOrEqual(t, generation, int32(2))
	case <-time.After(5 * time.Second):
		t.Fatal("Desktop Protocol did not reconnect after the offline revoke interval")
	}
	require.Eventually(t, func() bool {
		session.mu.Lock()
		defer session.mu.Unlock()
		return session.process == nil
	}, 3*time.Second, 10*time.Millisecond)
	_, err = token.Default().ValidateExecToken(
		"agent-1",
		"agentteam-main",
		strings.TrimSpace(string(tokenValue)),
	)
	require.Error(t, err)
	require.Eventually(t, func() bool {
		hint, current, publishedRevision, published, ok := dispatcher.extensionCatalogState(key)
		return ok && hint.Version == 8 && published &&
			publishedRevision == 8 && current.ETag() == `"ext-v1-empty"` &&
			current.Empty()
	}, 3*time.Second, 10*time.Millisecond)
}

func machineCatalogResponse(version int64) map[string]any {
	return map[string]any{
		"type": "machine:agent:list:response",
		"agents": []map[string]any{{
			"id": "agent-1", "name": "extension-agent", "displayName": "Extension Agent",
			"runtime": "codex", "status": "idle", "activity": "",
			"extensionCatalog": types.ExtensionCatalogHint{Version: version},
		}},
	}
}