package gateway

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

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

	type observedRequest struct {
		method        string
		path          string
		escapedPath   string
		query         string
		body          string
		routeID       string
		workspaceID   string
		authorization string
	}
	observed := make(chan observedRequest, 2)
	workspaceUpstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		body, _ := io.ReadAll(request.Body)
		observed <- observedRequest{
			method:        request.Method,
			path:          request.URL.Path,
			escapedPath:   request.URL.EscapedPath(),
			query:         request.URL.RawQuery,
			body:          string(body),
			routeID:       request.Header.Get(api.HeaderRouteWorkspaceID),
			workspaceID:   request.Header.Get(api.HeaderWorkspaceID),
			authorization: request.Header.Get("Authorization"),
		}
		writer.WriteHeader(http.StatusNoContent)
	}))
	t.Cleanup(workspaceUpstream.Close)

	handler := newTestHandler(t, "http://127.0.0.1:1", map[string]string{"workspace-default": workspaceUpstream.URL}, fakeResolver{
		target: workspace.WorkspaceRouteTarget{WorkspaceID: "wsp_1", WorkspaceGroupID: "workspace-default", PlacementState: workspace.PlacementStateActive},
	})
	for _, test := range []struct {
		name            string
		method          string
		url             string
		body            string
		wantPath        string
		wantEscapedPath string
		wantQuery       string
	}{
		{
			name:            "skill libraries",
			method:          http.MethodGet,
			url:             "/api/v1/skills?include=builtin&workspaceId=wsp_1",
			wantPath:        "/api/v1/workspaces/wsp_1/skills",
			wantEscapedPath: "/api/v1/workspaces/wsp_1/skills",
			wantQuery:       "include=builtin&workspaceId=wsp_1",
		},
		{
			name:            "knowledge base escaped suffix and body",
			method:          http.MethodPost,
			url:             "/api/v1/kb/sources/source%2Fbranch?limit=10&workspaceId=wsp_1",
			body:            `{"name":"source"}`,
			wantPath:        "/api/v1/workspaces/wsp_1/kb/sources/source/branch",
			wantEscapedPath: "/api/v1/workspaces/wsp_1/kb/sources/source%2Fbranch",
			wantQuery:       "limit=10&workspaceId=wsp_1",
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			request := httptest.NewRequest(test.method, test.url, strings.NewReader(test.body))
			request.Header.Set(api.HeaderRouteWorkspaceID, "wsp_1")
			request.Header.Set("Authorization", "Bearer desktop-session")
			recorder := httptest.NewRecorder()
			handler.ServeHTTP(recorder, request)
			require.Equal(t, http.StatusNoContent, recorder.Code, recorder.Body.String())

			seen := <-observed
			require.Equal(t, test.method, seen.method)
			require.Equal(t, test.wantPath, seen.path)
			require.Equal(t, test.wantEscapedPath, seen.escapedPath)
			require.Equal(t, test.wantQuery, seen.query)
			require.Equal(t, test.body, seen.body)
			require.Equal(t, "wsp_1", seen.routeID)
			require.Equal(t, "wsp_1", seen.workspaceID)
			require.Equal(t, "Bearer desktop-session", seen.authorization)
		})
	}
}

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

	for _, path := range []string{"/api/v1/skills", "/api/v1/skills/agents/bindings", "/api/v1/kb", "/api/llm/healthz"} {
		require.True(t, isLegacyDesktopWorkspaceRoute(path), path)
	}
	for _, path := range []string{"/api/v1/skills-extra", "/api/v1/kb-backup", "/api/llm-admin"} {
		require.False(t, isLegacyDesktopWorkspaceRoute(path), path)
	}
}

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

	var upstreamCalls atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		upstreamCalls.Add(1)
	}))
	t.Cleanup(upstream.Close)
	handler := newTestHandler(t, upstream.URL, map[string]string{"workspace-default": upstream.URL}, fakeResolver{
		target: workspace.WorkspaceRouteTarget{WorkspaceID: "wsp_1", WorkspaceGroupID: "workspace-default", PlacementState: workspace.PlacementStateActive},
	})

	request := httptest.NewRequest(http.MethodGet, "/api/v1/skills", nil)
	request.Header.Set(api.HeaderRouteWorkspaceID, "wsp/1")
	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusBadRequest, recorder.Code)
	require.Contains(t, recorder.Body.String(), "invalid route workspace identity")
	require.Zero(t, upstreamCalls.Load())
}

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

	var upstreamCalls atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		upstreamCalls.Add(1)
	}))
	t.Cleanup(upstream.Close)
	handler := newTestHandler(t, upstream.URL, map[string]string{"workspace-default": upstream.URL}, fakeResolver{
		target: workspace.WorkspaceRouteTarget{WorkspaceID: "wsp_1", WorkspaceGroupID: "workspace-default", PlacementState: workspace.PlacementStateActive},
	})

	for _, path := range []string{
		"/api/v1/kb?workspaceId=wsp_1",
		"/api/llm/healthz?workspace_id=wsp_1",
	} {
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, path, nil))
		require.Equal(t, http.StatusBadRequest, recorder.Code, path)
		require.Contains(t, recorder.Body.String(), api.HeaderRouteWorkspaceID+" is required", path)
	}
	require.Zero(t, upstreamCalls.Load())
}

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

	var upstreamCalls atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		upstreamCalls.Add(1)
	}))
	t.Cleanup(upstream.Close)
	handler := newTestHandler(t, upstream.URL, map[string]string{"workspace-default": upstream.URL}, fakeResolver{
		target: workspace.WorkspaceRouteTarget{WorkspaceID: "wsp_1", WorkspaceGroupID: "workspace-default", PlacementState: workspace.PlacementStateActive},
	})

	request := httptest.NewRequest(http.MethodGet, "/api/v1/kb?workspace_id=wsp_2", nil)
	request.Header.Set(api.HeaderRouteWorkspaceID, "wsp_1")
	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusBadRequest, recorder.Code)
	require.Contains(t, recorder.Body.String(), "conflicting route workspace identity")
	require.Zero(t, upstreamCalls.Load())
}

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

	workspaceUpstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		require.Equal(t, "wsp_1", request.Header.Get(api.HeaderRouteWorkspaceID))
		require.Equal(t, "wsp_1", request.Header.Get(api.HeaderWorkspaceID))
		require.Equal(t, "/api/llm/healthz", request.URL.Path)
		writer.WriteHeader(http.StatusNoContent)
	}))
	t.Cleanup(workspaceUpstream.Close)

	handler, err := NewHandler(fakeResolver{
		target: workspace.WorkspaceRouteTarget{WorkspaceID: "wsp_1", WorkspaceGroupID: "workspace-default", PlacementState: workspace.PlacementStateActive},
	}, Config{
		CentralUpstream:    "http://127.0.0.1:1",
		WorkspaceUpstreams: map[string]string{"workspace-default": workspaceUpstream.URL},
		IdentityResolver:   fakeWorkspaceIdentityResolver{workspaceID: "wsp_1"},
	})
	require.NoError(t, err)

	request := httptest.NewRequest(http.MethodGet, "/api/llm/healthz", nil)
	request.Header.Set("Authorization", "Bearer agentteam-session")
	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusBadRequest, recorder.Code)
	require.Contains(t, recorder.Body.String(), api.HeaderRouteWorkspaceID+" is required")
}

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

	paths := []string{"/api/llm/healthz"}
	for _, path := range paths {
		path := path
		t.Run(path, func(t *testing.T) {
			t.Parallel()
			workspaceUpstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
				t.Fatal("business-header-only request reached workspace upstream")
			}))
			t.Cleanup(workspaceUpstream.Close)

			handler := newTestHandler(t, "http://127.0.0.1:1", map[string]string{"workspace-default": workspaceUpstream.URL}, fakeResolver{target: workspace.WorkspaceRouteTarget{WorkspaceID: "wsp_1", WorkspaceGroupID: "workspace-default", PlacementState: workspace.PlacementStateActive}})

			recorder := httptest.NewRecorder()
			request := httptest.NewRequest(http.MethodPost, path, nil)
			request.Header.Set(api.HeaderWorkspaceID, "wsp_1")
			handler.ServeHTTP(recorder, request)

			require.Equal(t, http.StatusBadRequest, recorder.Code, path)
			require.Contains(t, recorder.Body.String(), api.HeaderRouteWorkspaceID+" is required")
		})
	}
}

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

	workspaceUpstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("business-header-only request reached workspace upstream")
	}))
	t.Cleanup(workspaceUpstream.Close)

	handler := newTestHandler(t, "http://127.0.0.1:1", map[string]string{"workspace-default": workspaceUpstream.URL}, fakeResolver{target: workspace.WorkspaceRouteTarget{WorkspaceID: "wsp_1", WorkspaceGroupID: "workspace-default", PlacementState: workspace.PlacementStateActive}})

	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodGet, "/api/llm/healthz", nil)
	request.Header.Set(api.HeaderWorkspaceID, "wsp_1")
	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusBadRequest, recorder.Code)
	require.Contains(t, recorder.Body.String(), api.HeaderRouteWorkspaceID+" is required")
}

