package agent

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

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/api"
	storeagent "github.com/tidbcloud/AgentTeam/internal/store/agent"
)

type fakeAgentThreadAgentAuth struct{}

func (fakeAgentThreadAgentAuth) AuthenticateAgent(context.Context, string, string) (storeagent.AuthenticatedAgent, error) {
	return storeagent.AuthenticatedAgent{AgentID: "agt_000000000000kb", WorkspaceID: "ws_1", MemberID: "0000000agent"}, nil
}

type fakeAgentThreadContextStore struct {
	attestation AgentThreadAttestation
	err         error
	target      *AgentThreadTarget
}

func (store fakeAgentThreadContextStore) ResolveAgentThreadAttestation(_ context.Context, _ storeagent.AuthenticatedAgent, target AgentThreadTarget) (AgentThreadAttestation, error) {
	if store.target != nil {
		*store.target = target
	}
	return store.attestation, store.err
}

type fakeAgentThreadGateway struct {
	requests []fakeAgentThreadGatewayRequest
	err      error
}

type fakeAgentThreadGatewayRequest struct {
	method string
	path   string
	body   json.RawMessage
}

func (client *fakeAgentThreadGateway) ForwardAgentThreadUpstream(_ context.Context, _ AgentThreadAttestation, method string, path string, request json.RawMessage) (json.RawMessage, error) {
	client.requests = append(client.requests, fakeAgentThreadGatewayRequest{method: method, path: path, body: request})
	if client.err != nil {
		return nil, client.err
	}
	return json.RawMessage(`{"hits":[{"title":"Wiki MCP","citation_uri":"kb://kb_1/sources/source_1/wiki/wiki-mcp"}],"total":1,"has_more":false}`), nil
}

func TestAgentThreadHandlerUpstreamAcceptsEveryAgentDeliveryTarget(t *testing.T) {
	gin.SetMode(gin.TestMode)
	for _, targetType := range []string{"channel", "dm", "thread", "task"} {
		t.Run(targetType, func(t *testing.T) {
			var resolved AgentThreadTarget
			gateway := &fakeAgentThreadGateway{}
			handler := &AgentThreadHandler{
				auth: fakeAgentThreadAgentAuth{},
				contexts: fakeAgentThreadContextStore{
					target: &resolved,
					attestation: AgentThreadAttestation{
						WorkspaceID: "ws_1", TargetKey: "scope_1", TargetType: targetType, MessageID: "msg_1",
						ThreadChannelID: "scope_1", AgentID: "agt_000000000000kb", AgentMemberID: "0000000agent", ActingUserID: "usr_00000000000001",
					},
				},
				gateway: gateway,
			}
			router := gin.New()
			router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/agent-thread/upstream", handler.Upstream)

			body := `{"target_key":"scope_1","target_type":"` + targetType + `","message_id":"msg_1","path":"/api/v1/kb/retrieval/knowledge/context","body":{}}`
			req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/agents/agt_000000000000kb/runtime/agent-thread/upstream", bytes.NewBufferString(body))
			req.Header.Set("Authorization", "Bearer token")
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
			require.Equal(t, AgentThreadTarget{TargetKey: "scope_1", TargetType: targetType, MessageID: "msg_1"}, resolved)
			require.Len(t, gateway.requests, 1)
		})
	}
}

func TestAgentThreadHandlerUpstreamAllowsMessageOnlyTargetResolution(t *testing.T) {
	gin.SetMode(gin.TestMode)
	var resolved AgentThreadTarget
	gateway := &fakeAgentThreadGateway{}
	handler := &AgentThreadHandler{
		auth: fakeAgentThreadAgentAuth{},
		contexts: fakeAgentThreadContextStore{
			target: &resolved,
			attestation: AgentThreadAttestation{
				WorkspaceID: "ws_1", TargetKey: "chn_1", TargetType: "channel", MessageID: "msg_1",
				ThreadChannelID: "chn_1", AgentID: "agt_000000000000kb", AgentMemberID: "0000000agent", ActingUserID: "usr_00000000000001",
			},
		},
		gateway: gateway,
	}
	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/agent-thread/upstream", handler.Upstream)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/agents/agt_000000000000kb/runtime/agent-thread/upstream", bytes.NewBufferString(`{"message_id":"msg_1","path":"/api/v1/kb/retrieval/knowledge/search","body":{"query":"mcp"}}`))
	req.Header.Set("Authorization", "Bearer token")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
	require.Equal(t, AgentThreadTarget{MessageID: "msg_1"}, resolved)
	require.Len(t, gateway.requests, 1)
}

