// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	"github.com/stacklok/toolhive/pkg/authserver/storage"
	storagemocks "github.com/stacklok/toolhive/pkg/authserver/storage/mocks"
	"github.com/stacklok/toolhive/pkg/authserver/upstream"
	upstreammocks "github.com/stacklok/toolhive/pkg/authserver/upstream/mocks"
)

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

	newExpiry := time.Now().Add(1 * time.Hour)
	sessionBound := time.Now().Add(7 * 24 * time.Hour)

	baseExpired := &storage.UpstreamTokens{
		ProviderID:       "github",
		AccessToken:      "old-access",
		RefreshToken:     "old-refresh",
		IDToken:          "old-id-token",
		ExpiresAt:        time.Now().Add(-1 * time.Hour),
		SessionExpiresAt: sessionBound,
		UserID:           "user-123",
		UpstreamSubject:  "upstream-sub-456",
		ClientID:         "client-abc",
	}

	tests := []struct {
		name           string
		sessionID      string
		expired        *storage.UpstreamTokens
		setupProvider  func(*testing.T, *upstreammocks.MockOAuth2Provider)
		setupStorage   func(*testing.T, *storagemocks.MockUpstreamTokenStorage)
		wantErr        bool
		wantErrContain string
		checkResult    func(*testing.T, *storage.UpstreamTokens)
	}{
		{
			name:      "successful refresh with token rotation",
			sessionID: "session-1",
			expired:   baseExpired,
			setupProvider: func(_ *testing.T, p *upstreammocks.MockOAuth2Provider) {
				p.EXPECT().RefreshTokens(gomock.Any(), "old-refresh", "upstream-sub-456").
					Return(&upstream.Tokens{
						AccessToken:  "new-access",
						RefreshToken: "new-refresh",
						IDToken:      "new-id-token",
						ExpiresAt:    newExpiry,
					}, nil)
			},
			setupStorage: func(_ *testing.T, s *storagemocks.MockUpstreamTokenStorage) {
				s.EXPECT().StoreUpstreamTokens(gomock.Any(), "session-1", "github", gomock.Any()).
					DoAndReturn(func(_ context.Context, _, _ string, tokens *storage.UpstreamTokens) error {
						// Verify binding fields are preserved from expired tokens
						assert.Equal(t, "github", tokens.ProviderID)
						assert.Equal(t, "user-123", tokens.UserID)
						assert.Equal(t, "upstream-sub-456", tokens.UpstreamSubject)
						assert.Equal(t, "client-abc", tokens.ClientID)
						// Verify new token values
						assert.Equal(t, "new-access", tokens.AccessToken)
						assert.Equal(t, "new-refresh", tokens.RefreshToken)
						assert.Equal(t, "new-id-token", tokens.IDToken)
						assert.Equal(t, newExpiry, tokens.ExpiresAt)
						return nil
					})
			},
			checkResult: func(t *testing.T, result *storage.UpstreamTokens) {
				t.Helper()
				assert.Equal(t, "new-access", result.AccessToken)
				assert.Equal(t, "new-refresh", result.RefreshToken)
				assert.Equal(t, "new-id-token", result.IDToken)
				assert.Equal(t, newExpiry, result.ExpiresAt)
				// Binding fields preserved
				assert.Equal(t, "github", result.ProviderID)
				assert.Equal(t, "user-123", result.UserID)
				assert.Equal(t, "upstream-sub-456", result.UpstreamSubject)
				assert.Equal(t, "client-abc", result.ClientID)
			},
		},
		{
			name:      "provider does not rotate refresh token - keeps old one",
			sessionID: "session-2",
			expired:   baseExpired,
			setupProvider: func(_ *testing.T, p *upstreammocks.MockOAuth2Provider) {
				p.EXPECT().RefreshTokens(gomock.Any(), "old-refresh", "upstream-sub-456").
					Return(&upstream.Tokens{
						AccessToken:  "new-access",
						RefreshToken: "", // Provider did not rotate
						IDToken:      "new-id-token",
						ExpiresAt:    newExpiry,
					}, nil)
			},
			setupStorage: func(_ *testing.T, s *storagemocks.MockUpstreamTokenStorage) {
				s.EXPECT().StoreUpstreamTokens(gomock.Any(), "session-2", "github", gomock.Any()).
					DoAndReturn(func(_ context.Context, _, _ string, tokens *storage.UpstreamTokens) error {
						assert.Equal(t, "old-refresh", tokens.RefreshToken)
						return nil
					})
			},
			checkResult: func(t *testing.T, result *storage.UpstreamTokens) {
				t.Helper()
				assert.Equal(t, "new-access", result.AccessToken)
				assert.Equal(t, "old-refresh", result.RefreshToken)
			},
		},
		{
			// Regression for the refresh-path bound. SessionExpiresAt must be carried
			// forward unchanged so a refresh that returns a token with zero ExpiresAt
			// (provider stops asserting expires_in) still has a storage TTL bound.
			// Without this, the row would be stored with no TTL and leak indefinitely.
			name:      "preserves SessionExpiresAt when provider omits expires_in",
			sessionID: "session-bound",
			expired:   baseExpired,
			setupProvider: func(_ *testing.T, p *upstreammocks.MockOAuth2Provider) {
				p.EXPECT().RefreshTokens(gomock.Any(), "old-refresh", "upstream-sub-456").
					Return(&upstream.Tokens{
						AccessToken:  "new-access",
						RefreshToken: "new-refresh",
						// ExpiresAt intentionally zero — provider omitted expires_in.
					}, nil)
			},
			setupStorage: func(_ *testing.T, s *storagemocks.MockUpstreamTokenStorage) {
				s.EXPECT().StoreUpstreamTokens(gomock.Any(), "session-bound", "github", gomock.Any()).
					DoAndReturn(func(_ context.Context, _, _ string, tokens *storage.UpstreamTokens) error {
						assert.Equal(t, sessionBound, tokens.SessionExpiresAt,
							"refresher must carry SessionExpiresAt forward unchanged")
						assert.True(t, tokens.ExpiresAt.IsZero(),
							"new ExpiresAt should be zero (provider omitted expires_in)")
						return nil
					})
			},
			checkResult: func(t *testing.T, result *storage.UpstreamTokens) {
				t.Helper()
				assert.Equal(t, sessionBound, result.SessionExpiresAt,
					"returned tokens must also carry SessionExpiresAt forward")
				assert.True(t, result.ExpiresAt.IsZero())
			},
		},
		{
			// Defensive re-anchor for legacy data. Pre-PR Redis rows decode
			// SessionExpiresAt as zero (the field was not persisted). If such a
			// row is refreshed and the upstream rotation also drops expires_in,
			// both bounds are zero — the row would be stored without any TTL,
			// and the Memory backend would retain it indefinitely. The refresher
			// must re-anchor SessionExpiresAt to now+RefreshTokenLifespan so the
			// row carries a storage TTL bound forward.
			name:      "re-anchors SessionExpiresAt when legacy row and provider both omit expiry",
			sessionID: "session-legacy",
			expired: &storage.UpstreamTokens{
				ProviderID:       "github",
				AccessToken:      "old-access",
				RefreshToken:     "old-refresh",
				IDToken:          "old-id-token",
				ExpiresAt:        time.Time{}, // legacy row decoded with zero expiry
				SessionExpiresAt: time.Time{}, // legacy row missing the field entirely
				UserID:           "user-123",
				UpstreamSubject:  "upstream-sub-456",
				ClientID:         "client-abc",
			},
			setupProvider: func(_ *testing.T, p *upstreammocks.MockOAuth2Provider) {
				p.EXPECT().RefreshTokens(gomock.Any(), "old-refresh", "upstream-sub-456").
					Return(&upstream.Tokens{
						AccessToken:  "new-access",
						RefreshToken: "new-refresh",
						IDToken:      "new-id-token",
						// ExpiresAt intentionally zero — provider also omitted expires_in.
					}, nil)
			},
			setupStorage: func(_ *testing.T, s *storagemocks.MockUpstreamTokenStorage) {
				s.EXPECT().StoreUpstreamTokens(gomock.Any(), "session-legacy", "github", gomock.Any()).
					DoAndReturn(func(_ context.Context, _, _ string, tokens *storage.UpstreamTokens) error {
						assert.False(t, tokens.SessionExpiresAt.IsZero(),
							"refresher must re-anchor SessionExpiresAt for legacy zero/zero rows")
						assert.True(t, tokens.ExpiresAt.IsZero(),
							"new ExpiresAt should be zero (provider omitted expires_in)")
						// The re-anchor uses the configured lifespan (24h in this test).
						assert.WithinDuration(t, time.Now().Add(24*time.Hour), tokens.SessionExpiresAt, time.Minute,
							"re-anchored SessionExpiresAt should be ~now+RefreshTokenLifespan")
						return nil
					})
			},
			checkResult: func(t *testing.T, result *storage.UpstreamTokens) {
				t.Helper()
				assert.False(t, result.SessionExpiresAt.IsZero(),
					"returned tokens must also carry the re-anchored SessionExpiresAt")
				assert.True(t, result.ExpiresAt.IsZero())
			},
		},
		{
			name:           "nil expired tokens returns error",
			sessionID:      "session-3",
			expired:        nil,
			setupProvider:  func(_ *testing.T, _ *upstreammocks.MockOAuth2Provider) {},
			setupStorage:   func(_ *testing.T, _ *storagemocks.MockUpstreamTokenStorage) {},
			wantErr:        true,
			wantErrContain: "expired tokens are required",
		},
		{
			name:      "empty refresh token returns error",
			sessionID: "session-4",
			expired: &storage.UpstreamTokens{
				ProviderID:      "github",
				AccessToken:     "old-access",
				RefreshToken:    "",
				UserID:          "user-123",
				UpstreamSubject: "upstream-sub-456",
				ClientID:        "client-abc",
			},
			setupProvider:  func(_ *testing.T, _ *upstreammocks.MockOAuth2Provider) {},
			setupStorage:   func(_ *testing.T, _ *storagemocks.MockUpstreamTokenStorage) {},
			wantErr:        true,
			wantErrContain: "no refresh token available",
		},
		{
			name:      "unknown provider returns error",
			sessionID: "session-unknown",
			expired: &storage.UpstreamTokens{
				ProviderID:      "unknown-provider",
				AccessToken:     "old-access",
				RefreshToken:    "old-refresh",
				UserID:          "user-123",
				UpstreamSubject: "upstream-sub-456",
				ClientID:        "client-abc",
			},
			setupProvider:  func(_ *testing.T, _ *upstreammocks.MockOAuth2Provider) {},
			setupStorage:   func(_ *testing.T, _ *storagemocks.MockUpstreamTokenStorage) {},
			wantErr:        true,
			wantErrContain: "no upstream provider configured",
		},
		{
			name:      "provider refresh fails returns error",
			sessionID: "session-5",
			expired:   baseExpired,
			setupProvider: func(_ *testing.T, p *upstreammocks.MockOAuth2Provider) {
				p.EXPECT().RefreshTokens(gomock.Any(), "old-refresh", "upstream-sub-456").
					Return(nil, errors.New("upstream IDP unavailable"))
			},
			setupStorage:   func(_ *testing.T, _ *storagemocks.MockUpstreamTokenStorage) {},
			wantErr:        true,
			wantErrContain: "upstream token refresh failed",
		},
		{
			name:      "storage fails after refresh - returns refreshed tokens anyway",
			sessionID: "session-6",
			expired:   baseExpired,
			setupProvider: func(_ *testing.T, p *upstreammocks.MockOAuth2Provider) {
				p.EXPECT().RefreshTokens(gomock.Any(), "old-refresh", "upstream-sub-456").
					Return(&upstream.Tokens{
						AccessToken:  "new-access",
						RefreshToken: "new-refresh",
						IDToken:      "new-id-token",
						ExpiresAt:    newExpiry,
					}, nil)
			},
			setupStorage: func(_ *testing.T, s *storagemocks.MockUpstreamTokenStorage) {
				s.EXPECT().StoreUpstreamTokens(gomock.Any(), "session-6", "github", gomock.Any()).
					Return(errors.New("redis connection lost"))
			},
			checkResult: func(t *testing.T, result *storage.UpstreamTokens) {
				t.Helper()
				// Tokens should still be returned despite storage failure
				assert.Equal(t, "new-access", result.AccessToken)
				assert.Equal(t, "new-refresh", result.RefreshToken)
				assert.Equal(t, "new-id-token", result.IDToken)
				assert.Equal(t, newExpiry, result.ExpiresAt)
				// Binding fields preserved
				assert.Equal(t, "github", result.ProviderID)
				assert.Equal(t, "user-123", result.UserID)
				assert.Equal(t, "upstream-sub-456", result.UpstreamSubject)
				assert.Equal(t, "client-abc", result.ClientID)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)

			mockProvider := upstreammocks.NewMockOAuth2Provider(ctrl)
			mockStorage := storagemocks.NewMockUpstreamTokenStorage(ctrl)

			tt.setupProvider(t, mockProvider)
			tt.setupStorage(t, mockStorage)

			refresher := &upstreamTokenRefresher{
				providers:            map[string]upstream.OAuth2Provider{"github": mockProvider},
				storage:              mockStorage,
				refreshTokenLifespan: 24 * time.Hour,
			}

			result, err := refresher.RefreshAndStore(context.Background(), tt.sessionID, tt.expired)

			if tt.wantErr {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tt.wantErrContain)
				assert.Nil(t, result)
				return
			}

			require.NoError(t, err)
			require.NotNil(t, result)
			if tt.checkResult != nil {
				tt.checkResult(t, result)
			}
		})
	}
}