func TestHandlerRejectsLegacyAgentRuntimeRoute(t *testing.T) {
	t.Parallel()
	handler, err := NewHandler(fakeResolver{}, Config{
		CentralUpstream: "http://127.0.0.1:1", WorkspaceUpstreams: map[string]string{"group-a": "http://127.0.0.1:2"},
		IdentityResolver: fakeWorkspaceIdentityResolver{workspaceID: "wsp_a"},
	})
	require.NoError(t, err)
	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodGet, "/internal/agent/agt_1/send", nil)
	request.Header.Set(api.HeaderWorkspaceID, "wsp_b")
	handler.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusNotFound, recorder.Code)
	require.Contains(t, recorder.Body.String(), "unknown gateway route")
}

func TestHandlerFailsClosedWhenLegacyWorkspaceHeaderIsMissing(t *testing.T) {
	t.Parallel()
	handler, err := NewHandler(fakeResolver{}, Config{
		CentralUpstream: "http://127.0.0.1:1", WorkspaceUpstreams: map[string]string{"group-a": "http://127.0.0.1:2"},
		IdentityResolver: fakeWorkspaceIdentityResolver{err: ErrWorkspaceIdentityUnresolved},
	})
	require.NoError(t, err)
	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/api/v1/skills", nil))
	require.Equal(t, http.StatusBadRequest, recorder.Code)
	require.Contains(t, recorder.Body.String(), api.HeaderRouteWorkspaceID+" is required")
}

func TestHandlerDoesNotResolveRejectedLegacyAgentRuntime(t *testing.T) {
	t.Parallel()
	handler, err := NewHandler(fakeResolver{}, Config{
		CentralUpstream: "http://127.0.0.1:1", WorkspaceUpstreams: map[string]string{"group-a": "http://127.0.0.1:2"},
		IdentityResolver: fakeWorkspaceIdentityResolver{err: ErrWorkspaceIdentityUnauthorized},
	})
	require.NoError(t, err)
	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/internal/agent/agt_1/send", nil))
	require.Equal(t, http.StatusNotFound, recorder.Code)
	require.Contains(t, recorder.Body.String(), "unknown gateway route")
}