package lookup

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

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

type fakeResolver struct {
	workspaceID string
	resolved    bool
	err         error
}

type countingResolver struct{ calls int }

func (resolver *countingResolver) ResolveWorkspaceRouteIdentity(context.Context, string, string, string, string, string, string, string) (string, bool, error) {
	resolver.calls++
	return "wsp_a", true, nil
}

func (resolver fakeResolver) ResolveWorkspaceRouteIdentity(context.Context, string, string, string, string, string, string, string) (string, bool, error) {
	return resolver.workspaceID, resolver.resolved, resolver.err
}

type fakeMachineAuthenticator struct {
	identity agent.MachineIdentity
	err      error
}

func (fake fakeMachineAuthenticator) AuthenticateMachine(context.Context, string) (agent.MachineIdentity, error) {
	return fake.identity, fake.err
}

type fakeAgentAuthenticator struct {
	identity agent.AuthenticatedAgent
	err      error
}

func (fake fakeAgentAuthenticator) AuthenticateAgent(context.Context, string, string) (agent.AuthenticatedAgent, error) {
	return fake.identity, fake.err
}

type fakeAppRunnerAuthenticator struct {
	workspaceID string
	resolved    bool
	err         error
}

func (fake fakeAppRunnerAuthenticator) ResolveAppRuntimeWorkspaceByBearerToken(context.Context, string) (string, bool, error) {
	return fake.workspaceID, fake.resolved, fake.err
}

func TestLegacyRuntimeResolverUsesCredentialWorkspace(t *testing.T) {
	t.Parallel()
	resolver := NewLegacyRuntimeResolver(
		fakeMachineAuthenticator{identity: agent.MachineIdentity{WorkspaceID: "wsp_machine"}},
		fakeAgentAuthenticator{identity: agent.AuthenticatedAgent{WorkspaceID: "wsp_agent"}},
		fakeAppRunnerAuthenticator{workspaceID: "wsp_app", resolved: true},
	)
	for _, test := range []struct {
		path          string
		rawQuery      string
		authorization string
		workspaceID   string
	}{
		{path: "/daemon/connect", rawQuery: "key=machine-key", workspaceID: "wsp_machine"},
		{path: "/internal/agent/agt_1/send", authorization: "Bearer agent-token", workspaceID: "wsp_agent"},
		{path: "/internal/app/runtime/heartbeat", authorization: "Bearer app-token", workspaceID: "wsp_app"},
	} {
		workspaceID, resolved, err := resolver.ResolveWorkspaceRouteIdentity(
			context.Background(), http.MethodPost, test.path, "", test.rawQuery, test.authorization, "", "",
		)
		require.NoError(t, err)
		require.True(t, resolved, test.path)
		require.Equal(t, test.workspaceID, workspaceID, test.path)
	}
}

func TestLegacyRuntimeResolverRejectsMissingOrInvalidCredential(t *testing.T) {
	t.Parallel()
	resolver := NewLegacyRuntimeResolver(
		fakeMachineAuthenticator{err: storeerr.ErrUnauthorized},
		fakeAgentAuthenticator{err: storeerr.ErrUnauthorized},
		fakeAppRunnerAuthenticator{err: storeerr.ErrUnauthorized},
	)
	for _, test := range []struct {
		path          string
		rawQuery      string
		authorization string
	}{
		{path: "/daemon/connect"},
		{path: "/daemon/connect", rawQuery: "key=invalid"},
		{path: "/internal/agent/agt_1/send"},
		{path: "/internal/agent/agt_1/send", authorization: "Bearer invalid"},
		{path: "/internal/app/runtime/heartbeat"},
		{path: "/internal/app/runtime/heartbeat", authorization: "Bearer invalid"},
	} {
		workspaceID, resolved, err := resolver.ResolveWorkspaceRouteIdentity(
			context.Background(), http.MethodPost, test.path, "", test.rawQuery, test.authorization, "", "",
		)
		require.ErrorIs(t, err, storeerr.ErrUnauthorized, test.path)
		require.False(t, resolved, test.path)
		require.Empty(t, workspaceID, test.path)
	}
}

func TestWorkspaceRouteIdentityEndpointRequiresInternalCredentialAndFailsClosed(t *testing.T) {
	t.Parallel()
	gin.SetMode(gin.ReleaseMode)
	for _, test := range []struct {
		name     string
		token    string
		resolver fakeResolver
		want     int
	}{
		{name: "missing service token", resolver: fakeResolver{workspaceID: "wsp_a", resolved: true}, want: http.StatusUnauthorized},
		{name: "unresolved", token: "secret", resolver: fakeResolver{}, want: http.StatusUnprocessableEntity},
		{name: "public credential rejected", token: "secret", resolver: fakeResolver{err: api.ErrUnauthorized}, want: http.StatusUnauthorized},
		{name: "resolved", token: "secret", resolver: fakeResolver{workspaceID: "wsp_a", resolved: true}, want: http.StatusOK},
	} {
		t.Run(test.name, func(t *testing.T) {
			router := gin.New()
			RegisterRoutes(router, Dependencies{Resolver: test.resolver, InternalAdminToken: "secret"})
			request := httptest.NewRequest(http.MethodPost, "/internal/meshxruntime/workspace-route-identity", strings.NewReader(`{"method":"GET","path":"/owner/repo.git/info/refs","authorization":"Basic value"}`))
			request.Header.Set("Content-Type", "application/json")
			if test.token != "" {
				request.Header.Set(api.HeaderInternalAdminToken, test.token)
			}
			recorder := httptest.NewRecorder()
			router.ServeHTTP(recorder, request)
			require.Equal(t, test.want, recorder.Code)
			if test.want == http.StatusOK {
				require.JSONEq(t, `{"workspaceId":"wsp_a"}`, recorder.Body.String())
			}
			if test.name == "public credential rejected" {
				require.Contains(t, recorder.Body.String(), "workspace_identity_unauthorized")
			}
		})
	}
}

func TestWorkspaceRouteIdentityEndpointRejectsInvalidRawPathAndOversizeBeforeResolver(t *testing.T) {
	gin.SetMode(gin.ReleaseMode)
	resolver := &countingResolver{}
	router := gin.New()
	RegisterRoutes(router, Dependencies{Resolver: resolver, InternalAdminToken: "secret"})
	for _, body := range []string{
		`{"method":"GET","path":"/api/v3/repos/owner/repo/archive/refs/tags/v1.zip","rawPath":"/api/v3/repos/owner/repo/archive/refs/tags/v2.zip"}`,
		`{"method":"GET","path":"/api/v3/user","authorization":"` + strings.Repeat("x", maxIdentityBodyBytes) + `"}`,
	} {
		request := httptest.NewRequest(http.MethodPost, "/internal/meshxruntime/workspace-route-identity", strings.NewReader(body))
		request.Header.Set("Content-Type", "application/json")
		request.Header.Set(api.HeaderInternalAdminToken, "secret")
		recorder := httptest.NewRecorder()
		router.ServeHTTP(recorder, request)
		require.Equal(t, http.StatusBadRequest, recorder.Code)
	}
	require.Zero(t, resolver.calls)
}