package forwarder

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api/agent"
)

func TestAgentThreadGatewayUsesActingUserIdentityForUpstreamJWTWhenPresent(t *testing.T) {
	var seenAuthorization string
	var seenActingUserAuthorization string
	var seenAttestation string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seenAuthorization = r.Header.Get("Authorization")
		seenActingUserAuthorization = r.Header.Get(AgentThreadActingUserVerifiedJWTHeader)
		seenAttestation = r.Header.Get(AgentThreadAttestationHeader)
		_ = json.NewEncoder(w).Encode(map[string]bool{"ok": true})
	}))
	defer upstream.Close()

	gateway, err := NewAgentThreadGateway(Dependencies{
		BaseURL:            upstream.URL,
		HTTPClient:         upstream.Client(),
		TrustedJWTAudience: "meshx",
		Now:                func() time.Time { return time.Unix(1700000000, 0) },
	})
	require.NoError(t, err)

	var raw json.RawMessage
	raw, err = gateway.ForwardAgentThreadUpstream(context.Background(), agent.AgentThreadAttestation{
		WorkspaceID:       "wsp_1",
		ThreadChannelID:   "thr_1",
		AgentID:           "agt_reader",
		AgentMemberID:     "mem_agent",
		ActingUserID:      "usr_human",
		ActingDisplayName: "Human Reviewer",
	}, http.MethodPost, "/api/v1/kb/retrieval/knowledge/search", json.RawMessage(`{"operation":"ping"}`))
	require.NoError(t, err)
	require.JSONEq(t, `{"ok":true}`, string(raw))
	require.NotEmpty(t, seenAttestation)
	attestationJSON := decodeRawURLBase64(t, seenAttestation)
	require.Contains(t, attestationJSON, `"workspaceId":"wsp_1"`)
	require.Contains(t, attestationJSON, `"threadChannelId":"thr_1"`)
	require.Contains(t, attestationJSON, `"agentId":"agt_reader"`)
	require.NotContains(t, attestationJSON, "workspace_id")
	require.NotContains(t, attestationJSON, "thread_channel_id")
	require.NotContains(t, attestationJSON, "agent_id")
	require.NotContains(t, attestationJSON, "knowledge_base_ids")
	require.NotContains(t, attestationJSON, "query")

	// When ActingUserID is present, the Authorization JWT must carry the human
	// subject (not the agent). meshx HandleAuth reads Authorization, so this is
	// what makes meshx see the human caller.
	claims := decodeClaims(t, strings.TrimPrefix(seenAuthorization, "Bearer "))
	require.Equal(t, "usr_human", claims["sub"])
	require.Equal(t, "usr_human", claims["agentteam_user_id"])
	require.Equal(t, "Human Reviewer", claims["name"])
	require.Equal(t, "usr-human", claims["preferred_username"])

	// The X-AgentTeam-Acting-User-Verified-JWT header still carries the human JWT.
	actingClaims := decodeClaims(t, strings.TrimPrefix(seenActingUserAuthorization, "Bearer "))
	require.Equal(t, "usr_human", actingClaims["sub"])
	require.Equal(t, "usr_human", actingClaims["agentteam_user_id"])
	require.Equal(t, "Human Reviewer", actingClaims["name"])
	require.Equal(t, "usr-human", actingClaims["preferred_username"])
}

func TestAgentThreadGatewayUsesAgentIdentityWhenNoActingUser(t *testing.T) {
	var seenAuthorization string
	var seenActingUserAuthorization string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seenAuthorization = r.Header.Get("Authorization")
		seenActingUserAuthorization = r.Header.Get(AgentThreadActingUserVerifiedJWTHeader)
		_ = json.NewEncoder(w).Encode(map[string]bool{"ok": true})
	}))
	defer upstream.Close()

	gateway, err := NewAgentThreadGateway(Dependencies{
		BaseURL:            upstream.URL,
		HTTPClient:         upstream.Client(),
		TrustedJWTAudience: "meshx",
		Now:                func() time.Time { return time.Unix(1700000000, 0) },
	})
	require.NoError(t, err)

	// No ActingUserID: the agent identity must be used for Authorization,
	// preserving existing behavior for non-human-triggered upstream calls.
	var raw json.RawMessage
	raw, err = gateway.ForwardAgentThreadUpstream(context.Background(), agent.AgentThreadAttestation{
		WorkspaceID:     "wsp_1",
		ThreadChannelID: "thr_1",
		AgentID:         "agt_reader",
		AgentMemberID:   "mem_agent",
	}, http.MethodPost, "/api/v1/kb/retrieval/search", json.RawMessage(`{}`))
	require.NoError(t, err)
	require.JSONEq(t, `{"ok":true}`, string(raw))

	claims := decodeClaims(t, strings.TrimPrefix(seenAuthorization, "Bearer "))
	require.Equal(t, "agt_reader", claims["sub"])
	require.Equal(t, "agt_reader", claims["agentteam_user_id"])

	// No acting-user header when ActingUserID is absent.
	require.Empty(t, seenActingUserAuthorization)
}

func TestAgentThreadGatewayUsesGenericUpstreamPathValidation(t *testing.T) {
	var seenPath string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seenPath = r.URL.Path
		_ = json.NewEncoder(w).Encode(map[string]bool{"ok": true})
	}))
	defer upstream.Close()

	gateway, err := NewAgentThreadGateway(Dependencies{BaseURL: upstream.URL, HTTPClient: upstream.Client()})
	require.NoError(t, err)

	var raw json.RawMessage
	raw, err = gateway.ForwardAgentThreadUpstream(context.Background(), agent.AgentThreadAttestation{WorkspaceID: "wsp_1", ThreadChannelID: "thr_1", AgentID: "agt_reader"}, http.MethodPost, "/api/v1/kb/retrieval/knowledge/search", json.RawMessage(`{}`))
	require.NoError(t, err)
	require.Equal(t, "/api/v1/kb/retrieval/knowledge/search", seenPath)
	require.JSONEq(t, `{"ok":true}`, string(raw))

	_, err = gateway.ForwardAgentThreadUpstream(context.Background(), agent.AgentThreadAttestation{WorkspaceID: "wsp_1", ThreadChannelID: "thr_1", AgentID: "agt_reader"}, http.MethodPost, "https://meshx.invalid/api/v1/kb/retrieval/knowledge/search", json.RawMessage(`{}`))
	require.Error(t, err)
	require.Contains(t, err.Error(), "agent thread upstream path is invalid")
}

func TestAgentThreadGatewayRejectsMissingWorkspaceBeforeUpstream(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("missing-workspace request reached upstream")
	}))
	t.Cleanup(upstream.Close)
	gateway, err := NewAgentThreadGateway(Dependencies{BaseURL: upstream.URL, HTTPClient: upstream.Client()})
	require.NoError(t, err)

	_, err = gateway.ForwardAgentThreadUpstream(context.Background(), agent.AgentThreadAttestation{AgentID: "agt_reader"}, http.MethodPost, "/api/v1/kb/retrieval/knowledge/search", json.RawMessage(`{}`))
	require.ErrorContains(t, err, "workspace is required")
}

func decodeRawURLBase64(t *testing.T, value string) string {
	t.Helper()
	payload, err := base64.RawURLEncoding.DecodeString(value)
	require.NoError(t, err)
	return string(payload)
}