package dispatcher

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"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/api/exec"
	"github.com/tidbcloud/AgentTeam/desktop/internal/promptcompiler"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

func TestResolveInjectsAuthorizedExtensionsOnlyThroughExec(t *testing.T) {
	t.Setenv("AGENTTEAM_CONFIG_DIR", t.TempDir())
	_ = initTokenForTest(t)
	snapshot, api, closeServer := extensionTestCatalog(t, 7)
	defer closeServer()

	d := mustNewDispatcher("http://127.0.0.1:3456")
	t.Cleanup(d.Close)
	d.serverConnections.Register(extensionTestServerConnection{
		mcpTestServerConnection: mcpTestServerConnection{id: "agentteam-main"},
		api:                     api,
	})
	hint := types.ExtensionCatalogHint{Version: 7}
	d.ObserveAgentExtensionCatalogHints("agentteam-main", []types.MachineAgentInfo{{
		ID: "agent-1", ExtensionCatalog: &hint,
	}})

	plc := execTestLaunchContext("agentteam-main")
	plc.mode = promptModeSession
	require.NoError(t, d.resolve(plc))
	t.Cleanup(func() { d.abortSessionLaunch(plc.sess, plc.startCfg) })
	require.NotNil(t, plc.startCfg.Exec)
	require.Equal(t, plc.startCfg.Exec.BaseURL, plc.startCfg.Env[exec.EnvBaseURL])
	require.Equal(t, plc.startCfg.Exec.TokenFile, plc.startCfg.Env[exec.EnvTokenFile])
	require.NotContains(t, plc.startCfg.Env, "AGENTTEAM_EXTENSION_BASE_URL")
	require.NotContains(t, plc.startCfg.Env, "AGENTTEAM_EXTENSION_TOKEN_FILE")
	require.NotContains(t, plc.startCfg.Env, "AGENTTEAM_EXTENSION_TOKEN")
	require.FileExists(t, plc.startCfg.Exec.TokenFile)
	require.True(t, plc.instructions.extension.Enabled)
	require.Equal(t, promptcompiler.ExtensionCatalogGuidanceTemplateData{Actions: []promptcompiler.ExtensionCatalogGuidanceAction{{ExtensionName: "github", ActionName: "list"}}}, plc.instructions.extension.Guidance)
	require.NotContains(t, testLaunchSurface(plc.startCfg, "baseInstructions"), "Catalog version")
	require.NotContains(t, testLaunchSurface(plc.startCfg, "baseInstructions"), snapshot.ETag())
	require.Contains(t, testLaunchSurface(plc.startCfg, "baseInstructions"), "agentteam exec <extension-name> <action> --help")
}

func TestResolveOmitsExtensionEnvWithoutAuthorizedCatalog(t *testing.T) {
	t.Setenv("AGENTTEAM_CONFIG_DIR", t.TempDir())
	_ = initTokenForTest(t)
	d := mustNewDispatcher("http://127.0.0.1:3456")
	t.Cleanup(d.Close)
	d.serverConnections.Register(execTestServerConnection{
		mcpTestServerConnection: mcpTestServerConnection{id: "agentteam-main"},
	})

	plc := execTestLaunchContext("agentteam-main")
	plc.mode = promptModeSession
	require.NoError(t, d.resolve(plc))
	t.Cleanup(func() { d.abortSessionLaunch(plc.sess, plc.startCfg) })
	require.NotContains(t, plc.startCfg.Env, "AGENTTEAM_EXTENSION_BASE_URL")
	require.NotContains(t, plc.startCfg.Env, "AGENTTEAM_EXTENSION_TOKEN_FILE")
	require.NotContains(t, plc.startCfg.Env, "AGENTTEAM_EXTENSION_TOKEN")
	require.False(t, plc.instructions.extension.Enabled)
	require.NotContains(t, testLaunchSurface(plc.startCfg, "baseInstructions"), "platform.agentteam.extension")
}