func TestAgentThreadHandlerRejectsAttestationWorkspaceBeforeGateway(t *testing.T) {
	gin.SetMode(gin.TestMode)
	for _, test := range []struct {
		name        string
		workspaceID string
	}{
		{name: "missing"},
		{name: "conflicting", workspaceID: "ws_other"},
	} {
		t.Run(test.name, func(t *testing.T) {
			gateway := &fakeAgentThreadGateway{}
			handler := &AgentThreadHandler{
				auth: fakeAgentThreadAgentAuth{},
				contexts: fakeAgentThreadContextStore{attestation: AgentThreadAttestation{
					WorkspaceID: test.workspaceID, TargetKey: "chn_1", AgentID: "agt_000000000000kb", AgentMemberID: "0000000agent",
				}},
				gateway: gateway,
			}
			router := gin.New()
			router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/agent-thread/upstream", handler.Upstream)
			req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/agents/agt_000000000000kb/runtime/agent-thread/upstream", bytes.NewBufferString(`{"target_key":"chn_1","path":"/api/v1/kb/retrieval/knowledge/search","body":{}}`))
			req.Header.Set("Authorization", "Bearer token")
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, http.StatusForbidden, rec.Code, rec.Body.String())
			require.Empty(t, gateway.requests)
		})
	}
}

func TestAgentThreadHandlerReturnsStructuredFailureStage(t *testing.T) {
	gin.SetMode(gin.TestMode)
	tests := []struct {
		name       string
		contextErr error
		gatewayErr error
		status     int
		code       string
		stage      string
	}{
		{name: "target missing", contextErr: api.ErrNotFound, status: http.StatusNotFound, code: "target_scope_not_found", stage: "target_scope_resolution"},
		{name: "target forbidden", contextErr: api.ErrForbidden, status: http.StatusForbidden, code: "target_scope_not_authorized", stage: "target_scope_resolution"},
		{name: "meshx missing", gatewayErr: api.ErrNotFound, status: http.StatusNotFound, code: "meshx_upstream_not_found", stage: "meshx_upstream"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			handler := &AgentThreadHandler{
				auth: fakeAgentThreadAgentAuth{},
				contexts: fakeAgentThreadContextStore{
					err:         test.contextErr,
					attestation: AgentThreadAttestation{WorkspaceID: "ws_1", TargetKey: "chn_1", ThreadChannelID: "chn_1", AgentID: "agt_000000000000kb", AgentMemberID: "0000000agent"},
				},
				gateway: &fakeAgentThreadGateway{err: test.gatewayErr},
			}
			router := gin.New()
			router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/agent-thread/upstream", handler.Upstream)
			req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/agents/agt_000000000000kb/runtime/agent-thread/upstream", bytes.NewBufferString(`{"target_key":"chn_1","target_type":"channel","message_id":"msg_1","path":"/api/v1/kb/retrieval/knowledge/context","body":{}}`))
			req.Header.Set("Authorization", "Bearer token")
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, test.status, rec.Code, rec.Body.String())
			var response struct {
				Code    string `json:"code"`
				Details struct {
					Stage string `json:"stage"`
				} `json:"details"`
			}
			require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &response))
			require.Equal(t, test.code, response.Code)
			require.Equal(t, test.stage, response.Details.Stage)
		})
	}
}

