package oidc

import (
	"context"
	"slices"
	"strings"
	"time"

	"github.com/zitadel/oidc/v3/pkg/oidc"
	"github.com/zitadel/oidc/v3/pkg/op"
	"golang.org/x/text/language"

	"github.com/zitadel/zitadel/internal/api/authz"
	"github.com/zitadel/zitadel/internal/api/oidc/sign"
	"github.com/zitadel/zitadel/internal/domain"
	"github.com/zitadel/zitadel/internal/telemetry/tracing"
	"github.com/zitadel/zitadel/internal/zerrors"
)

const (
	UserIDTokenType oidc.TokenType = "urn:zitadel:params:oauth:token-type:user_id"

	// TokenTypeNA is set when the returned Token Exchange access token value can't be used as an access token.
	// For example, when it is an ID Token.
	// See [RFC 8693, section 2.2.1, token_type](https://www.rfc-editor.org/rfc/rfc8693#section-2.2.1)
	TokenTypeNA = "N_A"
)

func init() {
	oidc.AllTokenTypes = append(oidc.AllTokenTypes, UserIDTokenType)
}

func (s *Server) TokenExchange(ctx context.Context, r *op.ClientRequest[oidc.TokenExchangeRequest]) (_ *op.Response, err error) {
	resp, err := s.tokenExchange(ctx, r)
	if err != nil {
		return nil, oidcError(ctx, err)
	}
	return resp, nil
}

func (s *Server) tokenExchange(ctx context.Context, r *op.ClientRequest[oidc.TokenExchangeRequest]) (_ *op.Response, err error) {
	ctx, span := tracing.NewSpan(ctx)
	defer func() { span.EndWithError(err) }()

	if len(r.Data.Resource) > 0 {
		return nil, oidc.ErrInvalidTarget().WithDescription("resource parameter not supported")
	}

	client, ok := r.Client.(*Client)
	if !ok {
		// not supposed to happen, but just preventing a panic if it does.
		return nil, zerrors.ThrowInternal(nil, "OIDC-eShi5", "Error.Internal")
	}

	subjectToken, err := s.verifyExchangeToken(ctx, client, r.Data.SubjectToken, r.Data.SubjectTokenType, oidc.AllTokenTypes...)
	if err != nil {
		return nil, oidc.ErrInvalidRequest().WithParent(err).WithReturnParentToClient(authz.GetFeatures(ctx).DebugOIDCParentError).WithDescription("subject_token invalid")
	}

	actorToken := subjectToken // see [createExchangeTokens] comment.
	// actorPath: a separate actor_token was verified (impersonation/delegation).
	// Standard exchange leaves actorToken == subjectToken; this branch runs for
	// user_id/JWT subjects or when the client sends actor_token explicitly.
	actorPath := false
	if subjectToken.tokenType == UserIDTokenType || subjectToken.tokenType == oidc.JWTTokenType || r.Data.ActorToken != "" {
		actorPath = true
		if !authz.GetInstance(ctx).EnableImpersonation() {
			return nil, zerrors.ThrowPermissionDenied(nil, "OIDC-Fae5w", "Errors.TokenExchange.Impersonation.PolicyDisabled")
		}
		actorToken, err = s.verifyExchangeToken(ctx, client, r.Data.ActorToken, r.Data.ActorTokenType, oidc.AccessTokenType, oidc.IDTokenType, oidc.RefreshTokenType)
		if err != nil {
			return nil, oidc.ErrInvalidRequest().WithParent(err).WithReturnParentToClient(authz.GetFeatures(ctx).DebugOIDCParentError).WithDescription("actor_token invalid")
		}
		ctx = authz.SetCtxData(ctx, authz.CtxData{
			UserID: actorToken.userID,
			OrgID:  actorToken.resourceOwner,
		})
	}

	audience, err := validateTokenExchangeAudience(r.Data.Audience, subjectToken.audience, actorToken.audience)
	if err != nil {
		return nil, err
	}
	// scopelessActorPath: actorPath with a subject that cannot carry scopes
	// (user_id or id_token). Subject-data scopes (email, profile, …) are taken
	// from the impersonated user at mint time; only authorization scopes must
	// appear on subject_token ∪ actor_token. Access/JWT subjects always use
	// union validation even when their scope claim is empty.
	scopelessActorPath := actorPath && (subjectToken.tokenType == UserIDTokenType || subjectToken.tokenType == oidc.IDTokenType)
	scopes, err := validateTokenExchangeScopes(client, r.Data.Scopes, subjectToken.scopes, actorToken.scopes, scopelessActorPath)
	if err != nil {
		return nil, err
	}

	resp, err := s.createExchangeTokens(ctx, r.Data.RequestedTokenType, client, subjectToken, actorToken, audience, scopes)
	if err != nil {
		return nil, err
	}

	return op.NewResponse(resp), nil
}