func TestCatalogHintChangeRevokesOldRuntimeLeaseAfterChangedContentFetch(t *testing.T) {
	manager, err := token.NewKeyManager(t.TempDir() + "/config.yaml")
	require.NoError(t, err)
	snapshot, api, closeServer := extensionTestCatalog(t, 7)
	defer closeServer()
	d := &Dispatcher{
		serverURL: "http://127.0.0.1:3456", execAuthorizations: exec.NewRuntimeAuthorizationStore(manager),
		sessionManager: NewSessionManager(),
	}
	d.initServices()
	hint := types.ExtensionCatalogHint{Version: 7}
	d.ObserveAgentExtensionCatalogHints("agentteam-main", []types.MachineAgentInfo{{
		ID: "agent-1", ExtensionCatalog: &hint,
	}})
	require.NoError(t, d.RefreshAgentExtensionCatalog(context.Background(), "agentteam-main", "agent-1", api))

	plc := execTestLaunchContext("agentteam-main")
	plc.startCfg.LaunchID = "process-generation-1"
	cfg, err := d.activateExec(
		context.Background(),
		plc,
		"session-1",
		extensionExecResolved(t, api, snapshot),
	)
	require.NoError(t, err)
	value, err := os.ReadFile(cfg.TokenFile)
	require.NoError(t, err)
	driver := newBlockingDispatcherManagedDriver()
	process, err := agent.NewProcess(
		context.Background(),
		driver,
		drivercontract.LaunchConfig{
			AgentID: "agent-1", WorkDir: t.TempDir(), Exec: cfg,
		},
		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,
		ephemeral: false, durableDelivery: false,
	}
	d.sessionManager.Set("agent-1", session)

	revokedHint := types.ExtensionCatalogHint{Version: 8}
	d.ObserveAgentExtensionCatalogHints("agentteam-main", []types.MachineAgentInfo{{
		ID: "agent-1", ExtensionCatalog: &revokedHint,
	}})
	_, err = manager.ValidateExecToken("agent-1", "agentteam-main", strings.TrimSpace(string(value)))
	require.NoError(t, err, "revision alone is only a notification and must not replace the process")
	session.mu.Lock()
	require.NotNil(t, session.process)
	session.mu.Unlock()

	revokedServer := httptest.NewServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		require.Equal(t, snapshot.ETag(), request.Header.Get("If-None-Match"))
		writer.Header().Set("Content-Type", "application/json")
		writer.Header().Set("ETag", `"ext-v1-v8-empty"`)
		fmt.Fprint(writer, `{"extensions":[]}`)
	}))
	defer revokedServer.Close()
	api.ServerURL = revokedServer.URL
	api.HTTPClient = revokedServer.Client()
	require.NoError(t, d.RefreshAgentExtensionCatalog(
		context.Background(),
		"agentteam-main",
		"agent-1",
		api,
	))

	_, cached, publishedRevision, published, ok := d.extensionCatalogState(
		extensionCatalogKey{serverID: "agentteam-main", agentID: "agent-1"},
	)
	require.True(t, ok)
	require.True(t, published)
	require.EqualValues(t, 8, publishedRevision)
	require.True(t, cached.Empty())
	require.Equal(t, `"ext-v1-v8-empty"`, cached.ETag())
	_, err = manager.ValidateExecToken("agent-1", "agentteam-main", strings.TrimSpace(string(value)))
	require.Error(t, err, "changed fetched content must revoke the old process capability")
	require.Eventually(t, func() bool {
		_, err := os.Stat(cfg.TokenFile)
		return os.IsNotExist(err)
	}, time.Second, time.Millisecond)
	session.mu.Lock()
	require.Nil(t, session.process, "changed fetched content must retire the old process generation")
	session.mu.Unlock()
}