func TestAgentThreadHandlerAttestationReturnsAgentTeamScopeWithoutUpstreamLookup(t *testing.T) {
	gin.SetMode(gin.TestMode)
	gateway := &fakeAgentThreadGateway{}
	handler := &AgentThreadHandler{
		auth: fakeAgentThreadAgentAuth{},
		contexts: fakeAgentThreadContextStore{attestation: AgentThreadAttestation{
			WorkspaceID:       "ws_1",
			TargetKey:         "thr_kb",
			TargetType:        "thread",
			MessageID:         "msg_kb",
			ThreadChannelID:   "thr_kb",
			AgentID:           "agt_000000000000kb",
			AgentMemberID:     "0000000agent",
			ActingUserID:      "usr_00000000000001",
			ActingMemberID:    "0000000human",
			ActingDisplayName: "Human",
		}},
		gateway: gateway,
	}
	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/agent-thread/attestation", handler.Attestation)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/agents/agt_000000000000kb/runtime/agent-thread/attestation", bytes.NewBufferString(`{"thread_channel_id":"thr_kb"}`))
	req.Header.Set("Authorization", "Bearer token")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Empty(t, gateway.requests)
	var body map[string]any
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
	require.NotContains(t, rec.Body.String(), "kb_ids")
	require.NotContains(t, rec.Body.String(), "query")
	require.NotContains(t, body, "knowledge"+"_bases")
	require.NotContains(t, rec.Body.String(), `"repo_full_name"`)
	require.Equal(t, "thr_kb", body["target_key"])
	require.Equal(t, "thread", body["target_type"])
	require.Equal(t, "msg_kb", body["message_id"])
}

func TestAgentThreadHandlerUpstreamForwardsRawBodyWithThreadAttestation(t *testing.T) {
	gin.SetMode(gin.TestMode)
	gateway := &fakeAgentThreadGateway{}
	handler := &AgentThreadHandler{
		auth: fakeAgentThreadAgentAuth{},
		contexts: fakeAgentThreadContextStore{attestation: AgentThreadAttestation{
			WorkspaceID:       "ws_1",
			ThreadChannelID:   "thr_kb",
			AgentID:           "agt_000000000000kb",
			AgentMemberID:     "0000000agent",
			ActingUserID:      "usr_00000000000001",
			ActingMemberID:    "0000000human",
			ActingDisplayName: "Human",
		}},
		gateway: gateway,
	}
	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/agent-thread/upstream", handler.Upstream)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/agents/agt_000000000000kb/runtime/agent-thread/upstream", bytes.NewBufferString(`{"thread_channel_id":"thr_kb","method":"POST","path":"/api/v1/kb/retrieval/knowledge/search","body":{"query":"mcp","limit":5,"kb_ids":["kb_1"]}}`))
	req.Header.Set("Authorization", "Bearer token")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Len(t, gateway.requests, 1)
	require.Equal(t, "POST", gateway.requests[0].method)
	require.Equal(t, "/api/v1/kb/retrieval/knowledge/search", gateway.requests[0].path)
	requireAgentThreadBody(t, gateway.requests[0].body, map[string]any{
		"kb_ids": []any{"kb_1"},
		"query":  "mcp",
		"limit":  float64(5),
	})
	require.Contains(t, rec.Body.String(), `"title":"Wiki MCP"`)
}

func TestAgentThreadHandlerDoesNotTreatSelectedKBIDsAsAgentTeamBoundary(t *testing.T) {
	gin.SetMode(gin.TestMode)
	gateway := &fakeAgentThreadGateway{}
	handler := &AgentThreadHandler{
		auth:     fakeAgentThreadAgentAuth{},
		contexts: fakeAgentThreadContextStore{attestation: AgentThreadAttestation{WorkspaceID: "ws_1", ThreadChannelID: "thr_kb", AgentID: "agt_000000000000kb", AgentMemberID: "0000000agent", ActingUserID: "usr_00000000000001"}},
		gateway:  gateway,
	}
	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/agent-thread/upstream", handler.Upstream)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/ws_1/agents/agt_000000000000kb/runtime/agent-thread/upstream", bytes.NewBufferString(`{"thread_channel_id":"thr_kb","path":"/api/v1/kb/retrieval/knowledge/search","body":{"query":"mcp","kb_ids":["kb_2"]}}`))
	req.Header.Set("Authorization", "Bearer token")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Len(t, gateway.requests, 1)
	requireAgentThreadBody(t, gateway.requests[0].body, map[string]any{
		"query":  "mcp",
		"kb_ids": []any{"kb_2"},
	})
}

func requireAgentThreadBody(t *testing.T, payload json.RawMessage, expected map[string]any) {
	t.Helper()
	var body map[string]any
	require.NoError(t, json.Unmarshal(payload, &body))
	for key, value := range expected {
		require.Equal(t, value, body[key], key)
	}
}