package forwarder

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

	"github.com/stretchr/testify/require"

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

func TestForwardJSONReservedHeadersOverrideCustomHeaders(t *testing.T) {
	var seenAuthorization string
	var seenTrustedJWT string
	var seenWorkspaceID string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		seenAuthorization = request.Header.Get("Authorization")
		seenTrustedJWT = request.Header.Get(trustedJWTHeader)
		seenWorkspaceID = request.Header.Get(api.HeaderWorkspaceID)
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:        "/test",
		WorkspaceID: "wsp_typed",
		Identity:    TrustedIdentity{UserID: "usr_1"},
		Headers: map[string]string{
			"authorization":            "Bearer custom",
			"x-agentteam-verified-jwt": "Bearer custom",
			"x-workspace-id":           "wsp_custom",
		},
	}, nil)
	require.NoError(t, err)
	require.NotEqual(t, "Bearer custom", seenAuthorization)
	require.Equal(t, seenAuthorization, seenTrustedJWT)
	require.Equal(t, "wsp_typed", seenWorkspaceID)
}

func TestForwardJSONDoesNotForwardCustomWorkspaceWithoutTypedWorkspace(t *testing.T) {
	var seenWorkspaceID string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		seenWorkspaceID = request.Header.Get(api.HeaderWorkspaceID)
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:     "/test",
		Identity: TrustedIdentity{UserID: "usr_1"},
		Headers:  map[string]string{"x-workspace-id": "wsp_custom"},
	}, nil)
	require.NoError(t, err)
	require.Empty(t, seenWorkspaceID)
}

func TestForwardJSONUsesCompatibleUpstreamErrorFields(t *testing.T) {
	testCases := []struct {
		name        string
		body        string
		wantCode    string
		wantMessage string
	}{
		{
			name:        "code and message",
			body:        `{"code":"role_conflict","message":"role already exists"}`,
			wantCode:    "role_conflict",
			wantMessage: "role already exists",
		},
		{
			name:        "MeshX error and message",
			body:        `{"error":"role_conflict","message":"role already exists"}`,
			wantCode:    "role_conflict",
			wantMessage: "role already exists",
		},
		{
			name:        "error only",
			body:        `{"error":"role_conflict"}`,
			wantCode:    "role_conflict",
			wantMessage: "role_conflict",
		},
		{
			name:        "code fallback message",
			body:        `{"code":"role_conflict"}`,
			wantCode:    "role_conflict",
			wantMessage: "role_conflict",
		},
		{
			name:        "code takes precedence over error",
			body:        `{"code":"role_conflict","error":"legacy_conflict"}`,
			wantCode:    "role_conflict",
			wantMessage: "legacy_conflict",
		},
		{
			name:        "status text fallback",
			body:        `{}`,
			wantCode:    "",
			wantMessage: "Conflict",
		},
	}

	for _, testCase := range testCases {
		t.Run(testCase.name, func(t *testing.T) {
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				w.WriteHeader(http.StatusConflict)
				_, _ = w.Write([]byte(testCase.body))
			}))
			defer upstream.Close()

			handler := newForwardJSONTestHandler(t, upstream)
			err := handler.ForwardJSON(context.Background(), JSONRequest{
				Path:     "/test",
				Identity: TrustedIdentity{UserID: "usr_1"},
			}, nil)
			require.Error(t, err)
			var upstreamErr UpstreamHTTPError
			require.True(t, errors.As(err, &upstreamErr))
			require.Equal(t, testCase.wantCode, upstreamErr.Code)
			require.Equal(t, testCase.wantMessage, upstreamErr.Message)
		})
	}
}

func TestForwardJSONRequiresExactlyOneJSONValueForSuccessfulResponses(t *testing.T) {
	testCases := []struct {
		name string
		body string
	}{
		{name: "empty", body: ""},
		{name: "whitespace only", body: " \n\t"},
		{name: "non JSON", body: "not-json"},
		{name: "trailing JSON value", body: `{"ok":true} {"extra":true}`},
		{name: "trailing malformed data", body: `{"ok":true} trailing`},
	}

	for _, testCase := range testCases {
		for _, responseTarget := range []string{"decode", "discard"} {
			t.Run(testCase.name+"/"+responseTarget, func(t *testing.T) {
				upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
					w.WriteHeader(http.StatusOK)
					_, _ = w.Write([]byte(testCase.body))
				}))
				defer upstream.Close()

				handler := newForwardJSONTestHandler(t, upstream)
				var response any
				var target any = &response
				if responseTarget == "discard" {
					target = nil
				}
				err := handler.ForwardJSON(context.Background(), JSONRequest{
					Path:     "/test",
					Identity: TrustedIdentity{UserID: "usr_1"},
				}, target)
				require.Error(t, err)
				require.Contains(t, err.Error(), "decode upstream JSON response")
			})
		}
	}
}

