package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

func TestCodexChildAgentTeamMCPToolEffectGateLoadsInventoryAnnotations(t *testing.T) {
	t.Parallel()

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		_, err := io.WriteString(w, `{
			"jsonrpc":"2.0",
			"id":1,
			"result":{"tools":[
				{"name":"message_history_read","inputSchema":{"type":"object"},"annotations":{"readOnlyHint":true}},
				{"name":"message_send","inputSchema":{"type":"object"}}
			]}
		}`)
		require.NoError(t, err)
	}))
	defer server.Close()

	driver := &CodexAppServerDriver{
		launchGeneration: 1,
		codexLaunchState: codexLaunchState{
			mcpConfigs: map[string]types.NamedMCPConfig{
				"agentteam": {Name: "agentteam", URL: server.URL},
			},
			httpClient: server.Client(),
		},
		codexThreadState: codexThreadState{threadID: "thread-primary"},
	}

	_, err := driver.dynamicToolsForMCPServers(context.Background())
	require.NoError(t, err)
	require.NoError(t, driver.authorizeMCPToolEffect("agentteam", "message_history_read", "thread-child"))
	require.ErrorContains(t, driver.authorizeMCPToolEffect("agentteam", "message_send", "thread-child"), "return result to leader")
	require.ErrorContains(t, driver.authorizeMCPToolEffect("agentteam", "future_write", "thread-child"), "return result to leader")
}

func TestCodexChildAgentTeamMCPToolEffectGateUsesTrustedInventory(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name       string
		threadID   string
		tool       string
		definition *appServerMCPToolDef
		wantCalls  int32
		wantError  string
	}{
		{
			name:     "primary may call write tool",
			threadID: "thread-primary",
			tool:     "message_send",
			definition: &appServerMCPToolDef{
				Name: "message_send",
			},
			wantCalls: 1,
		},
		{
			name:     "child may call annotated read tool",
			threadID: "thread-child",
			tool:     "message_history_read",
			definition: &appServerMCPToolDef{
				Name:        "message_history_read",
				Annotations: &types.MCPToolAnnotation{ReadOnlyHint: true},
			},
			wantCalls: 1,
		},
		{
			name:     "child may not call known write tool",
			threadID: "thread-child",
			tool:     "message_send",
			definition: &appServerMCPToolDef{
				Name: "message_send",
			},
			wantError: "return result to leader",
		},
		{
			name:     "child may not set reactions",
			threadID: "thread-child",
			tool:     "message_reaction_set",
			definition: &appServerMCPToolDef{
				Name: "message_reaction_set",
			},
			wantError: "return result to leader",
		},
		{
			name:     "child may not upload attachments",
			threadID: "thread-child",
			tool:     "message_attachment_upload",
			definition: &appServerMCPToolDef{
				Name: "message_attachment_upload",
			},
			wantError: "return result to leader",
		},
		{
			name:      "child may not call unknown future tool",
			threadID:  "thread-child",
			tool:      "future_write",
			wantError: "return result to leader",
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()

			var calls atomic.Int32
			client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
				calls.Add(1)
				return codexChildGateHTTPResult(), nil
			})}
			driver := &CodexAppServerDriver{
				launchGeneration: 1,
				codexLaunchState: codexLaunchState{
					mcpConfigs: map[string]types.NamedMCPConfig{
						"agentteam": {Name: "agentteam", URL: "http://agentteam.invalid"},
					},
				},
				codexThreadState: codexThreadState{threadID: "thread-primary"},
			}
			if test.definition != nil {
				require.True(t, driver.replaceTrustedMCPToolEffects(
					driver.launchGeneration,
					"agentteam",
					[]appServerMCPToolDef{*test.definition},
				))
			}

			_, err := driver.callMCPToolWithBinding(
				context.Background(),
				client,
				types.NamedMCPConfig{Name: "agentteam", URL: "http://agentteam.invalid"},
				"agentteam",
				test.tool,
				json.RawMessage(`{}`),
				test.threadID,
			)
			if test.wantError != "" {
				require.ErrorContains(t, err, test.wantError)
			} else {
				require.NoError(t, err)
			}
			require.Equal(t, test.wantCalls, calls.Load(), "gate must run before AgentTeam MCP HTTP")
		})
	}
}