// verifyExchangeToken verifies the passed token based on the token type. It is safe to pass both from the request as-is.
// A list of allowed token types must be passed to determine which types are trusted at a particular stage of the token exchange.
func (s *Server) verifyExchangeToken(ctx context.Context, client *Client, token string, tokenType oidc.TokenType, allowed ...oidc.TokenType) (*exchangeToken, error) {
	if token == "" {
		return nil, zerrors.ThrowInvalidArgument(nil, "OIDC-lei0O", "Errors.TokenExchange.Token.Missing")
	}
	if tokenType == "" {
		return nil, zerrors.ThrowInvalidArgument(nil, "OIDC-sei9V", "Errors.TokenExchange.Token.TypeMissing")
	}
	if !slices.Contains(allowed, tokenType) {
		return nil, zerrors.ThrowInvalidArgument(nil, "OIDC-OZ1ie", "Errors.TokenExchange.Token.TypeNotAllowed")
	}

	switch tokenType {
	case oidc.AccessTokenType:
		token, err := s.verifyAccessToken(ctx, token)
		if err != nil {
			return nil, zerrors.ThrowPermissionDenied(err, "OIDC-Osh3t", "Errors.TokenExchange.Token.Invalid")
		}
		if !token.isPAT {
			if err = validateIntrospectionAudience(token.audience, client.GetID(), client.client.ProjectID); err != nil {
				return nil, zerrors.ThrowPermissionDenied(err, "OIDC-zi9Y0", "Errors.TokenExchange.Token.Invalid")
			}
		}
		if token.isPAT {
			if err = s.assertClientScopesForPAT(ctx, token, client.GetID(), client.client.ProjectID); err != nil {
				return nil, err
			}
		}
		return accessToExchangeToken(token, op.IssuerFromContext(ctx)), nil

	case oidc.IDTokenType:
		verifier := op.NewIDTokenHintVerifier(op.IssuerFromContext(ctx), s.idTokenHintKeySet)
		claims, err := op.VerifyIDTokenHint[*oidc.IDTokenClaims](ctx, token, verifier)
		if err != nil {
			return nil, zerrors.ThrowPermissionDenied(err, "OIDC-Rei0f", "Errors.TokenExchange.Token.Invalid")
		}
		resourceOwner, ok := claims.Claims[ClaimResourceOwnerID].(string)
		if !ok || resourceOwner == "" {
			user, err := s.query.GetUserByID(ctx, false, token)
			if err != nil {
				return nil, zerrors.ThrowPermissionDenied(err, "OIDC-aD0Oo", "Errors.TokenExchange.Token.Invalid")
			}
			resourceOwner = user.ResourceOwner
		}

		return idTokenClaimsToExchangeToken(claims, resourceOwner), nil

	case oidc.JWTTokenType:
		var (
			resourceOwner     string
			preferredLanguage *language.Tag
		)
		verifier := op.NewJWTProfileVerifierKeySet(keySetMap(client.client.PublicKeys), op.IssuerFromContext(ctx), time.Hour, client.client.ClockSkew, s.jwtProfileUserCheck(ctx, &resourceOwner, &preferredLanguage))
		jwt, err := op.VerifyJWTAssertion(ctx, token, verifier)
		if err != nil {
			return nil, zerrors.ThrowPermissionDenied(err, "OIDC-eiS6o", "Errors.TokenExchange.Token.Invalid")
		}
		return jwtToExchangeToken(jwt, resourceOwner, preferredLanguage), nil

	case UserIDTokenType:
		user, err := s.query.GetUserByID(ctx, false, token)
		if err != nil {
			return nil, zerrors.ThrowPermissionDenied(err, "OIDC-Nee6r", "Errors.TokenExchange.Token.Invalid")
		}
		return userToExchangeToken(user), nil

	case oidc.RefreshTokenType:
		fallthrough
	default:
		return nil, zerrors.ThrowInvalidArgument(nil, "OIDC-oda4R", "Errors.TokenExchange.Token.TypeNotSupported")
	}
}

