package gateway

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/agsbridge/workspaceurl"
	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
	workspaceroute "github.com/tidbcloud/AgentTeam/internal/workspace/router"
)

func TestStandardReverseProxyStreamsWorkspaceAGSGitWire(t *testing.T) {
	requestChunkSeen := make(chan struct{})
	releaseResponse := make(chan struct{})
	observed := make(chan *http.Request, 1)
	workspace := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		path, err := workspaceurl.Parse(request)
		require.NoError(t, err)
		require.NoError(t, path.Strip(request, "wsp_1"))
		first := make([]byte, len("first-chunk"))
		_, err = io.ReadFull(request.Body, first)
		require.NoError(t, err)
		require.Equal(t, "first-chunk", string(first))
		close(requestChunkSeen)
		rest, err := io.ReadAll(request.Body)
		require.NoError(t, err)
		require.Len(t, rest, 1<<20)
		observed <- request.Clone(request.Context())

		writer.Header().Set("Content-Type", "application/x-git-upload-pack-result")
		writer.Header().Set("X-Git-Protocol", "v2")
		writer.WriteHeader(http.StatusCreated)
		_, err = writer.Write([]byte("response-first\n"))
		require.NoError(t, err)
		writer.(http.Flusher).Flush()
		<-releaseResponse
		_, err = writer.Write([]byte("response-last\n"))
		require.NoError(t, err)
	}))
	t.Cleanup(workspace.Close)
	workspaceURL, err := url.Parse(workspace.URL)
	require.NoError(t, err)
	gatewayServer := httptest.NewServer(newSanitizedReverseProxy(workspaceURL))
	t.Cleanup(gatewayServer.Close)

	bodyReader, bodyWriter := io.Pipe()
	request, err := http.NewRequest(http.MethodPost, gatewayServer.URL+"/api/v1/workspaces/wsp_1/ags/acme/repo.git/git-upload-pack/../opaque//part?service=git-upload-pack", bodyReader)
	require.NoError(t, err)
	request.Header.Set("Content-Type", "application/x-git-upload-pack-request")
	request.Header.Set("Authorization", "Basic opaque")
	responseReady := make(chan *http.Response, 1)
	requestErr := make(chan error, 1)
	go func() {
		response, err := gatewayServer.Client().Do(request)
		if err != nil {
			requestErr <- err
			return
		}
		responseReady <- response
	}()

	_, err = bodyWriter.Write([]byte("first-chunk"))
	require.NoError(t, err)
	select {
	case <-requestChunkSeen:
	case <-time.After(2 * time.Second):
		t.Fatal("workspace upstream did not receive the first Git body chunk before request completion")
	}
	_, err = io.Copy(bodyWriter, bytes.NewReader(make([]byte, 1<<20)))
	require.NoError(t, err)
	require.NoError(t, bodyWriter.Close())

	var response *http.Response
	select {
	case response = <-responseReady:
	case err := <-requestErr:
		t.Fatal(err)
	case <-time.After(2 * time.Second):
		t.Fatal("Git response headers were buffered")
	}
	t.Cleanup(func() { _ = response.Body.Close() })
	require.Equal(t, "application/x-git-upload-pack-result", response.Header.Get("Content-Type"))
	require.Equal(t, http.StatusCreated, response.StatusCode)
	require.Equal(t, "v2", response.Header.Get("X-Git-Protocol"))
	reader := bufio.NewReader(response.Body)
	firstResponse, err := reader.ReadString('\n')
	require.NoError(t, err)
	require.Equal(t, "response-first\n", firstResponse)
	close(releaseResponse)
	lastResponse, err := reader.ReadString('\n')
	require.NoError(t, err)
	require.Equal(t, "response-last\n", lastResponse)

	got := <-observed
	require.Equal(t, http.MethodPost, got.Method)
	require.Equal(t, "/acme/repo.git/git-upload-pack/../opaque//part", got.URL.Path)
	require.Equal(t, "service=git-upload-pack", got.URL.RawQuery)
	require.Equal(t, "application/x-git-upload-pack-request", got.Header.Get("Content-Type"))
	require.Equal(t, "Basic opaque", got.Header.Get("Authorization"))
	require.EqualValues(t, -1, got.ContentLength)
	require.Contains(t, got.TransferEncoding, "chunked")
	require.False(t, strings.HasPrefix(got.URL.Path, "/api/v1/workspaces/"))
}

