package gateway

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/api/identity/lookup"
	"github.com/tidbcloud/AgentTeam/internal/config"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
	"github.com/tidbcloud/AgentTeam/internal/workspace/router"
)

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

	var workspaceHits atomic.Int32
	workspaceServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		workspaceHits.Add(1)
		require.Equal(t, "wsp_a", request.Header.Get(api.HeaderRouteWorkspaceID))
		require.Equal(t, "wsp_a", request.Header.Get(api.HeaderWorkspaceID))
		require.Empty(t, request.Header.Get(config.DefaultTrustedJWTHeader))
		require.Empty(t, request.Header.Get("X-AgentTeam-Acting-User-Verified-JWT"))
		require.Empty(t, request.Header.Get("X-MeshX-Acting-Principal"))
		writer.WriteHeader(http.StatusNoContent)
	}))
	t.Cleanup(workspaceServer.Close)

	identity := &candidateIdentityResolver{memberships: map[string]bool{"wsp_a": true}}
	placement := &seamPlacementResolver{targets: map[string]workspace.WorkspaceRouteTarget{
		"wsp_a": {WorkspaceID: "wsp_a", WorkspaceGroupID: "group-a", PlacementState: workspace.PlacementStateActive},
	}}
	handler, err := NewHandler(placement, Config{
		CentralUpstream:    "http://127.0.0.1:1",
		WorkspaceUpstreams: map[string]string{"group-a": workspaceServer.URL},
		IdentityResolver:   identity,
	})
	require.NoError(t, err)

	for _, test := range []struct {
		method, path, authorization string
	}{
		{http.MethodGet, "/api/v1/workspaces/wsp_a/ags/api/v3/repos/acme/docs", "Bearer ags_rat_repo"},
		{http.MethodGet, "/api/v1/workspaces/wsp_a/ags/acme/docs/raw/main/README.md", ""},
		{http.MethodGet, "/api/v1/workspaces/wsp_a/ags/acme/docs.git/info/refs?service=git-upload-pack", "Bearer ags_rat_git"},
		{http.MethodGet, "/api/v1/workspaces/wsp_a/ags/avatars/alice", ""},
	} {
		request := httptest.NewRequest(test.method, test.path, nil)
		if test.authorization != "" {
			request.Header.Set("Authorization", test.authorization)
		}
		request.Header.Set(config.DefaultTrustedJWTHeader, "Bearer forged")
		request.Header.Set("X-AgentTeam-Acting-User-Verified-JWT", "Bearer forged")
		request.Header.Set("X-MeshX-Acting-Principal", "forged")
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, request)
		require.Equal(t, http.StatusNoContent, recorder.Code, test.path)
	}

	require.Zero(t, identity.calls.Load(), "path-routed credentials and public reads must not call session identity")
	require.EqualValues(t, 4, placement.calls.Load())
	require.EqualValues(t, 4, workspaceHits.Load())
}

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

	identity := &candidateIdentityResolver{memberships: map[string]bool{"wsp_a": true}}
	placement := &seamPlacementResolver{}
	handler, err := NewHandler(placement, Config{
		CentralUpstream:    "http://127.0.0.1:1",
		WorkspaceUpstreams: map[string]string{"group-a": "http://127.0.0.1:2"},
		IdentityResolver:   identity,
	})
	require.NoError(t, err)

	for _, path := range []string{
		"/api/v1/workspaces/wsp_a/ags/api/v3/user",
		"/api/v1/workspaces/wsp_a/ags/api/ext/v1/user/tokens",
		"/api/v1/workspaces/wsp_a/ags/api/v3/openapi.json",
		"/api/v1/workspaces/wsp_a/ags/api/ext/v1/analytics/events",
		"/api/v1/workspaces/wsp_a/ags/avatars/alice/extra",
	} {
		request := httptest.NewRequest(http.MethodGet, path, nil)
		request.Header.Set("Authorization", "Bearer "+testSessionToken())
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, request)
		require.Equal(t, http.StatusNotFound, recorder.Code, path)
	}
	require.Zero(t, identity.calls.Load())
	require.Zero(t, placement.calls.Load())
}

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

	_, err := NewHandler(&seamPlacementResolver{}, Config{
		CentralUpstream:    "http://127.0.0.1:1",
		WorkspaceUpstreams: map[string]string{"group-a": "http://127.0.0.1:2"},
		IdentityResolver:   &candidateIdentityResolver{},
		TrustedJWTHeader:   "X-Client-Controlled-Identity",
	})
	require.EqualError(t, err, `gateway trusted identity header: split deployment MESHX_AUTH_JWT_HEADER must be "X-AgentTeam-Verified-JWT"`)
}

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

	workspaceServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		require.Equal(t, "/api/v1/workspaces/wsp_a/ags/api/v3/user/repos", request.URL.Path)
		require.Empty(t, request.Header.Get(config.DefaultTrustedJWTHeader))
		require.Empty(t, request.Header.Get("X-AgentTeam-Acting-User-Verified-JWT"))
		require.Empty(t, request.Header.Get("X-MeshX-Acting-Principal"))
		writer.WriteHeader(http.StatusNoContent)
	}))
	t.Cleanup(workspaceServer.Close)

	identity := &candidateIdentityResolver{memberships: map[string]bool{"wsp_a": true}}
	placement := &seamPlacementResolver{targets: map[string]workspace.WorkspaceRouteTarget{
		"wsp_a": {WorkspaceID: "wsp_a", WorkspaceGroupID: "group-a", PlacementState: workspace.PlacementStateActive},
	}}
	handler, err := NewHandler(placement, Config{
		CentralUpstream: "http://127.0.0.1:1", WorkspaceUpstreams: map[string]string{"group-a": workspaceServer.URL}, IdentityResolver: identity,
	})
	require.NoError(t, err)

	request := httptest.NewRequest(http.MethodGet, "/api/v3/user/repos", nil)
	request.Header.Set("Authorization", "Bearer "+testSessionToken())
	request.Header.Set(api.HeaderRouteWorkspaceID, "wsp_a")
	request.Header.Set(config.DefaultTrustedJWTHeader, "Bearer forged")
	request.Header.Set("X-AgentTeam-Acting-User-Verified-JWT", "Bearer forged")
	request.Header.Set("X-MeshX-Acting-Principal", "forged")
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusNoContent, recorder.Code)
	require.EqualValues(t, 1, identity.calls.Load())
	require.EqualValues(t, 1, placement.calls.Load())
}

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

	var groupAHits, groupBHits, defaultHits, businessHits atomic.Int32
	workspaceServer := func(workspaceID, groupID string, hits *atomic.Int32) *httptest.Server {
		server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
			hits.Add(1)
			require.Equal(t, workspaceID, request.Header.Get(api.HeaderRouteWorkspaceID))
			require.Equal(t, workspaceID, request.Header.Get(api.HeaderWorkspaceID))
			require.Equal(t, groupID, request.Header.Get(router.HeaderResolvedWorkspaceGroupID))
			require.Equal(t, "/api/v1/workspaces/"+workspaceID+"/ags/api/v3/repos/acme/docs", request.URL.Path)
			businessHits.Add(1)
			http.NotFound(writer, request)
		}))
		t.Cleanup(server.Close)
		return server
	}
	groupA := workspaceServer("wsp_a", "group-a", &groupAHits)
	groupB := workspaceServer("wsp_b", "group-b", &groupBHits)
	defaultGroup := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { defaultHits.Add(1) }))
	t.Cleanup(defaultGroup.Close)
	central := httptest.NewServer(http.NotFoundHandler())
	t.Cleanup(central.Close)

	identity := &candidateIdentityResolver{memberships: map[string]bool{"wsp_a": true, "wsp_b": true}}
	placement := &seamPlacementResolver{targets: map[string]workspace.WorkspaceRouteTarget{
		"wsp_a": {WorkspaceID: "wsp_a", WorkspaceGroupID: "group-a"},
		"wsp_b": {WorkspaceID: "wsp_b", WorkspaceGroupID: "group-b"},
	}}
	handler, err := NewHandler(placement, Config{
		CentralUpstream: central.URL,
		WorkspaceUpstreams: map[string]string{
			"group-a": groupA.URL, "group-b": groupB.URL, "default": defaultGroup.URL,
		},
		IdentityResolver: identity,
	})
	require.NoError(t, err)

	for _, workspaceID := range []string{"wsp_a", "wsp_b"} {
		request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/"+workspaceID+"/ags/api/v3/repos/acme/docs", nil)
		request.Header.Set("Authorization", "Bearer "+testSessionToken())
		request.Header.Set(api.HeaderRouteWorkspaceID, "wsp_forged")
		request.Header.Set(api.HeaderWorkspaceID, "wsp_forged")
		request.Header.Set(router.HeaderResolvedWorkspaceGroupID, "group-forged")
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, request)
		require.Equal(t, http.StatusNotFound, recorder.Code)
	}

	require.EqualValues(t, 2, identity.calls.Load())
	require.EqualValues(t, 2, placement.calls.Load())
	require.EqualValues(t, 1, groupAHits.Load())
	require.EqualValues(t, 1, groupBHits.Load())
	require.Zero(t, defaultHits.Load())
	require.EqualValues(t, 2, businessHits.Load())
}

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

	central := httptest.NewServer(http.NotFoundHandler())
	t.Cleanup(central.Close)
	workspaceHits := atomic.Int32{}
	workspaceServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { workspaceHits.Add(1) }))
	t.Cleanup(workspaceServer.Close)

	for _, test := range []struct {
		name       string
		identity   *candidateIdentityResolver
		wantStatus int
	}{
		{name: "membership rejected", identity: &candidateIdentityResolver{}, wantStatus: http.StatusUnauthorized},
		{name: "verified workspace mismatch", identity: &candidateIdentityResolver{forcedWorkspace: "wsp_b", memberships: map[string]bool{"wsp_a": true}}, wantStatus: http.StatusBadRequest},
	} {
		t.Run(test.name, func(t *testing.T) {
			placement := &seamPlacementResolver{targets: map[string]workspace.WorkspaceRouteTarget{"wsp_a": {WorkspaceID: "wsp_a", WorkspaceGroupID: "group-a"}}}
			handler, err := NewHandler(placement, Config{
				CentralUpstream:    central.URL,
				WorkspaceUpstreams: map[string]string{"group-a": workspaceServer.URL}, IdentityResolver: test.identity,
			})
			require.NoError(t, err)
			request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_a/ags/api/v3/repos/acme/docs", nil)
			request.Header.Set("Authorization", "Bearer "+testSessionToken())
			request.Header.Set(api.HeaderRouteWorkspaceID, "wsp_forged")
			request.Header.Set(api.HeaderWorkspaceID, "wsp_forged")
			request.Header.Set(router.HeaderResolvedWorkspaceGroupID, "group-forged")
			recorder := httptest.NewRecorder()

			handler.ServeHTTP(recorder, request)

			require.Equal(t, test.wantStatus, recorder.Code)
			require.EqualValues(t, 1, test.identity.calls.Load())
			require.Zero(t, placement.calls.Load())
			require.Zero(t, workspaceHits.Load())
		})
	}
}

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

	identity := &candidateIdentityResolver{memberships: map[string]bool{"wsp_a": true}}
	placement := &seamPlacementResolver{}
	handler, err := NewHandler(placement, Config{
		CentralUpstream: "http://127.0.0.1:1", WorkspaceUpstreams: map[string]string{"group-a": "http://127.0.0.1:2"}, IdentityResolver: identity,
	})
	require.NoError(t, err)

	for _, values := range [][]string{
		{"Bearer " + testSessionToken(), "Bearer forged"},
		{"Bearer " + testSessionToken() + ",Bearer forged"},
		{"Bearer malformed"},
	} {
		request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_a/ags/api/v3/repos/acme/docs", nil)
		for _, value := range values {
			request.Header.Add("Authorization", value)
		}
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, request)
		require.Contains(t, []int{http.StatusBadRequest, http.StatusUnauthorized}, recorder.Code)
	}
	require.Zero(t, identity.calls.Load())
	require.Zero(t, placement.calls.Load())
}