// jwtProfileUserCheck finds the user by subject (user ID) and sets the resourceOwner through the pointer.
// preferred Language is set only if it was defined for a Human user, else the pointed pointer remains nil.
func (s *Server) jwtProfileUserCheck(ctx context.Context, resourceOwner *string, preferredLanguage **language.Tag) op.JWTProfileVerifierOption {
	return op.SubjectCheck(func(request *oidc.JWTTokenRequest) error {
		user, err := s.query.GetUserByID(ctx, false, request.Subject)
		if err != nil {
			return zerrors.ThrowPermissionDenied(err, "OIDC-Nee6r", "Errors.TokenExchange.Token.Invalid")
		}
		*resourceOwner = user.ResourceOwner
		if user.Human != nil && !user.Human.PreferredLanguage.IsRoot() {
			*preferredLanguage = &user.Human.PreferredLanguage
		}
		return nil
	})
}

// validateTokenExchangeScopes picks the scope rule set for this exchange.
// scopelessActorPath selects impersonation rules (user_id/id_token subjects
// on the actor path); otherwise every requested scope must ⊆ subject ∪ actor,
// with OrgRoleIDScope handled as a downscoping filter (see
// applyOrgRoleIDScopeDownscoping).
func validateTokenExchangeScopes(
	client *Client,
	requestedScopes, subjectScopes, actorScopes []string,
	scopelessActorPath bool,
) ([]string, error) {
	requestedScopes = normalizeRequestedScopes(requestedScopes)
	// Reject a client-requested bare prefix only. Input tokens may still carry
	// it (authorize accepts it); empty-request fallback inherits that ceiling.
	for _, scope := range requestedScopes {
		if scope == domain.OrgRoleIDScope {
			return nil, oidc.ErrInvalidScope().
				WithDescription("scope %q requires an organization id", scope)
		}
	}

	if scopelessActorPath {
		return validateImpersonationTokenExchangeScopes(client, requestedScopes, subjectScopes, actorScopes)
	}

	return validateUnionTokenExchangeScopes(client, requestedScopes, subjectScopes, actorScopes)
}

// validateUnionTokenExchangeScopes applies when the requested scopes must be a subset
// of the scopes present on the subject and/or actor tokens (union).
// This covers standard exchange (actor == subject) and actor-path exchanges with
// access/JWT subject tokens. Actor-path exchanges with scope-less subjects
// (user_id, id_token) use validateImpersonationTokenExchangeScopes instead.
//
// Non-restriction scopes must be present on subject ∪ actor. OrgRoleIDScope is
// validated separately against the subject token's filter ceiling so a broad
// subject can be narrowed, but a previously narrowed subject filter cannot be
// widened (actor filters are ignored on this path).
func validateUnionTokenExchangeScopes(
	client *Client,
	requestedScopes, subjectScopes, actorScopes []string,
) ([]string, error) {
	if len(requestedScopes) == 0 {
		requestedScopes = subjectScopes
	}
	if len(requestedScopes) == 0 {
		requestedScopes = actorScopes
	}

	for _, scope := range requestedScopes {
		if isTokenExchangeRestrictionScope(scope) {
			continue
		}
		if !scopeInUnion(scope, subjectScopes, actorScopes) {
			return nil, oidc.ErrInvalidScope().
				WithDescription("scope %q not found in subject or actor token", scope)
		}
	}

	requestedScopes, err := applyOrgRoleIDScopeDownscoping(requestedScopes, subjectScopes)
	if err != nil {
		return nil, err
	}

	return op.ValidateAuthReqScopes(client, requestedScopes)
}

