package authn

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/go-jose/go-jose/v4"
	"github.com/go-jose/go-jose/v4/jwt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNewTokenExchangeClient(t *testing.T) {
	t.Run("should not be able to create client wihtout token", func(t *testing.T) {
		_, err := NewTokenExchangeClient(TokenExchangeConfig{
			TokenExchangeURL: "some-url",
		})
		require.Error(t, err)

	})

	t.Run("should not be able to create client wihtout url", func(t *testing.T) {
		_, err := NewTokenExchangeClient(TokenExchangeConfig{
			Token: "some-token",
		})
		require.Error(t, err)
	})

	t.Run("should be able to create client", func(t *testing.T) {
		_, err := NewTokenExchangeClient(TokenExchangeConfig{
			Token:            "some-token",
			TokenExchangeURL: "some-url",
		})
		require.NoError(t, err)
	})
}

func Test_TokenExchangeClient_Exchange(t *testing.T) {
	expiresIn := 10 * time.Minute
	setup := func(srv *httptest.Server, opts ...ExchangeClientOpts) *TokenExchangeClient {
		c, err := NewTokenExchangeClient(TokenExchangeConfig{
			Token:            "some-token",
			TokenExchangeURL: srv.URL,
		}, opts...)
		require.NoError(t, err)
		return c
	}

	t.Run("should return error if namespace is empty", func(t *testing.T) {
		c := setup(httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})))
		res, err := c.Exchange(context.Background(), TokenExchangeRequest{})
		assert.ErrorIs(t, err, ErrMissingNamespace)
		assert.Nil(t, res)

	})

	t.Run("should return error if audiences is empty", func(t *testing.T) {
		c := setup(httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {})))
		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*"})
		assert.ErrorIs(t, err, ErrMissingAudiences)
		assert.Nil(t, res)
	})

	t.Run("should return error for unexpected server response", func(t *testing.T) {
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusUnauthorized)
			_, _ = w.Write([]byte("{}"))
		})))
		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}})
		assert.ErrorIs(t, err, ErrInvalidExchangeResponse)
		assert.Nil(t, res)
	})

	t.Run("should cache and return token on successful sign request", func(t *testing.T) {
		var calls int
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls++
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
			bytes.NewBuffer([]byte(`{}`))
			json.NewEncoder(&bytes.Buffer{})
		})))

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)

		// same namespace and audiences should load token from cache
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)

		// different namespace should issue new token request
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "stacks-1", Audiences: []string{"some-service"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 2, calls)

		// different audiences should issue new token request
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service-2"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 3, calls)
	})

	t.Run("should cache and return token on successful exchange request", func(t *testing.T) {
		var calls int
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls++
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
			bytes.NewBuffer([]byte(`{}`))
			json.NewEncoder(&bytes.Buffer{})
		})))

		tokenToBeExchanged := signAccessToken(t, expiresIn)

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}, SubjectToken: tokenToBeExchanged})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)

		// same namespace and audiences should load token from cache
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}, SubjectToken: tokenToBeExchanged})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)

		// different namespace should issue new token request
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "stacks-1", Audiences: []string{"some-service"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 2, calls)

		// different audiences should issue new token request
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service-2"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 3, calls)

		// different subjectToken should issue new token request
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service-2"}, SubjectToken: signAccessToken(t, expiresIn)})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 4, calls)
	})

	t.Run("should include expiration in request when provided", func(t *testing.T) {
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			requestBody, err := io.ReadAll(r.Body)
			require.NoError(t, err)
			w.WriteHeader(http.StatusOK)
			var request TokenExchangeRequest
			err = json.Unmarshal(requestBody, &request)
			require.NoError(t, err)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, time.Duration(*request.ExpiresIn)*time.Second) + `"}}`))
		})))

		expiresIn := 30
		res, err := c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace: "*",
			Audiences: []string{"some-service"},
			ExpiresIn: &expiresIn,
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)

		resultToken, err := jwt.ParseSigned(res.Token, []jose.SignatureAlgorithm{jose.ES256})
		require.NoError(t, err)
		var claims jwt.Claims
		err = resultToken.UnsafeClaimsWithoutVerification(&claims)
		require.NoError(t, err)
		expectedExpiry := time.Now().Add(time.Duration(expiresIn) * time.Second)
		require.InDelta(t, expectedExpiry.Unix(), claims.Expiry.Time().Unix(), 1)
	})

	t.Run("should include subject in request when provided", func(t *testing.T) {
		var got TokenExchangeRequest
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			requestBody, err := io.ReadAll(r.Body)
			require.NoError(t, err)
			require.NoError(t, json.Unmarshal(requestBody, &got))
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
		})))

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace: "stacks-1",
			Audiences: []string{"some-service"},
			Subject: &TokenExchangeSubject{
				Identifier: "1",
				Type:       "user",
				Namespace:  "stacks-1",
				Email:      "user@example.com",
				Username:   "user",
			},
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)

		require.Empty(t, got.SubjectToken)
		require.NotNil(t, got.Subject)
		require.Equal(t, "1", got.Subject.Identifier)
		require.Equal(t, "user", got.Subject.Type)
		require.Equal(t, "stacks-1", got.Subject.Namespace)
		require.Equal(t, "user@example.com", got.Subject.Email)
		require.Equal(t, "user", got.Subject.Username)
	})

	t.Run("should return error when subject and subjectToken are both set", func(t *testing.T) {
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
		})))

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace:    "stacks-1",
			Audiences:    []string{"some-service"},
			SubjectToken: signAccessToken(t, expiresIn),
			Subject: &TokenExchangeSubject{
				Identifier: "1",
				Type:       "user",
				Namespace:  "stacks-1",
			},
		})
		require.ErrorIs(t, err, ErrMutuallyExclusiveSubject)
		require.Nil(t, res)
	})

	t.Run("should not reuse cache for different subjects", func(t *testing.T) {
		var calls int
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			calls++
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
		})))

		subject1 := &TokenExchangeSubject{Identifier: "1", Type: "user", Namespace: "stacks-1"}
		subject2 := &TokenExchangeSubject{Identifier: "2", Type: "user", Namespace: "stacks-1"}

		_, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "stacks-1", Audiences: []string{"some-service"}, Subject: subject1})
		assert.NoError(t, err)
		require.Equal(t, 1, calls)

		// Same subject should reuse the cache.
		_, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "stacks-1", Audiences: []string{"some-service"}, Subject: subject1})
		assert.NoError(t, err)
		require.Equal(t, 1, calls)

		// Different subject should miss the cache.
		_, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "stacks-1", Audiences: []string{"some-service"}, Subject: subject2})
		assert.NoError(t, err)
		require.Equal(t, 2, calls)

		// Subjects differing only by Sub should miss the cache.
		subject3 := &TokenExchangeSubject{Sub: "user:1", Identifier: "abc", Type: "user", Namespace: "stacks-1"}
		subject4 := &TokenExchangeSubject{Sub: "user:2", Identifier: "abc", Type: "user", Namespace: "stacks-1"}

		_, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "stacks-1", Audiences: []string{"some-service"}, Subject: subject3})
		assert.NoError(t, err)
		require.Equal(t, 3, calls)

		_, err = c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "stacks-1", Audiences: []string{"some-service"}, Subject: subject4})
		assert.NoError(t, err)
		require.Equal(t, 4, calls)
	})

	t.Run("should use an alternate cache if provided", func(t *testing.T) {
		testcache := &testCache{data: make(map[string][]byte)}

		var calls int
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls++
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
			bytes.NewBuffer([]byte(`{}`))
			json.NewEncoder(&bytes.Buffer{})
		})), WithTokenExchangeClientCache(testcache))

		tokenToBeExchanged := signAccessToken(t, expiresIn)

		res1, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}, SubjectToken: tokenToBeExchanged})
		assert.NoError(t, err)
		assert.NotNil(t, res1)
		require.Equal(t, 1, calls)
		require.Len(t, testcache.data, 1)

		// same namespace and audiences should load token from cache
		res2, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}, SubjectToken: tokenToBeExchanged})
		assert.NoError(t, err)
		assert.NotNil(t, res2)
		require.Equal(t, 1, calls)
		require.Len(t, testcache.data, 1)
		require.Equal(t, res1, res2)

		// This is only testing that the cache is used, so we do not repeat the other cases here.
	})

	t.Run("should use stable subject plus actor chain for cache key", func(t *testing.T) {
		var calls int
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls++
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
		})))

		chain := &ActorClaims{
			Subject: "access-policy:mycap",
			Actor: &ActorClaims{
				Subject: "service:gateway",
				Actor: &ActorClaims{
					Subject: "user:1",
				},
			},
		}

		token1 := signSubjectToken(t, 10*time.Minute, "access-policy:mycap", chain)
		token2 := signSubjectToken(t, 11*time.Minute, "access-policy:mycap", chain)

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace:    "*",
			Audiences:    []string{"some-service"},
			SubjectToken: token1,
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)

		// Same stable subject + actor lineage should reuse cache, even if token body changed.
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace:    "*",
			Audiences:    []string{"some-service"},
			SubjectToken: token2,
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)
	})

	t.Run("should use stable subject for first hop without actor chain", func(t *testing.T) {
		var calls int
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls++
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
		})))

		token1 := signSubjectToken(t, 10*time.Minute, "access-policy:mycap", nil)
		token2 := signSubjectToken(t, 11*time.Minute, "access-policy:mycap", nil)

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace:    "*",
			Audiences:    []string{"some-service"},
			SubjectToken: token1,
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)

		// No actor chain yet; stable top-level subject should still reuse cache.
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace:    "*",
			Audiences:    []string{"some-service"},
			SubjectToken: token2,
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)
	})

	t.Run("should not reuse cache for different actor chain", func(t *testing.T) {
		var calls int
		c := setup(httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls++
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
		})))

		tokenUser1 := signSubjectToken(t, 10*time.Minute, "access-policy:mycap", &ActorClaims{
			Subject: "access-policy:mycap",
			Actor: &ActorClaims{
				Subject: "service:gateway",
				Actor: &ActorClaims{
					Subject: "user:1",
				},
			},
		})
		tokenUser2 := signSubjectToken(t, 10*time.Minute, "access-policy:mycap", &ActorClaims{
			Subject: "access-policy:mycap",
			Actor: &ActorClaims{
				Subject: "service:gateway",
				Actor: &ActorClaims{
					Subject: "user:2",
				},
			},
		})

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace:    "*",
			Audiences:    []string{"some-service"},
			SubjectToken: tokenUser1,
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 1, calls)

		// Different actor lineage should create a different cache key.
		res, err = c.Exchange(context.Background(), TokenExchangeRequest{
			Namespace:    "*",
			Audiences:    []string{"some-service"},
			SubjectToken: tokenUser2,
		})
		assert.NoError(t, err)
		assert.NotNil(t, res)
		require.Equal(t, 2, calls)
	})

	t.Run("should retry requests and return token on successful sign request", func(t *testing.T) {
		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
			bytes.NewBuffer([]byte(`{}`))
			json.NewEncoder(&bytes.Buffer{})
		}))

		c, err := NewTokenExchangeClient(TokenExchangeConfig{
			Token:            "some-token",
			TokenExchangeURL: srv.URL,
		})
		require.NoError(t, err)

		var countRequests = 0
		// brokenTransport returns a network error on the first 2 requests.
		brokenTransport := &http.Transport{
			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
				if countRequests < 2 {
					countRequests++
					return nil, fmt.Errorf("network error")
				}
				countRequests = 0
				return net.Dial(network, addr)
			},
		}
		c.client.Transport = brokenTransport

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
	})

	t.Run("should retry requests and return error after 3 failed requests", func(t *testing.T) {
		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
			bytes.NewBuffer([]byte(`{}`))
			json.NewEncoder(&bytes.Buffer{})
		}))

		c, err := NewTokenExchangeClient(TokenExchangeConfig{
			Token:            "some-token",
			TokenExchangeURL: srv.URL,
		})
		require.NoError(t, err)

		var countRequests = 0
		// brokenTransport returns a network error on the first 3 requests.
		brokenTransport := &http.Transport{
			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
				if countRequests < 3 {
					countRequests++
					return nil, fmt.Errorf("network error")
				}
				countRequests = 0
				return net.Dial(network, addr)
			},
		}
		c.client.Transport = brokenTransport

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}})
		assert.Error(t, err)
		assert.Nil(t, res)
	})

	t.Run("should retry failed requests with transient errors, like status code 500x, and return token on successful sign request", func(t *testing.T) {
		var countRequests = 0

		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			countRequests++
			if countRequests < 3 {
				w.WriteHeader(http.StatusInternalServerError)
				_, _ = w.Write([]byte("{}"))
				return
			}

			require.Equal(t, r.Header.Get("Authorization"), "Bearer some-token")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(`{"data": {"token": "` + signAccessToken(t, expiresIn) + `"}}`))
			bytes.NewBuffer([]byte(`{}`))
			json.NewEncoder(&bytes.Buffer{})
		}))

		c, err := NewTokenExchangeClient(TokenExchangeConfig{
			Token:            "some-token",
			TokenExchangeURL: srv.URL,
		})
		require.NoError(t, err)

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}})
		assert.NoError(t, err)
		assert.NotNil(t, res)
	})

	t.Run("should not retry failed requests with permanent errors, like status code 400x", func(t *testing.T) {
		var countRequests = 0

		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			countRequests++
			w.WriteHeader(http.StatusUnauthorized)
			_, _ = w.Write([]byte("{}"))
		}))

		c, err := NewTokenExchangeClient(TokenExchangeConfig{
			Token:            "some-token",
			TokenExchangeURL: srv.URL,
		})
		require.NoError(t, err)

		res, err := c.Exchange(context.Background(), TokenExchangeRequest{Namespace: "*", Audiences: []string{"some-service"}})
		assert.ErrorIs(t, err, ErrInvalidExchangeResponse)
		assert.Nil(t, res)
		assert.Equal(t, countRequests, 1)
	})
}

func signAccessToken(t *testing.T, expiresIn time.Duration) string {
	signer, err := jose.NewSigner(jose.SigningKey{
		Algorithm: jose.ES256,
		Key:       testPrivateKey,
	}, nil)
	require.NoError(t, err)

	token, err := jwt.Signed(signer).
		Claims(&jwt.Claims{Expiry: jwt.NewNumericDate(time.Now().Add(expiresIn))}).
		Serialize()

	require.NoError(t, err)
	return token
}

func signSubjectToken(t *testing.T, expiresIn time.Duration, subject string, actor *ActorClaims) string {
	signer, err := jose.NewSigner(jose.SigningKey{
		Algorithm: jose.ES256,
		Key:       testPrivateKey,
	}, nil)
	require.NoError(t, err)

	customClaims := AccessTokenClaims{
		Actor: actor,
	}

	token, err := jwt.Signed(signer).
		Claims(&jwt.Claims{
			Subject: subject,
			Expiry:  jwt.NewNumericDate(time.Now().Add(expiresIn)),
		}).
		Claims(&customClaims).
		Serialize()
	require.NoError(t, err)

	return token
}