func TestCatalogHintDecreaseAndResetStillFetchAuthoritativeContent(t *testing.T) {
	for _, nextVersion := range []int64{7, 0} {
		t.Run(fmt.Sprintf("revision-%d", nextVersion), func(t *testing.T) {
			snapshot, api, closeServer := extensionTestCatalog(t, 8)
			defer closeServer()
			d := &Dispatcher{serverURL: "http://127.0.0.1:3456"}
			d.initServices()
			initial := types.ExtensionCatalogHint{Version: 8}
			d.ObserveAgentExtensionCatalogHints(
				"agentteam-main",
				[]types.MachineAgentInfo{{
					ID:               "agent-1",
					ExtensionCatalog: &initial,
				}},
			)
			require.NoError(t, d.RefreshAgentExtensionCatalog(
				context.Background(),
				"agentteam-main",
				"agent-1",
				api,
			))

			var requests int
			authoritative := httptest.NewServer(http.HandlerFunc(func(
				writer http.ResponseWriter,
				request *http.Request,
			) {
				requests++
				require.Equal(
					t,
					snapshot.ETag(),
					request.Header.Get("If-None-Match"),
				)
				writer.Header().Set("Content-Type", "application/json")
				writer.Header().Set("ETag", `"ext-v1-empty"`)
				fmt.Fprint(writer, `{"extensions":[]}`)
			}))
			defer authoritative.Close()
			api.ServerURL = authoritative.URL
			api.HTTPClient = authoritative.Client()

			next := types.ExtensionCatalogHint{Version: nextVersion}
			d.ObserveAgentExtensionCatalogHints(
				"agentteam-main",
				[]types.MachineAgentInfo{{
					ID:               "agent-1",
					ExtensionCatalog: &next,
				}},
			)
			require.NoError(t, d.RefreshAgentExtensionCatalog(
				context.Background(),
				"agentteam-main",
				"agent-1",
				api,
			))

			require.Equal(t, 1, requests)
			hint, cached, publishedRevision, published, ok :=
				d.extensionCatalogState(extensionCatalogKey{
					serverID: "agentteam-main",
					agentID:  "agent-1",
				})
			require.True(t, ok)
			require.Equal(t, nextVersion, hint.Version)
			require.True(t, published)
			require.Equal(t, nextVersion, publishedRevision)
			require.True(t, cached.Empty())
			require.Equal(t, `"ext-v1-empty"`, cached.ETag())
		})
	}
}

func TestExecActivationRejectsCatalogChangedBeforeAuthorizationPublication(t *testing.T) {
	manager, err := token.NewKeyManager(t.TempDir() + "/config.yaml")
	require.NoError(t, err)
	snapshot, api, closeServer := extensionTestCatalog(t, 7)
	defer closeServer()
	d := &Dispatcher{
		serverURL:          "http://127.0.0.1:3456",
		execAuthorizations: exec.NewRuntimeAuthorizationStore(manager),
	}
	d.initServices()
	key := extensionCatalogKey{serverID: "agentteam-main", agentID: "agent-1"}
	d.desired[key] = types.ExtensionCatalogHint{Version: 8}

	plc := execTestLaunchContext("agentteam-main")
	plc.startCfg.LaunchID = "process-generation-1"
	config, err := d.activateExec(
		context.Background(),
		plc,
		"session-1",
		extensionExecResolved(t, api, snapshot),
	)

	require.Nil(t, config)
	require.ErrorContains(t, err, "Catalog changed during activation")
	d.mu.Lock()
	require.Empty(t, d.authorizations)
	require.Zero(t, d.pendingActivations)
	d.mu.Unlock()
}

// An Extension failure must never abort agent startup: the launch proceeds
// with built-in Exec actions only and the failure stays operator-visible.
func TestResolveSurvivesExtensionCatalogFetchFailure(t *testing.T) {
	t.Setenv("AGENTTEAM_CONFIG_DIR", t.TempDir())
	_ = initTokenForTest(t)
	server := httptest.NewServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		_ *http.Request,
	) {
		http.Error(writer, "catalog unavailable", http.StatusInternalServerError)
	}))
	defer server.Close()
	api := &apiclient.Client{
		ServerURL: server.URL, AgentID: "agent-1", AuthToken: "agent-token",
		RouteWorkspaceID: "ws-1", HTTPClient: server.Client(),
	}

	d := mustNewDispatcher("http://127.0.0.1:3456")
	t.Cleanup(d.Close)
	d.serverConnections.Register(extensionTestServerConnection{
		mcpTestServerConnection: mcpTestServerConnection{id: "agentteam-main"},
		api:                     api,
	})
	hint := types.ExtensionCatalogHint{Version: 7}
	d.ObserveAgentExtensionCatalogHints("agentteam-main", []types.MachineAgentInfo{{
		ID: "agent-1", ExtensionCatalog: &hint,
	}})

	plc := execTestLaunchContext("agentteam-main")
	plc.mode = promptModeSession
	require.NoError(t, d.resolve(plc), "agent startup must survive an Extension Catalog fetch failure")
	t.Cleanup(func() { d.abortSessionLaunch(plc.sess, plc.startCfg) })
	require.False(t, plc.instructions.extension.Enabled, "no Extension guidance without a Catalog")
	require.Contains(t, testLaunchSurface(plc.startCfg, "baseInstructions"), "agentteam exec", "built-in Exec stays compiled")
}