// validateImpersonationTokenExchangeScopes applies when the subject token
// cannot carry scopes (user_id, id_token). The actor token authorizes the
// request; impersonation permission is checked later in CreateOIDCSession.
//
// Scope classes:
//   - authorization scopes: privilege/audience/roles — must be on input tokens
//   - restriction scopes: OrgRoleIDScope filter — downscope-only vs actor filters
//   - subject-data scopes: user claims (email, profile, …) — client allowlist only
func validateImpersonationTokenExchangeScopes(
	client *Client,
	requestedScopes, subjectScopes, actorScopes []string,
) ([]string, error) {
	// Subject has no scopes to inherit; fall back to actor.
	if len(requestedScopes) == 0 {
		requestedScopes = actorScopes
	}

	for _, scope := range requestedScopes {
		// Prevent privilege escalation: offline_access, audiences, roles, etc.
		// must still come from subject or actor token.
		if isTokenExchangeAuthorizationScope(scope) && !scopeInUnion(scope, subjectScopes, actorScopes) {
			return nil, oidc.ErrInvalidScope().
				WithDescription("scope %q not found in subject or actor token", scope)
		}

		// Restriction and subject-data scopes skip the union check —
		// restriction scopes are validated below; subject-data by client allowlist.
	}

	// Subject cannot carry scopes; the actor filter is the only available ceiling.
	requestedScopes, err := applyOrgRoleIDScopeDownscoping(requestedScopes, actorScopes)
	if err != nil {
		return nil, err
	}

	return op.ValidateAuthReqScopes(client, requestedScopes)
}

func normalizeRequestedScopes(scopes []string) []string {
	// Space-delimited empty scope produces a single "" entry.
	return slices.DeleteFunc(scopes, func(s string) bool {
		return s == ""
	})
}

func scopeInUnion(scope string, subjectScopes, actorScopes []string) bool {
	return slices.Contains(subjectScopes, scope) || slices.Contains(actorScopes, scope)
}

// isTokenExchangeRestrictionScope identifies scopes that only filter token
// claims. They may be newly requested when the input tokens are unfiltered,
// but must not expand an existing filter — see applyOrgRoleIDScopeDownscoping.
func isTokenExchangeRestrictionScope(scope string) bool {
	return strings.HasPrefix(scope, domain.OrgRoleIDScope)
}

// orgRoleIDScopes returns the OrgRoleIDScope entries from scopes, preserving order
// and dropping duplicates.
func orgRoleIDScopes(scopes []string) []string {
	var out []string
	for _, scope := range scopes {
		if !isTokenExchangeRestrictionScope(scope) {
			continue
		}
		if !slices.Contains(out, scope) {
			out = append(out, scope)
		}
	}
	return out
}

// applyOrgRoleIDScopeDownscoping enforces that OrgRoleIDScope can only narrow
// role-org filters relative to ceilingScopes (the caller's chosen ceiling):
//   - union / standard exchange passes the subject token scopes
//   - impersonation (scopeless subject) passes the actor token scopes
//
// Rules:
//   - ceiling unfiltered → any requested filter is allowed (broad → narrow)
//   - ceiling filtered → requested filters must be ⊆ that set
//   - ceiling filtered but request omits the filter → inherit the ceiling
//     (omitting would otherwise re-widen role claims to all granted orgs)
func applyOrgRoleIDScopeDownscoping(requested, ceilingScopes []string) ([]string, error) {
	allowed := orgRoleIDScopes(ceilingScopes)
	requestedFilters := orgRoleIDScopes(requested)

	if len(allowed) == 0 {
		return requested, nil
	}
	if len(requestedFilters) == 0 {
		return append(slices.Clone(requested), allowed...), nil
	}
	for _, scope := range requestedFilters {
		if !slices.Contains(allowed, scope) {
			return nil, oidc.ErrInvalidScope().
				WithDescription("scope %q expands org role filter beyond the permitted downscope set", scope)
		}
	}
	return requested, nil
}