func TestCodexChildAgentTeamMCPToolEffectInventoryReplacementFailsClosed(t *testing.T) {
	t.Parallel()

	driver := &CodexAppServerDriver{
		launchGeneration: 1,
		codexThreadState: codexThreadState{threadID: "thread-primary"},
	}
	require.True(t, driver.replaceTrustedMCPToolEffects(1, "agentteam", []appServerMCPToolDef{
		{Name: "removed_read", Annotations: &types.MCPToolAnnotation{ReadOnlyHint: true}},
		{Name: "changed_read", Annotations: &types.MCPToolAnnotation{ReadOnlyHint: true}},
	}))
	require.NoError(t, driver.authorizeMCPToolEffect("agentteam", "removed_read", "thread-child"))
	require.NoError(t, driver.authorizeMCPToolEffect("agentteam", "changed_read", "thread-child"))

	require.True(t, driver.replaceTrustedMCPToolEffects(1, "agentteam", []appServerMCPToolDef{
		{Name: "changed_read"},
		{Name: "current_read", Annotations: &types.MCPToolAnnotation{ReadOnlyHint: true}},
	}))
	require.ErrorContains(t, driver.authorizeMCPToolEffect("agentteam", "removed_read", "thread-child"), "return result to leader")
	require.ErrorContains(t, driver.authorizeMCPToolEffect("agentteam", "changed_read", "thread-child"), "return result to leader")
	require.NoError(t, driver.authorizeMCPToolEffect("agentteam", "current_read", "thread-child"))
}

func TestCodexChildAgentTeamMCPToolEffectInventoryRejectsOldGenerationPublish(t *testing.T) {
	t.Parallel()

	started := make(chan struct{})
	release := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		close(started)
		<-release
		_, _ = io.WriteString(w, `{
			"jsonrpc":"2.0",
			"id":1,
			"result":{"tools":[
				{"name":"stale_read","inputSchema":{"type":"object"},"annotations":{"readOnlyHint":true}}
			]}
		}`)
	}))
	defer server.Close()

	driver := &CodexAppServerDriver{
		launchGeneration: 1,
		codexLaunchState: codexLaunchState{
			mcpConfigs: map[string]types.NamedMCPConfig{
				"agentteam": {Name: "agentteam", URL: server.URL},
			},
			httpClient: server.Client(),
		},
		codexThreadState: codexThreadState{threadID: "thread-primary"},
	}

	result := make(chan error, 1)
	go func() {
		_, err := driver.dynamicToolsForMCPServers(context.Background())
		result <- err
	}()
	<-started

	driver.mu.Lock()
	driver.launchGeneration = 2
	driver.trustedMCPToolReadOnly = make(map[string]bool)
	driver.mu.Unlock()
	require.True(t, driver.replaceTrustedMCPToolEffects(2, "agentteam", []appServerMCPToolDef{{Name: "stale_read"}}))
	close(release)

	require.ErrorContains(t, <-result, "generation was superseded")
	require.ErrorContains(t, driver.authorizeMCPToolEffect("agentteam", "stale_read", "thread-child"), "return result to leader",
		"the delayed generation-1 readOnly annotation must not overwrite generation 2")
}

func TestCodexChildAgentTeamMCPWriteDeniedForBothToolCallTransports(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name   string
		method string
		params string
	}{
		{
			name:   "mcp server tool call",
			method: "mcpServer/tool/call",
			params: `{"threadId":"thread-child","turnId":"turn-child","callId":"call-child","server":"agentteam","tool":"message_send","arguments":{}}`,
		},
		{
			name:   "dynamic tool call",
			method: "item/tool/call",
			params: `{"threadId":"thread-child","turnId":"turn-child","callId":"call-child","tool":"mcp_agentteam_message_send","arguments":{}}`,
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()

			var calls atomic.Int32
			driver := &CodexAppServerDriver{
				launchGeneration: 1,
				codexLaunchState: codexLaunchState{
					mcpConfigs: map[string]types.NamedMCPConfig{
						"agentteam": {Name: "agentteam", URL: "http://agentteam.invalid"},
					},
					httpClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
						calls.Add(1)
						return codexChildGateHTTPResult(), nil
					})},
				},
				codexThreadState: codexThreadState{
					threadID: "thread-primary",
					toolActivities: map[string]toolActivityRef{
						"call-child": {
							threadID: "thread-child",
							turnID:   "turn-child",
							callID:   "call-child",
						},
					},
				},
			}
			require.True(t, driver.replaceTrustedMCPToolEffects(
				driver.launchGeneration,
				"agentteam",
				[]appServerMCPToolDef{{Name: "message_send"}},
			))

			result, rpcErr := driver.executeServerRequestWithActivityForGeneration(
				context.Background(),
				1,
				"request-child",
				test.method,
				json.RawMessage(test.params),
			)
			if test.method == "item/tool/call" {
				require.Nil(t, rpcErr)
				response, ok := result.(appServerDynamicToolCallResponse)
				require.True(t, ok)
				require.False(t, response.Success)
				require.Len(t, response.ContentItems, 1)
				require.Contains(t, response.ContentItems[0].Text, "return result to leader")
			} else {
				require.NotNil(t, rpcErr)
				require.Contains(t, rpcErr.Message, "return result to leader")
			}
			require.Zero(t, calls.Load(), "child write must be rejected before AgentTeam MCP HTTP")
		})
	}
}

func codexChildGateHTTPResult() *http.Response {
	body := []byte(`{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"ok"}]}}`)
	return &http.Response{
		StatusCode: http.StatusOK,
		Header:     make(http.Header),
		Body:       io.NopCloser(bytes.NewReader(body)),
	}
}