// Copyright © 2023 Ory Corp
// SPDX-License-Identifier: Apache-2.0

package hook_test

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/ory/kratos/hydra"
	"github.com/ory/kratos/pkg/testhelpers"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidwall/gjson"

	"github.com/ory/kratos/driver/config"
	"github.com/ory/kratos/identity"
	"github.com/ory/kratos/pkg"
	"github.com/ory/kratos/selfservice/flow"
	"github.com/ory/kratos/selfservice/flow/recovery"
	"github.com/ory/kratos/selfservice/flow/registration"
	"github.com/ory/kratos/selfservice/flow/verification"
	"github.com/ory/kratos/selfservice/hook"
	"github.com/ory/kratos/session"
	"github.com/ory/kratos/x"
	"github.com/ory/x/randx"
)

func TestSessionIssuer(t *testing.T) {
	ctx := context.Background()
	conf, reg := pkg.NewFastRegistryWithMocks(t)
	reg.SetHydra(hydra.NewFake())
	conf.MustSet(ctx, config.ViperKeyPublicBaseURL, "http://localhost/")
	testhelpers.SetDefaultIdentitySchema(conf, "file://./stub/stub.schema.json")

	var r http.Request
	h := hook.NewSessionIssuer(reg)

	t.Run("method=sign-up", func(t *testing.T) {
		t.Run("flow=browser", func(t *testing.T) {
			w := httptest.NewRecorder()
			s := testhelpers.CreateSession(t, reg)
			f := &registration.Flow{Type: flow.TypeBrowser}

			require.NoError(t, h.ExecutePostRegistrationPostPersistHook(w, &r,
				f, &session.Session{ID: s.ID, Identity: s.Identity, Token: randx.MustString(12, randx.AlphaLowerNum)}))

			require.Empty(t, f.ContinueWithItems)

			got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
			require.NoError(t, err)
			assert.Equal(t, s.ID, got.ID)
			assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

			assert.Contains(t, w.Header().Get("Set-Cookie"), config.DefaultSessionCookieName)
		})

		t.Run("flow=api", func(t *testing.T) {
			w := httptest.NewRecorder()

			i := identity.NewIdentity(config.DefaultIdentityTraitsSchemaID)
			s := &session.Session{
				ID:              x.NewUUID(),
				Identity:        i,
				Token:           randx.MustString(12, randx.AlphaLowerNum),
				LogoutToken:     randx.MustString(12, randx.AlphaLowerNum),
				AuthenticatedAt: time.Now().UTC(),
			}
			f := &registration.Flow{Type: flow.TypeAPI}

			require.NoError(t, reg.PrivilegedIdentityPool().CreateIdentity(context.Background(), i))
			require.NoError(t, reg.SessionPersister().UpsertSession(ctx, s))

			err := h.ExecutePostRegistrationPostPersistHook(w, &http.Request{Header: http.Header{"Accept": {"application/json"}}}, f, s)
			require.ErrorIs(t, err, registration.ErrHookAbortFlow, "%+v", err)
			require.Len(t, f.ContinueWithItems, 1)

			st := f.ContinueWithItems[0]
			require.IsType(t, &flow.ContinueWithSetOrySessionToken{}, st)
			assert.NotEmpty(t, st.(*flow.ContinueWithSetOrySessionToken).OrySessionToken)

			got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
			require.NoError(t, err)
			assert.Equal(t, s.ID.String(), got.ID.String())
			assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

			assert.Empty(t, w.Header().Get("Set-Cookie"))
			body := w.Body.Bytes()
			assert.Equal(t, i.ID.String(), gjson.GetBytes(body, "identity.id").String())
			assert.Equal(t, s.ID.String(), gjson.GetBytes(body, "session.id").String())
			assert.Equal(t, got.Token, gjson.GetBytes(body, "session_token").String())
		})

		t.Run("flow=spa", func(t *testing.T) {
			w := httptest.NewRecorder()

			i := identity.NewIdentity(config.DefaultIdentityTraitsSchemaID)
			s := &session.Session{
				ID:              x.NewUUID(),
				Identity:        i,
				Token:           randx.MustString(12, randx.AlphaLowerNum),
				LogoutToken:     randx.MustString(12, randx.AlphaLowerNum),
				AuthenticatedAt: time.Now().UTC(),
			}
			f := &registration.Flow{Type: flow.TypeBrowser}

			require.NoError(t, reg.PrivilegedIdentityPool().CreateIdentity(context.Background(), i))
			require.NoError(t, reg.SessionPersister().UpsertSession(ctx, s))

			err := h.ExecutePostRegistrationPostPersistHook(w, &http.Request{Header: http.Header{"Accept": {"application/json"}}}, f, s)
			require.ErrorIs(t, err, registration.ErrHookAbortFlow, "%+v", err)
			require.Empty(t, f.ContinueWithItems)

			got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
			require.NoError(t, err)
			assert.Equal(t, s.ID.String(), got.ID.String())
			assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

			assert.NotEmpty(t, w.Header().Get("Set-Cookie"))
			body := w.Body.Bytes()
			assert.Equal(t, i.ID.String(), gjson.GetBytes(body, "identity.id").String())
			assert.Equal(t, s.ID.String(), gjson.GetBytes(body, "session.id").String())
			assert.Empty(t, gjson.GetBytes(body, "session_token").String())
		})

		t.Run("flow=api with oidc code exchange", func(t *testing.T) {
			// A native API OIDC flow with a session token exchange code redirects
			// via MaybeRedirectAPICodeFlow. Once the redirect is written, the hook
			// must abort the flow (ErrHookAbortFlow) so the executor's tail block
			// does not call MaybeRedirectAPICodeFlow a second time and issue a
			// redundant write to the GLOBAL session_token_exchanges table.
			w := httptest.NewRecorder()

			i := identity.NewIdentity(config.DefaultIdentityTraitsSchemaID)
			s := &session.Session{
				ID:              x.NewUUID(),
				Identity:        i,
				Token:           randx.MustString(12, randx.AlphaLowerNum),
				LogoutToken:     randx.MustString(12, randx.AlphaLowerNum),
				AuthenticatedAt: time.Now().UTC(),
			}
			s.CompletedLoginFor(identity.CredentialsTypeOIDC, identity.AuthenticatorAssuranceLevel1)
			f := &registration.Flow{ID: x.NewUUID(), Type: flow.TypeAPI}

			require.NoError(t, reg.PrivilegedIdentityPool().CreateIdentity(ctx, i))
			require.NoError(t, reg.SessionPersister().UpsertSession(ctx, s))

			ex, err := reg.SessionTokenExchangePersister().CreateSessionTokenExchanger(ctx, f.ID)
			require.NoError(t, err)

			req := httptest.NewRequest(http.MethodGet, "https://localhost/self-service/methods/oidc/callback/test", nil).WithContext(ctx)
			req.Header.Set("Accept", "application/json")
			err = h.ExecutePostRegistrationPostPersistHook(w, req, f, s)
			require.ErrorIs(t, err, registration.ErrHookAbortFlow, "%+v", err)

			// The redirect carries the return-to code, proving the OIDC handled
			// branch ran before aborting.
			assert.Equal(t, http.StatusSeeOther, w.Code)
			assert.Contains(t, w.Header().Get("Location"), "code="+ex.ReturnToCode)

			// The exchanger was bound to the session exactly once.
			got, _, err := reg.SessionTokenExchangePersister().CodeForFlow(ctx, f.ID)
			require.NoError(t, err)
			assert.Equal(t, ex.ReturnToCode, got.ReturnToCode)
			exchanged, err := reg.SessionTokenExchangePersister().GetExchangerFromCode(ctx, ex.InitCode, ex.ReturnToCode)
			require.NoError(t, err)
			assert.Equal(t, s.ID, exchanged.SessionID.UUID)
		})
	})

	t.Run("method=sign-up with oauth2", func(t *testing.T) {
		t.Run("flow=browser", func(t *testing.T) {
			w := httptest.NewRecorder()
			s := testhelpers.CreateSession(t, reg)
			f := &registration.Flow{
				Type:                 flow.TypeBrowser,
				OAuth2LoginChallenge: hydra.FakeValidLoginChallenge,
			}

			require.NoError(t, h.ExecutePostRegistrationPostPersistHook(w, &r,
				f, &session.Session{ID: s.ID, Identity: s.Identity, Token: randx.MustString(12, randx.AlphaLowerNum)}))

			require.Empty(t, f.ContinueWithItems)

			got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
			require.NoError(t, err)
			assert.Equal(t, s.ID, got.ID)
			assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

			assert.Contains(t, w.Header().Get("Set-Cookie"), config.DefaultSessionCookieName)
		})

		t.Run("flow=api", func(t *testing.T) {
			w := httptest.NewRecorder()

			i := identity.NewIdentity(config.DefaultIdentityTraitsSchemaID)
			s := &session.Session{
				ID:              x.NewUUID(),
				Identity:        i,
				Token:           randx.MustString(12, randx.AlphaLowerNum),
				LogoutToken:     randx.MustString(12, randx.AlphaLowerNum),
				AuthenticatedAt: time.Now().UTC(),
			}
			f := &registration.Flow{
				Type:                 flow.TypeAPI,
				OAuth2LoginChallenge: hydra.FakeValidLoginChallenge,
			}

			require.NoError(t, reg.PrivilegedIdentityPool().CreateIdentity(context.Background(), i))
			require.NoError(t, reg.SessionPersister().UpsertSession(ctx, s))

			err := h.ExecutePostRegistrationPostPersistHook(w, &http.Request{Header: http.Header{"Accept": {"application/json"}}}, f, s)
			require.ErrorIs(t, err, registration.ErrHookAbortFlow, "%+v", err)
			require.Len(t, f.ContinueWithItems, 1)

			st := f.ContinueWithItems[0]
			require.IsType(t, &flow.ContinueWithSetOrySessionToken{}, st)
			assert.NotEmpty(t, st.(*flow.ContinueWithSetOrySessionToken).OrySessionToken)

			got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
			require.NoError(t, err)
			assert.Equal(t, s.ID.String(), got.ID.String())
			assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

			assert.Empty(t, w.Header().Get("Set-Cookie"))
			body := w.Body.Bytes()
			assert.Equal(t, i.ID.String(), gjson.GetBytes(body, "identity.id").String())
			assert.Equal(t, s.ID.String(), gjson.GetBytes(body, "session.id").String())
			assert.Equal(t, got.Token, gjson.GetBytes(body, "session_token").String())
		})

		t.Run("flow=spa", func(t *testing.T) {
			w := httptest.NewRecorder()

			i := identity.NewIdentity(config.DefaultIdentityTraitsSchemaID)
			s := &session.Session{
				ID:              x.NewUUID(),
				Identity:        i,
				Token:           randx.MustString(12, randx.AlphaLowerNum),
				LogoutToken:     randx.MustString(12, randx.AlphaLowerNum),
				AuthenticatedAt: time.Now().UTC(),
			}
			f := &registration.Flow{
				Type:                 flow.TypeBrowser,
				OAuth2LoginChallenge: hydra.FakeValidLoginChallenge,
				ContinueWithItems:    []flow.ContinueWith{flow.NewContinueWithRedirectBrowserTo("https://ory.sh")},
			}

			require.NoError(t, reg.PrivilegedIdentityPool().CreateIdentity(context.Background(), i))
			require.NoError(t, reg.SessionPersister().UpsertSession(ctx, s))

			err := h.ExecutePostRegistrationPostPersistHook(w, &http.Request{Header: http.Header{"Accept": {"application/json"}}}, f, s)
			require.ErrorIs(t, err, registration.ErrHookAbortFlow, "%+v", err)
			require.Len(t, f.ContinueWithItems, 1)
			require.EqualValues(t, flow.ContinueWithActionRedirectBrowserToString, f.ContinueWithItems[0].GetAction())
			require.EqualValues(t, hydra.FakePostLoginURL, f.ContinueWithItems[0].(*flow.ContinueWithRedirectBrowserTo).RedirectTo)

			got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
			require.NoError(t, err)
			assert.Equal(t, s.ID.String(), got.ID.String())
			assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

			assert.NotEmpty(t, w.Header().Get("Set-Cookie"))
			body := w.Body.Bytes()
			assert.Equal(t, i.ID.String(), gjson.GetBytes(body, "identity.id").String())
			assert.Equal(t, s.ID.String(), gjson.GetBytes(body, "session.id").String())
			assert.Empty(t, gjson.GetBytes(body, "session_token").String())
		})
	})

	t.Run("method=sign-up with oauth2 and verification enabled", func(t *testing.T) {
		w := httptest.NewRecorder()

		i := identity.NewIdentity(config.DefaultIdentityTraitsSchemaID)
		s := &session.Session{
			ID:              x.NewUUID(),
			Identity:        i,
			Token:           randx.MustString(12, randx.AlphaLowerNum),
			LogoutToken:     randx.MustString(12, randx.AlphaLowerNum),
			AuthenticatedAt: time.Now().UTC(),
		}
		vf := &verification.Flow{
			ID: x.NewUUID(),
		}
		f := &registration.Flow{
			Type:                 flow.TypeBrowser,
			OAuth2LoginChallenge: hydra.FakeValidLoginChallenge,
			ContinueWithItems:    []flow.ContinueWith{flow.NewContinueWithVerificationUI(vf.ID, "some@ory.sh", "")},
		}

		require.NoError(t, reg.PrivilegedIdentityPool().CreateIdentity(context.Background(), i))
		require.NoError(t, reg.SessionPersister().UpsertSession(ctx, s))

		err := h.ExecutePostRegistrationPostPersistHook(w, &http.Request{Header: http.Header{"Accept": {"application/json"}}}, f, s)
		require.ErrorIs(t, err, registration.ErrHookAbortFlow, "%+v", err)
		require.Len(t, f.ContinueWithItems, 1)
		require.EqualValues(t, flow.ContinueWithActionShowVerificationUIString, f.ContinueWithItems[0].GetAction())

		got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
		require.NoError(t, err)
		assert.Equal(t, s.ID.String(), got.ID.String())
		assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

		assert.NotEmpty(t, w.Header().Get("Set-Cookie"))
		body := w.Body.Bytes()
		assert.Equal(t, i.ID.String(), gjson.GetBytes(body, "identity.id").String())
		assert.Equal(t, s.ID.String(), gjson.GetBytes(body, "session.id").String())
		assert.Empty(t, gjson.GetBytes(body, "session_token").String())
	})

	t.Run("method=sign-up with oauth2 and some other continue with item", func(t *testing.T) {
		w := httptest.NewRecorder()

		i := identity.NewIdentity(config.DefaultIdentityTraitsSchemaID)
		s := &session.Session{
			ID:              x.NewUUID(),
			Identity:        i,
			Token:           randx.MustString(12, randx.AlphaLowerNum),
			LogoutToken:     randx.MustString(12, randx.AlphaLowerNum),
			AuthenticatedAt: time.Now().UTC(),
		}
		vf := &recovery.Flow{
			ID: x.NewUUID(),
		}
		f := &registration.Flow{
			Type:                 flow.TypeBrowser,
			OAuth2LoginChallenge: hydra.FakeValidLoginChallenge,
			ContinueWithItems:    []flow.ContinueWith{flow.NewContinueWithRecoveryUI(vf)},
		}

		require.NoError(t, reg.PrivilegedIdentityPool().CreateIdentity(context.Background(), i))
		require.NoError(t, reg.SessionPersister().UpsertSession(ctx, s))

		err := h.ExecutePostRegistrationPostPersistHook(w, &http.Request{Header: http.Header{"Accept": {"application/json"}}}, f, s)
		require.ErrorIs(t, err, registration.ErrHookAbortFlow, "%+v", err)
		require.Len(t, f.ContinueWithItems, 2)
		require.EqualValues(t, flow.ContinueWithActionShowRecoveryUIString, f.ContinueWithItems[0].GetAction())
		require.EqualValues(t, hydra.FakePostLoginURL, f.ContinueWithItems[1].(*flow.ContinueWithRedirectBrowserTo).RedirectTo)

		got, err := reg.SessionPersister().GetSession(context.Background(), s.ID, session.ExpandNothing)
		require.NoError(t, err)
		assert.Equal(t, s.ID.String(), got.ID.String())
		assert.True(t, got.AuthenticatedAt.After(time.Now().Add(-time.Minute)))

		assert.NotEmpty(t, w.Header().Get("Set-Cookie"))
		body := w.Body.Bytes()
		assert.Equal(t, i.ID.String(), gjson.GetBytes(body, "identity.id").String())
		assert.Equal(t, s.ID.String(), gjson.GetBytes(body, "session.id").String())
		assert.Empty(t, gjson.GetBytes(body, "session_token").String())
	})
}