package trusted

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestDecodePayloadExtractsAgentTeamIdentityClaims(t *testing.T) {
	token := jwt(map[string]any{
		"sub":                "user-1",
		"email":              "user@example.com",
		"name":               "User One",
		"preferred_username": "uone",
		"groups":             []string{"dev", "ops"},
		"tenant":             "tenant-1",
	})

	identity, err := DecodePayload(token)
	require.NoError(t, err)
	require.Equal(t, "user-1", identity.Subject)
	require.Equal(t, "user@example.com", identity.Email)
	require.Equal(t, "User One", identity.Name)
	require.Equal(t, "uone", identity.PreferredUsername)
	require.Equal(t, []string{"dev", "ops"}, identity.Groups)
	require.Equal(t, "tenant-1", identity.Tenant)
}

func TestDecodePayloadRejectsMalformedJWT(t *testing.T) {
	_, err := DecodePayload("not-a-jwt")
	require.Error(t, err)
}

func jwt(claims map[string]any) string {
	header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"none"}`))
	payload, _ := json.Marshal(claims)
	return header + "." + base64.RawURLEncoding.EncodeToString(payload) + "."
}

func TestIdentityFromRequestReadsAuthorizationBearerByDefault(t *testing.T) {
	token := jwt(map[string]any{"sub": "user-1"})
	req := httptest.NewRequest("GET", "/", nil)
	req.Header.Set("Authorization", "Bearer "+token)

	identity, gotToken, err := IdentityFromRequest(req, Options{})
	require.NoError(t, err)
	require.Equal(t, "user-1", identity.Subject)
	require.Equal(t, token, gotToken)
}

func TestIdentityFromRequestRejectsNonBearerAuthorization(t *testing.T) {
	req := httptest.NewRequest("GET", "/", nil)
	req.Header.Set("Authorization", "Basic abc")

	_, _, err := IdentityFromRequest(req, Options{})
	require.Error(t, err)
	require.Contains(t, err.Error(), "missing AgentTeam-verified JWT")
}

func TestIdentityContextRoundTrip(t *testing.T) {
	ctx := WithIdentity(context.Background(), Identity{Subject: "user-1", Groups: []string{"dev"}})

	identity, ok := IdentityFromContext(ctx)
	require.True(t, ok)
	require.Equal(t, "user-1", identity.Subject)
	require.Equal(t, []string{"dev"}, identity.Groups)

	_, ok = IdentityFromContext(context.Background())
	require.False(t, ok)
}

func TestIdentityFromRequestReadsRawTokenFromTrustedHeader(t *testing.T) {
	token := jwt(map[string]any{"sub": "trusted-user"})
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("X-AgentTeam-Verified-JWT", token)

	identity, gotToken, err := IdentityFromRequest(req, Options{Header: "X-AgentTeam-Verified-JWT"})
	require.NoError(t, err)
	require.Equal(t, "trusted-user", identity.Subject)
	require.Equal(t, token, gotToken)
}

func TestIdentityFromRequestAllowsAuthorizationMockOnlyWhenConfigured(t *testing.T) {
	token := jwt(map[string]any{"sub": "mock-user"})
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("Authorization", "Bearer "+token)

	_, _, err := IdentityFromRequest(req, Options{Header: "X-AgentTeam-Verified-JWT"})
	require.Error(t, err)

	identity, gotToken, err := IdentityFromRequest(req, Options{Header: "X-AgentTeam-Verified-JWT", AllowMock: true})
	require.NoError(t, err)
	require.Equal(t, "mock-user", identity.Subject)
	require.Equal(t, token, gotToken)
}

func TestMiddlewareMapsAuthorizationAndAddsIdentityToContext(t *testing.T) {
	token := jwt(map[string]any{"sub": "user-1"})
	var mappedIdentity Identity
	var mappedToken string
	var handledIdentity Identity
	var authorization string
	handler := Middleware(Options{
		Enabled: true,
		Header:  "X-AgentTeam-Verified-JWT",
		TokenMapper: func(_ context.Context, identity Identity, token string) (string, error) {
			mappedIdentity = identity
			mappedToken = token
			return "mapped-token", nil
		},
	})(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		handledIdentity, _ = IdentityFromContext(r.Context())
		authorization = r.Header.Get("Authorization")
	}))
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("X-AgentTeam-Verified-JWT", token)
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)
	require.Equal(t, "user-1", mappedIdentity.Subject)
	require.Equal(t, token, mappedToken)
	require.Equal(t, "user-1", handledIdentity.Subject)
	require.Equal(t, "Bearer mapped-token", authorization)
}

func TestMiddlewareRejectsTokenMapperFailure(t *testing.T) {
	token := jwt(map[string]any{"sub": "user-1"})
	called := false
	handler := Middleware(Options{
		Enabled: true,
		TokenMapper: func(context.Context, Identity, string) (string, error) {
			return "", errors.New("mapping failed")
		},
	})(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true }))
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusUnauthorized, recorder.Code)
	require.Contains(t, recorder.Body.String(), "mapping failed")
	require.False(t, called)
}