func TestGatewayWorkspaceAGSSeamUsesRealSessionMembershipIdentityWire(t *testing.T) {
	gin.SetMode(gin.ReleaseMode)
	sessions := &wireCountingSessions{}
	memberships := &wireCountingMemberships{allowed: map[string]string{"wsp_a": "wsp_a"}}
	privateRouter := gin.New()
	lookup.RegisterRoutes(privateRouter, lookup.Dependencies{
		Resolver: lookup.NewAuthenticatedResolver(nil, sessions, memberships), InternalAdminToken: "private-token",
	})
	privateServer := httptest.NewServer(privateRouter)
	t.Cleanup(privateServer.Close)
	identityResolver, err := NewHTTPWorkspaceIdentityResolver(privateServer.URL+"/internal/meshxruntime/workspace-route-identity", "private-token", privateServer.Client())
	require.NoError(t, err)

	workspaceHits := atomic.Int32{}
	workspaceServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		workspaceHits.Add(1)
		writer.WriteHeader(http.StatusNotFound)
	}))
	t.Cleanup(workspaceServer.Close)
	placement := &seamPlacementResolver{targets: map[string]workspace.WorkspaceRouteTarget{"wsp_a": {WorkspaceID: "wsp_a", WorkspaceGroupID: "group-a"}}}
	handler, err := NewHandler(placement, Config{
		CentralUpstream: "http://127.0.0.1:1", WorkspaceUpstreams: map[string]string{"group-a": workspaceServer.URL}, IdentityResolver: identityResolver,
	})
	require.NoError(t, err)

	request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_a/ags/api/v3/repos/acme/docs", nil)
	request.Header.Set("Authorization", "Bearer "+testSessionToken())
	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusNotFound, recorder.Code)
	require.EqualValues(t, 1, sessions.calls.Load())
	require.EqualValues(t, 1, memberships.calls.Load())
	require.EqualValues(t, 1, placement.calls.Load())
	require.EqualValues(t, 1, workspaceHits.Load())

	for _, values := range [][]string{{"Bearer " + testSessionToken(), "Bearer forged"}, {"Bearer " + testSessionToken() + ",Bearer forged"}} {
		request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_a/ags/api/v3/repos/acme/docs", nil)
		for _, value := range values {
			request.Header.Add("Authorization", value)
		}
		handler.ServeHTTP(httptest.NewRecorder(), request)
	}
	require.EqualValues(t, 1, sessions.calls.Load())
	require.EqualValues(t, 1, memberships.calls.Load())
	require.EqualValues(t, 1, placement.calls.Load())
	require.EqualValues(t, 1, workspaceHits.Load())
}

