package dispatcher

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeevents"
)

func TestDispatcherSessionStateSubscriptionUnsubscribes(t *testing.T) {
	d := mustNewDispatcher("http://127.0.0.1:0")
	defer d.Close()

	sub := d.SubscribeSessionStateChanges("test-session-state")
	require.NotNil(t, sub)

	d.publishSessionStateChanged("agent-a")
	select {
	case got := <-sub.Events():
		require.Equal(t, runtimeevents.KindSessionStateChanged, got.Kind)
		require.Equal(t, "agent-a", got.Context.AgentID)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for session state event")
	}

	sub.Unsubscribe()
	d.publishSessionStateChanged("agent-a")
	select {
	case got := <-sub.Events():
		t.Fatalf("unexpected event after unsubscribe: %s", got.Kind)
	default:
	}
}

func TestDispatcherConsoleRefreshSubscriptionIncludesAgentCatalogChanges(t *testing.T) {
	d := mustNewDispatcher("http://127.0.0.1:0")
	defer d.Close()

	sub := d.SubscribeConsoleRefreshEvents("test-console-refresh")
	require.NotNil(t, sub)
	defer sub.Unsubscribe()

	d.PublishAgentCatalogChanged("agentteam-1")
	select {
	case got := <-sub.Events():
		require.Equal(t, runtimeevents.KindAgentCatalogChanged, got.Kind)
		require.Equal(t, "agentteam-1", got.Context.ServerID)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for agent catalog event")
	}

	d.publishSessionStateChanged("agent-a")
	select {
	case got := <-sub.Events():
		require.Equal(t, runtimeevents.KindSessionStateChanged, got.Kind)
		require.Equal(t, "agent-a", got.Context.AgentID)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for session state event")
	}
}

func TestDispatcherResetByAgentPublishesSessionStateEvent(t *testing.T) {
	d := mustNewDispatcher("http://127.0.0.1:0")
	defer d.Close()
	d.sessionManager.Set("agent-a", &interactiveSession{
		agentID: "agent-a",
		workDir: t.TempDir(),
	})

	sub := d.SubscribeSessionStateChanges("test-reset-session-state")
	require.NotNil(t, sub)
	defer sub.Unsubscribe()

	ok, err := d.ResetByAgent("agent-a")
	require.NoError(t, err)
	require.True(t, ok)

	select {
	case got := <-sub.Events():
		require.Equal(t, runtimeevents.KindSessionStateChanged, got.Kind)
		require.Equal(t, "agent-a", got.Context.AgentID)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for reset session state event")
	}
}

func TestDispatcherRestartByAgentRetiresSessionWithoutClearingWorkspace(t *testing.T) {
	d := mustNewDispatcher("http://127.0.0.1:0")
	defer d.Close()
	workDir := t.TempDir()
	d.sessionManager.Set("agent-a", &interactiveSession{
		sessionKey:    "agent-a",
		agentID:       "agent-a",
		serverAgentID: "server-agent-a",
		workDir:       workDir,
	})

	sub := d.SubscribeSessionStateChanges("test-restart-session-state")
	require.NotNil(t, sub)
	defer sub.Unsubscribe()

	require.True(t, d.RestartByAgent("server-agent-a"))
	require.DirExists(t, workDir)
	_, ok := d.Session("agent-a")
	require.False(t, ok)

	select {
	case got := <-sub.Events():
		require.Equal(t, runtimeevents.KindSessionStateChanged, got.Kind)
		require.Equal(t, "agent-a", got.Context.AgentID)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for restart session state event")
	}
}