// isTokenExchangeAuthorizationScope identifies scopes that control what a token
// may do (audiences, refresh, roles), not which user claims to embed.
// Restriction scopes are intentionally excluded — see isTokenExchangeRestrictionScope.
func isTokenExchangeAuthorizationScope(scope string) bool {
	switch scope {
	case oidc.ScopeOfflineAccess, ScopeProjectsRoles:
		return true
	}
	if strings.HasPrefix(scope, ScopeProjectRolePrefix) {
		return true // urn:zitadel:iam:org:project:role:*
	}
	if strings.HasPrefix(scope, domain.ProjectIDScope) {
		return true // urn:zitadel:iam:org:project:id:*:aud
	}
	if strings.HasPrefix(scope, domain.OrgIDScope) {
		return true
	}
	if strings.HasPrefix(scope, domain.OrgDomainPrimaryScope) {
		return true
	}
	if strings.HasPrefix(scope, domain.SelectIDPScope) {
		return true
	}
	return false
}

func validateTokenExchangeAudience(requestedAudience, subjectAudience, actorAudience []string) ([]string, error) {
	if len(requestedAudience) == 0 {
		if len(subjectAudience) > 0 {
			return subjectAudience, nil
		}
		if len(actorAudience) > 0 {
			return actorAudience, nil
		}
	}
	if slices.Equal(requestedAudience, subjectAudience) || slices.Equal(requestedAudience, actorAudience) {
		return requestedAudience, nil
	}

	//nolint:gocritic
	allowedAudience := append(subjectAudience, actorAudience...)
	for _, a := range requestedAudience {
		if !slices.Contains(allowedAudience, a) {
			return nil, oidc.ErrInvalidTarget().WithDescription("audience %q not found in subject or actor token", a)
		}
	}
	return requestedAudience, nil
}

// createExchangeTokens prepares the final tokens to be returned to the client.
// The subjectToken is used to set the new token's subject and resource owner.
// The actorToken is used to set the new token's auth time AMR and actor.
// Both tokens may point to the same object (subjectToken) in case of a regular Token Exchange.
// When the subject and actor Tokens point to different objects, the new tokens will be for impersonation / delegation.
func (s *Server) createExchangeTokens(ctx context.Context, tokenType oidc.TokenType, client *Client, subjectToken, actorToken *exchangeToken, audience, scopes []string) (_ *oidc.TokenExchangeResponse, err error) {
	getUserInfo := s.getUserInfo(subjectToken.userID, client.client.ProjectID, client.GetID(), client.client.ProjectRoleAssertion, client.IDTokenUserinfoClaimsAssertion(), scopes)
	getSigner := s.getSignerOnce()

	resp := &oidc.TokenExchangeResponse{
		Scopes: scopes,
	}

	reason := domain.TokenReasonExchange
	actor := actorToken.actor
	if subjectToken != actorToken {
		reason = domain.TokenReasonImpersonation
		actor = actorToken.nestedActor()
	}

	var sessionID string
	switch tokenType {
	case oidc.AccessTokenType, "":
		resp.AccessToken, resp.RefreshToken, sessionID, resp.ExpiresIn, err = s.createExchangeAccessToken(ctx, client, subjectToken.userID, subjectToken.resourceOwner, audience, scopes, actorToken.authMethods, actorToken.authTime, subjectToken.preferredLanguage, reason, actor)
		resp.TokenType = oidc.BearerToken
		resp.IssuedTokenType = oidc.AccessTokenType

	case oidc.JWTTokenType:
		resp.AccessToken, resp.RefreshToken, resp.ExpiresIn, err = s.createExchangeJWT(ctx, client, getUserInfo, client.client.AccessTokenRoleAssertion, getSigner, subjectToken.userID, subjectToken.resourceOwner, audience, scopes, actorToken.authMethods, actorToken.authTime, subjectToken.preferredLanguage, reason, actor)
		resp.TokenType = oidc.BearerToken
		resp.IssuedTokenType = oidc.JWTTokenType

	case oidc.IDTokenType:
		resp.AccessToken, resp.ExpiresIn, err = s.createIDToken(ctx, client, getUserInfo, client.client.IDTokenRoleAssertion, getSigner, "", resp.AccessToken, audience, actorToken.authMethods, actorToken.authTime, "", actor)
		resp.TokenType = TokenTypeNA
		resp.IssuedTokenType = oidc.IDTokenType

	case oidc.RefreshTokenType, UserIDTokenType:
		fallthrough
	default:
		err = zerrors.ThrowInvalidArgument(nil, "OIDC-wai5E", "Errors.TokenExchange.Token.TypeNotSupported")
	}
	if err != nil {
		return nil, err
	}

	if slices.Contains(scopes, oidc.ScopeOpenID) && tokenType != oidc.IDTokenType {
		resp.IDToken, _, err = s.createIDToken(ctx, client, getUserInfo, client.client.IDTokenRoleAssertion, getSigner, sessionID, resp.AccessToken, audience, actorToken.authMethods, actorToken.authTime, "", actor)
		if err != nil {
			return nil, err
		}
	}

	return resp, nil
}