func testSessionToken() string { return "sssssssssssssssssssssssssssssssssssssssssss" }

type candidateIdentityResolver struct {
	calls           atomic.Int32
	memberships     map[string]bool
	forcedWorkspace string
}

func (resolver *candidateIdentityResolver) ResolveWorkspaceIdentity(_ context.Context, request WorkspaceIdentityRequest) (RouteIdentity, error) {
	resolver.calls.Add(1)
	if request.Authority != "session_candidate" || !resolver.memberships[request.WorkspaceID] {
		return RouteIdentity{}, ErrWorkspaceIdentityUnauthorized
	}
	if resolver.forcedWorkspace != "" {
		return RouteIdentity{WorkspaceID: resolver.forcedWorkspace}, nil
	}
	return RouteIdentity{WorkspaceID: request.WorkspaceID}, nil
}

type seamPlacementResolver struct {
	calls   atomic.Int32
	targets map[string]workspace.WorkspaceRouteTarget
}

func (resolver *seamPlacementResolver) ResolveWorkspaceRouteTarget(_ context.Context, workspaceID string) (workspace.WorkspaceRouteTarget, error) {
	resolver.calls.Add(1)
	target, ok := resolver.targets[workspaceID]
	if !ok {
		return workspace.WorkspaceRouteTarget{}, fmt.Errorf("unexpected workspace %s", workspaceID)
	}
	return target, nil
}