func TestWorkspaceRouterPreservesWorkspaceAGSGitStreamingAndAuthority(t *testing.T) {
	var groupAHits, groupBHits, businessHits atomic.Int32
	var defaultHits atomic.Int32

	newWorkspace := func(workspaceID, groupID string, hits *atomic.Int32) *httptest.Server {
		return 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(workspaceroute.HeaderResolvedWorkspaceGroupID))
			path, err := workspaceurl.Parse(request)
			if err != nil || path.Strip(request, request.Header.Get(api.HeaderRouteWorkspaceID)) != nil {
				http.Error(writer, "workspace path does not match routed authority", http.StatusConflict)
				return
			}
			businessHits.Add(1)
			writer.WriteHeader(http.StatusNoContent)
		}))
	}
	groupA := newWorkspace("wsp_a", "group-a", &groupAHits)
	t.Cleanup(groupA.Close)
	groupB := newWorkspace("wsp_b", "group-b", &groupBHits)
	t.Cleanup(groupB.Close)
	defaultGroup := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		defaultHits.Add(1)
	}))
	t.Cleanup(defaultGroup.Close)
	router, err := workspaceroute.NewHandler(workspaceRouteResolver{targets: map[string]workspace.WorkspaceRouteTarget{
		"wsp_a": {WorkspaceID: "wsp_a", WorkspaceGroupID: "group-a"},
		"wsp_b": {WorkspaceID: "wsp_b", WorkspaceGroupID: "group-b"},
	}}, map[string]string{"group-a": groupA.URL, "group-b": groupB.URL, "workspace-default": defaultGroup.URL})
	require.NoError(t, err)

	verifier := &workspaceCandidateVerifier{memberships: map[string]bool{"wsp_a": true, "wsp_b": true}}
	authenticatedOwner := func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
			parsed, err := workspaceurl.Parse(request)
			if err != nil {
				http.NotFound(writer, request)
				return
			}
			verifiedWorkspace, ok := verifier.Verify(parsed.WorkspaceID)
			if !ok {
				http.Error(writer, "workspace candidate could not be verified", http.StatusUnauthorized)
				return
			}
			request.Header.Set(api.HeaderRouteWorkspaceID, verifiedWorkspace)
			next.ServeHTTP(writer, request)
		})
	}

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_a/ags/acme/repo.git/git-upload-pack?service=git-upload-pack", strings.NewReader("git-body"))
	request.Header.Set(api.HeaderRouteWorkspaceID, "wsp_forged")
	request.Header.Set(api.HeaderWorkspaceID, "wsp_forged")
	request.Header.Set(workspaceroute.HeaderResolvedWorkspaceGroupID, "group-forged")
	request.Header.Set("Authorization", "Basic opaque")
	request.Header.Set("Content-Type", "application/x-git-upload-pack-request")
	recorder := httptest.NewRecorder()
	authenticatedOwner(router).ServeHTTP(recorder, request)
	require.Equal(t, http.StatusNoContent, recorder.Code)
	require.EqualValues(t, 1, groupAHits.Load())
	require.Zero(t, groupBHits.Load())
	require.EqualValues(t, 1, businessHits.Load())
	require.EqualValues(t, 1, verifier.calls.Load())
	require.Zero(t, defaultHits.Load())

	// A second authoritative workspace selects exactly one different group.
	second := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_b/ags/acme/repo.git/info/refs?service=git-upload-pack", nil)
	secondRecorder := httptest.NewRecorder()
	authenticatedOwner(router).ServeHTTP(secondRecorder, second)
	require.Equal(t, http.StatusNoContent, secondRecorder.Code)
	require.EqualValues(t, 1, groupAHits.Load())
	require.EqualValues(t, 1, groupBHits.Load())
	require.EqualValues(t, 2, businessHits.Load())
	require.EqualValues(t, 2, verifier.calls.Load())
	require.Zero(t, defaultHits.Load())

	// An internal claim/placement anomaly is injected after authentication: path
	// A is routed with authoritative claim B. It reaches only Workspace B and
	// fails there before the AGS/business handler.
	mismatch := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_a/ags/acme/repo.git/info/refs?service=git-upload-pack", nil)
	mismatch.Header.Set(api.HeaderRouteWorkspaceID, "wsp_b")
	mismatchRecorder := httptest.NewRecorder()
	router.ServeHTTP(mismatchRecorder, mismatch)
	require.Equal(t, http.StatusConflict, mismatchRecorder.Code)
	require.EqualValues(t, 1, groupAHits.Load())
	require.EqualValues(t, 2, groupBHits.Load())
	require.EqualValues(t, 2, businessHits.Load())
	require.EqualValues(t, 2, verifier.calls.Load())
	require.Zero(t, defaultHits.Load())
}

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

func (verifier *workspaceCandidateVerifier) Verify(pathCandidate string) (string, bool) {
	verifier.calls.Add(1)
	return pathCandidate, verifier.memberships[pathCandidate]
}

type workspaceRouteResolver struct {
	targets map[string]workspace.WorkspaceRouteTarget
}

func (resolver workspaceRouteResolver) ResolveWorkspaceRouteTarget(_ context.Context, workspaceID string) (workspace.WorkspaceRouteTarget, error) {
	return resolver.targets[workspaceID], nil
}