func TestForwardJSONValidatesSuccessfulJSONWhenResponseIsNil(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusCreated)
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:     "/test",
		Identity: TrustedIdentity{UserID: "usr_1"},
	}, nil)
	require.NoError(t, err)
}

func TestForwardJSONDecodesExactlyOneJSONValueAtExpectedStatus(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusCreated)
		_, _ = w.Write([]byte(`{"id":"created"}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	var response struct {
		ID string `json:"id"`
	}
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:           "/test",
		Identity:       TrustedIdentity{UserID: "usr_1"},
		ExpectedStatus: http.StatusCreated,
	}, &response)
	require.NoError(t, err)
	require.Equal(t, "created", response.ID)
}

func TestForwardJSONRejectsUnexpectedNoContent(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:     "/test",
		Identity: TrustedIdentity{UserID: "usr_1"},
	}, nil)
	require.Error(t, err)
	require.Contains(t, err.Error(), "204")
}

func TestForwardJSONAcceptsExplicitlyExpectedNoContent(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:           "/test",
		Identity:       TrustedIdentity{UserID: "usr_1"},
		ExpectedStatus: http.StatusNoContent,
	}, nil)
	require.NoError(t, err)
}

func TestForwardJSONRejectsResponseTargetForExpectedNoContent(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	var response map[string]any
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:           "/test",
		Identity:       TrustedIdentity{UserID: "usr_1"},
		ExpectedStatus: http.StatusNoContent,
	}, &response)
	require.Error(t, err)
	require.Contains(t, err.Error(), "response")
}

func TestForwardJSONEnforcesExplicitExpectedStatus(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:           "/test",
		Identity:       TrustedIdentity{UserID: "usr_1"},
		ExpectedStatus: http.StatusCreated,
	}, nil)
	require.Error(t, err)
	require.Contains(t, err.Error(), "expected 201")
	require.Contains(t, err.Error(), "received 200")
}

func TestForwardJSONKeepsInternalPurposeOutOfHeaders(t *testing.T) {
	var seenPurposeHeader string
	var seenAuthorization string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		seenPurposeHeader = request.Header.Get("X-AgentTeam-Internal-Purpose")
		seenAuthorization = request.Header.Get("Authorization")
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:            "/test",
		Identity:        TrustedIdentity{UserID: "usr_1"},
		InternalPurpose: "approved_operation",
	}, nil)
	require.NoError(t, err)
	require.Empty(t, seenPurposeHeader)
	claims := decodeClaims(t, strings.TrimPrefix(seenAuthorization, "Bearer "))
	require.Equal(t, "approved_operation", claims["agentteam_internal_purpose"])
}

func TestForwardJSONProjectsExplicitOwnershipWorkspaceWithoutOwnedAgents(t *testing.T) {
	var seenClaims map[string]any
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		seenClaims = decodeClaims(t, strings.TrimPrefix(request.Header.Get("Authorization"), "Bearer "))
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:                      "/test",
		WorkspaceID:               "wsp_1",
		Identity:                  TrustedIdentity{UserID: "usr_1"},
		AgentOwnershipWorkspaceID: "wsp_1",
	}, nil)
	require.NoError(t, err)
	require.Equal(t, "wsp_1", seenClaims["agent_ownership_workspace_id"])
	require.NotContains(t, seenClaims, "owned_agent_ids")
	require.NotContains(t, seenClaims, "owned_agents")
}

func TestForwardJSONStripsCustomInternalPurposeHeader(t *testing.T) {
	var seenPurposeHeader string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		seenPurposeHeader = request.Header.Get("X-AgentTeam-Internal-Purpose")
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	handler := newForwardJSONTestHandler(t, upstream)
	err := handler.ForwardJSON(context.Background(), JSONRequest{
		Path:     "/test",
		Identity: TrustedIdentity{UserID: "usr_1"},
		Headers: map[string]string{
			"x-AgEnTtEaM-InTeRnAl-PuRpOsE": "attacker-controlled",
		},
	}, nil)
	require.NoError(t, err)
	require.Empty(t, seenPurposeHeader)
}

func newForwardJSONTestHandler(t *testing.T, upstream *httptest.Server) *Handler {
	t.Helper()
	handler, err := NewHandler(Dependencies{
		BaseURL:            upstream.URL,
		HTTPClient:         upstream.Client(),
		TrustedJWTAudience: "meshx",
	})
	require.NoError(t, err)
	return handler
}