// A degraded activation (desired Catalog enabled, authorization carries zero
// Extension actions) must activate instead of aborting: an empty snapshot has
// nothing stale to fail-close on.
func TestExecActivationAcceptsDegradedEmptyExtensionSnapshot(t *testing.T) {
	manager, err := token.NewKeyManager(t.TempDir() + "/config.yaml")
	require.NoError(t, err)
	_, api, closeServer := extensionTestCatalog(t, 7)
	defer closeServer()
	d := &Dispatcher{
		serverURL: "http://127.0.0.1:3456", execAuthorizations: exec.NewRuntimeAuthorizationStore(manager),
	}
	d.initServices()
	key := extensionCatalogKey{serverID: "agentteam-main", agentID: "agent-1"}
	d.desired[key] = types.ExtensionCatalogHint{Version: 7}

	plc := execTestLaunchContext("agentteam-main")
	plc.startCfg.LaunchID = "process-generation-1"
	config, err := d.activateExec(
		context.Background(),
		plc,
		"session-1",
		extensionExecResolved(t, api, exec.ExtensionSnapshot{}),
	)

	require.NoError(t, err, "a degraded (no-Extension) activation must not abort")
	require.NotNil(t, config)
	require.NotEmpty(t, config.AuthorizationID)
	t.Cleanup(func() {
		drained := d.revokeExecRuntimeAuthorization(config.AuthorizationID)
		<-drained
	})
}

func extensionTestCatalog(
	t *testing.T,
	version int64,
) (exec.ExtensionSnapshot, *apiclient.Client, func()) {
	t.Helper()
	etag := fmt.Sprintf(`"ext-v1-v%d"`, version)
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		if !strings.HasSuffix(request.URL.EscapedPath(), "/runtime/ext") {
			http.NotFound(writer, request)
			return
		}
		writer.Header().Set("ETag", etag)
		if request.Header.Get("If-None-Match") == etag {
			writer.WriteHeader(http.StatusNotModified)
			return
		}
		fmt.Fprint(
			writer,
			`{"extensions":[{"extensionName":"github","displayName":"GitHub","actions":[{"name":"list","summary":"List","description":"List issues","safety":{"class":"read"},"inputSchema":{"$schema":"https://json-schema.org/draft/2020-12/schema","type":"object"},"outputSchema":{"type":"object"},"example":{"state":"open"}}]}]}`,
		)
	}))
	api := &apiclient.Client{
		ServerURL: server.URL, AgentID: "agent-1", AuthToken: "agent-token",
		RouteWorkspaceID: "ws-1", HTTPClient: server.Client(),
	}
	snapshot, diagnostics, modified, err := exec.FetchExtensionSnapshot(
		context.Background(),
		api,
		exec.ExtensionSnapshot{},
	)
	require.NoError(t, err)
	require.Empty(t, diagnostics)
	require.True(t, modified)
	require.Equal(t, etag, snapshot.ETag())
	return snapshot, api, server.Close
}

func extensionExecResolved(
	t *testing.T,
	api *apiclient.Client,
	snapshot exec.ExtensionSnapshot,
) *execResolved {
	t.Helper()
	builtin, err := exec.DefaultCatalog()
	require.NoError(t, err)
	runtimeCatalog, err := exec.NewRuntimeCatalog(builtin, snapshot)
	require.NoError(t, err)
	return &execResolved{api: api, runtimeCatalog: runtimeCatalog}
}

type extensionTestServerConnection struct {
	mcpTestServerConnection
	api *apiclient.Client
}

func (connection extensionTestServerConnection) APIClient(agentID ...string) *apiclient.Client {
	copy := *connection.api
	if len(agentID) > 0 {
		copy.AgentID = strings.TrimSpace(agentID[0])
	}
	return &copy
}