func (s *Server) createExchangeAccessToken(
	ctx context.Context,
	client *Client,
	userID,
	resourceOwner string,
	audience,
	scope []string,
	authMethods []domain.UserAuthMethodType,
	authTime time.Time,
	preferredLanguage *language.Tag,
	reason domain.TokenReason,
	actor *domain.TokenActor,
) (accessToken, refreshToken, sessionID string, exp uint64, err error) {
	ctx, span := tracing.NewSpan(ctx)
	defer func() { span.EndWithError(err) }()

	session, err := s.command.CreateOIDCSession(ctx,
		userID,
		resourceOwner,
		client.client.ClientID,
		client.client.BackChannelLogoutURI,
		scope,
		audience,
		authMethods,
		authTime,
		"",
		preferredLanguage,
		nil,
		reason,
		actor,
		slices.Contains(scope, oidc.ScopeOfflineAccess),
		"",
		domain.OIDCResponseTypeUnspecified,
	)
	if err != nil {
		return "", "", "", 0, err
	}
	accessToken, err = op.CreateBearerToken(session.TokenID, userID, s.opCrypto)
	if err != nil {
		return "", "", "", 0, err
	}
	return accessToken, session.RefreshToken, session.SessionID, timeToOIDCExpiresIn(session.Expiration), nil
}

func (s *Server) createExchangeJWT(
	ctx context.Context,
	client *Client,
	getUserInfo userInfoFunc,
	roleAssertion bool,
	getSigner sign.SignerFunc,
	userID,
	resourceOwner string,
	audience,
	scope []string,
	authMethods []domain.UserAuthMethodType,
	authTime time.Time,
	preferredLanguage *language.Tag,
	reason domain.TokenReason,
	actor *domain.TokenActor,
) (accessToken string, refreshToken string, exp uint64, err error) {
	ctx, span := tracing.NewSpan(ctx)
	defer func() { span.EndWithError(err) }()

	session, err := s.command.CreateOIDCSession(ctx,
		userID,
		resourceOwner,
		client.client.ClientID,
		client.client.BackChannelLogoutURI,
		scope,
		audience,
		authMethods,
		authTime,
		"",
		preferredLanguage,
		nil,
		reason,
		actor,
		slices.Contains(scope, oidc.ScopeOfflineAccess),
		"",
		domain.OIDCResponseTypeUnspecified,
	)
	if err != nil {
		return "", "", 0, err
	}
	accessToken, err = s.createJWT(ctx, client, session, getUserInfo, roleAssertion, getSigner)
	if err != nil {
		return "", "", 0, err
	}
	return accessToken, session.RefreshToken, timeToOIDCExpiresIn(session.Expiration), nil
}