package runtimeidentity

import (
	"context"
	"encoding/base64"
	"errors"
	"strings"
	"testing"
	"time"

	featureauth "github.com/nexuscrm/auth"
	kernelerrors "github.com/nexuscrm/kernel/pkg/errors"
	semantic "github.com/pingcap-inc/nexuscrm-contract/pkg/runtime/federatedloginv1"
)

const federatedLoginTransactionTestState = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"

func (stub *authorizationOAuthStub) BeginSSOLoginTransaction(_ context.Context, redirectURI string) (*featureauth.SSOLoginTransaction, error) {
	stub.federatedBeginRedirect = redirectURI
	return stub.federatedBeginResult, stub.federatedBeginErr
}

func (stub *authorizationOAuthStub) CompleteSSOLoginTransaction(_ context.Context, state, redirectURI string) error {
	stub.federatedConsumeState = state
	stub.federatedConsumeRedirect = redirectURI
	return stub.federatedConsumeErr
}

func TestAuthorizationLifecycleProjectsTheExistingOneTimeFederatedLoginOwner(t *testing.T) {
	now := time.Unix(1_800_000_000, 0)
	oauth := &authorizationOAuthStub{federatedBeginResult: &featureauth.SSOLoginTransaction{
		State: federatedLoginTransactionTestState, ExpiresAt: now.Add(time.Minute),
	}}
	owner := newAuthorizationLifecycleTest(t, &sessionAuthStub{}, oauth, now)
	ctx := sessionTestContext(t)
	request := semantic.BeginRequest{
		AppInstallationID: "installation-1",
		RedirectURI:       "https://app.example.test/oauth/callback",
	}
	response, err := owner.BeginFederatedLogin(ctx, authorizationTestAuthority(), request)
	if err != nil || response.State != federatedLoginTransactionTestState ||
		response.ExpiresAtUnixSecond != now.Add(time.Minute).Unix() ||
		oauth.federatedBeginRedirect != request.RedirectURI {
		t.Fatalf("begin=%#v redirect=%q err=%v", response, oauth.federatedBeginRedirect, err)
	}
	if err := owner.ConsumeFederatedLogin(ctx, authorizationTestAuthority(), semantic.ConsumeRequest{
		AppInstallationID: request.AppInstallationID,
		State:             response.State,
		RedirectURI:       request.RedirectURI,
	}); err != nil {
		t.Fatal(err)
	}
	if oauth.federatedConsumeState != response.State || oauth.federatedConsumeRedirect != request.RedirectURI {
		t.Fatalf("consume state=%q redirect=%q", oauth.federatedConsumeState, oauth.federatedConsumeRedirect)
	}
}

func TestAuthorizationLifecycleKeepsFederatedStateFailuresNonDisclosing(t *testing.T) {
	now := time.Unix(1_800_000_000, 0)
	oauth := &authorizationOAuthStub{
		federatedConsumeErr: kernelerrors.NewUnauthorizedError("SSO login state is invalid"),
	}
	owner := newAuthorizationLifecycleTest(t, &sessionAuthStub{}, oauth, now)
	err := owner.ConsumeFederatedLogin(sessionTestContext(t), authorizationTestAuthority(), semantic.ConsumeRequest{
		AppInstallationID: "installation-1", State: federatedLoginTransactionTestState,
		RedirectURI: "https://app.example.test/oauth/callback",
	})
	var failure *semantic.Failure
	if !errors.As(err, &failure) || failure.Category != semantic.FailureUnauthenticated {
		t.Fatalf("failure=%#v err=%v", failure, err)
	}
}

func TestAuthorizationLifecycleAdmitsOnlyTheExactPreUpgradeFederatedStateShape(t *testing.T) {
	now := time.Unix(1_800_000_000, 0)
	legacyState := "ssologin." + base64.URLEncoding.EncodeToString(make([]byte, semantic.StateEntropyBytes))
	oauth := &authorizationOAuthStub{}
	owner := newAuthorizationLifecycleTest(t, &sessionAuthStub{}, oauth, now)
	err := owner.ConsumeFederatedLogin(sessionTestContext(t), authorizationTestAuthority(), semantic.ConsumeRequest{
		AppInstallationID: "installation-1", State: legacyState,
		RedirectURI: "https://app.example.test/oauth/callback",
	})
	if err != nil {
		t.Fatalf("exact pre-upgrade state was rejected before reaching its compatibility owner: %v", err)
	}
	if oauth.federatedConsumeState != legacyState {
		t.Fatalf("owner state=%q, want exact pre-upgrade state %q", oauth.federatedConsumeState, legacyState)
	}

	oauth.federatedConsumeState = ""
	err = owner.ConsumeFederatedLogin(sessionTestContext(t), authorizationTestAuthority(), semantic.ConsumeRequest{
		AppInstallationID: "installation-1", State: "ssologin." + strings.Repeat("A", semantic.MaxStateBytes),
		RedirectURI: "https://app.example.test/oauth/callback",
	})
	if err == nil || oauth.federatedConsumeState != "" {
		t.Fatalf("non-legacy prefixed state reached owner: state=%q err=%v", oauth.federatedConsumeState, err)
	}
}

func TestAuthorizationLifecycleNormalizesFederatedConsumeAdmissionFailures(t *testing.T) {
	now := time.Unix(1_800_000_000, 0)
	oauth := &authorizationOAuthStub{}
	owner := newAuthorizationLifecycleTest(t, &sessionAuthStub{}, oauth, now)
	authority := authorizationTestAuthority()
	authority.CallerWorkloadID = "other-appservice"
	err := owner.ConsumeFederatedLogin(sessionTestContext(t), authority, semantic.ConsumeRequest{
		AppInstallationID: "installation-1", State: federatedLoginTransactionTestState,
		RedirectURI: "https://app.example.test/oauth/callback",
	})
	var failure *semantic.Failure
	if !errors.As(err, &failure) || failure.Category != semantic.FailureUnauthenticated {
		t.Fatalf("failure=%#v err=%v", failure, err)
	}
	if oauth.federatedConsumeState != "" || oauth.federatedConsumeRedirect != "" {
		t.Fatalf("consume reached owner with state=%q redirect=%q", oauth.federatedConsumeState, oauth.federatedConsumeRedirect)
	}
}