func TestDispatcherRefreshRunningAgentConfigUpdatesSessionSnapshot(t *testing.T) {
	d := mustNewDispatcher("http://127.0.0.1:0")
	defer d.Close()
	d.sessionManager.Set("agent-a", &interactiveSession{
		sessionKey:    "agent-a",
		agentID:       "agent-a",
		serverAgentID: "agent-a",
		agent: agent.Config{
			ID:          "agent-a",
			Name:        "Old Agent",
			Description: "old description",
			Runtime:     "codex",
			Model:       "old-model",
		},
	})

	sub := d.SubscribeSessionStateChanges("test-refresh-running-agent-config")
	require.NotNil(t, sub)
	defer sub.Unsubscribe()

	changed := d.RefreshRunningAgentConfig(agent.Config{
		ID:          "agent-a",
		Name:        "New Agent",
		Description: "new description",
		Runtime:     "opencode",
		Model:       "new-model",
	})

	require.True(t, changed)
	snapshots := d.Sessions()
	require.Len(t, snapshots, 1)
	require.Equal(t, "New Agent", snapshots[0].AgentName)

	select {
	case got := <-sub.Events():
		require.Equal(t, runtimeevents.KindSessionStateChanged, got.Kind)
		require.Equal(t, "agent-a", got.Context.AgentID)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for session state event")
	}
}

func TestDispatcherRefreshRunningAgentConfigClearsStaleLLMCredentials(t *testing.T) {
	d := mustNewDispatcher("http://127.0.0.1:0")
	defer d.Close()
	d.sessionManager.Set("agent-a", &interactiveSession{
		sessionKey:    "agent-a",
		agentID:       "agent-a",
		serverAgentID: "agent-a",
		agent: agent.Config{
			ID:      "agent-a",
			Name:    "Agent",
			Runtime: "codex",
			LLMMode: string(agent.LLMModeInline),
			BaseURL: "https://agentteam.example.com/api/v1/workspaces/wsp_1/llm/v1",
			APIKey:  "stale-agentteam-key",
			Headers: map[string]string{"x-test": "1"},
			WireAPI: "responses",
		},
	})

	changed := d.RefreshRunningAgentConfig(agent.Config{
		ID:      "agent-a",
		Name:    "Agent",
		Runtime: "codex",
		LLMMode: string(agent.LLMModeInline),
		BaseURL: "https://agentteam.example.com/api/v1/workspaces/wsp_1/llm/v1",
		APIKey:  "",
		Headers: map[string]string{},
		WireAPI: "responses",
	})

	require.True(t, changed)
	snapshots := d.Sessions()
	require.Len(t, snapshots, 1)

	sess, ok := d.sessionManager.Get("agent-a")
	require.True(t, ok)
	require.Empty(t, sess.agent.APIKey)
	require.Equal(t, "https://agentteam.example.com/api/v1/workspaces/wsp_1/llm/v1", sess.agent.BaseURL)
	require.Equal(t, "responses", sess.agent.WireAPI)
	require.Empty(t, sess.agent.Headers)
}

func TestDispatcherRefreshRunningAgentConfigNameOnlyUpdatePreservesLLMCredentials(t *testing.T) {
	d := mustNewDispatcher("http://127.0.0.1:0")
	defer d.Close()
	d.sessionManager.Set("agent-a", &interactiveSession{
		sessionKey:    "agent-a",
		agentID:       "agent-a",
		serverAgentID: "agent-a",
		agent: agent.Config{
			ID:      "agent-a",
			Name:    "Old Agent",
			Runtime: "codex",
			LLMMode: string(agent.LLMModeInline),
			BaseURL: "https://agentteam.example.com/api/v1/workspaces/wsp_1/llm/v1",
			APIKey:  "keep-me",
			Headers: map[string]string{"x-test": "1"},
			WireAPI: "responses",
		},
	})

	changed := d.RefreshRunningAgentConfig(agent.Config{
		ID:   "agent-a",
		Name: "New Agent",
	})

	require.True(t, changed)
	sess, ok := d.sessionManager.Get("agent-a")
	require.True(t, ok)
	require.Equal(t, "keep-me", sess.agent.APIKey)
	require.Equal(t, "https://agentteam.example.com/api/v1/workspaces/wsp_1/llm/v1", sess.agent.BaseURL)
	require.Equal(t, map[string]string{"x-test": "1"}, sess.agent.Headers)
	require.Equal(t, "responses", sess.agent.WireAPI)
}