// Copyright (c) 2026 Sidero Labs, Inc.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.

package saml

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"maps"
	"net/http"
	"net/mail"
	"net/url"
	"slices"
	"strings"
	"time"

	"github.com/cosi-project/runtime/pkg/resource"
	"github.com/cosi-project/runtime/pkg/resource/kvutils"
	"github.com/cosi-project/runtime/pkg/safe"
	"github.com/cosi-project/runtime/pkg/state"
	"github.com/crewjam/saml"
	"github.com/crewjam/saml/samlsp"
	"github.com/google/uuid"
	"github.com/siderolabs/gen/xslices"
	"go.uber.org/zap"

	"github.com/siderolabs/omni/client/pkg/cosi/labels"
	"github.com/siderolabs/omni/client/pkg/omni/resources/auth"
	"github.com/siderolabs/omni/internal/pkg/auth/actor"
	"github.com/siderolabs/omni/internal/pkg/auth/role"
	"github.com/siderolabs/omni/internal/pkg/auth/user"
)

// sloSessionCookieTTL is how long the SLO cookie persists. It must outlive the browser
// session so that logout can send a LogoutRequest to the IdP even after a browser restart.
const sloSessionCookieTTL = 24 * time.Hour

// UserInfo describes user identity and fullname.
type UserInfo struct {
	// Identity ...
	Identity string
	// Fullname ...
	Fullname string
}

// NewSessionProvider creates a new SessionProvider.
func NewSessionProvider(state state.State, tracker samlsp.RequestTracker, logger *zap.Logger, attributeRules map[string]string) *SessionProvider {
	return &SessionProvider{
		state:          state,
		tracker:        tracker,
		logger:         logger,
		attributeRules: attributeRules,
	}
}

// SessionProvider is an implementation of SessionProvider that stores
// session tokens in the COSI state.
type SessionProvider struct {
	state          state.State
	tracker        samlsp.RequestTracker
	logger         *zap.Logger
	attributeRules map[string]string
}

// CreateSession is called when we have received a valid SAML assertion and
// should create a new session and do redirect.
//
//nolint:gocognit,gocyclo,cyclop
func (sp *SessionProvider) CreateSession(w http.ResponseWriter, r *http.Request, assertion *saml.Assertion) error {
	var hashInput strings.Builder

	if assertion.Subject == nil {
		return errors.New("no subject in the assertion")
	}

	for _, subjectConfirmation := range assertion.Subject.SubjectConfirmations {
		hashInput.WriteString(subjectConfirmation.SubjectConfirmationData.InResponseTo)
	}

	var logFields []zap.Field

	if sub := assertion.Subject; sub != nil {
		if nameID := sub.NameID; nameID != nil {
			logFields = append(logFields, zap.String("subject_id", nameID.Value))
		}
	}

	for _, attributeStatement := range assertion.AttributeStatements {
		for _, attr := range attributeStatement.Attributes {
			name := attr.FriendlyName
			if name == "" {
				name = attr.Name
			}

			var value string

			if len(attr.Values) > 0 {
				value = attr.Values[0].Value
			}

			logFields = append(logFields, zap.String(name, value))
		}
	}

	sp.logger.Info("new SAML assertion", logFields...)

	var session string

	if hashInput.String() == "" {
		session = uuid.New().String()
	} else {
		h := sha256.New()

		h.Write([]byte(hashInput.String()))

		session = hex.EncodeToString(h.Sum(nil))
	}

	query := r.URL.Query()

	if trackedRequestIndex := r.Form.Get("RelayState"); trackedRequestIndex != "" {
		trackedRequest, err := sp.tracker.GetTrackedRequest(r, trackedRequestIndex)
		if err != nil {
			return err
		}

		redirectURI, err := url.Parse(trackedRequest.URI)
		if err != nil {
			return err
		}

		query = redirectURI.Query()
	}

	query.Set("session", session)

	user, err := LocateUserInfo(assertion, sp.attributeRules)
	if err != nil {
		return err
	}

	query.Set("identity", user.Identity)
	query.Set("fullname", user.Fullname)

	if user.Identity == "" {
		return errors.New("couldn't find user identity in the SAML assertion")
	}

	samlAssertion := auth.NewSAMLAssertion(session)

	samlAssertion.TypedSpec().Value.Data, err = json.Marshal(assertion)
	if err != nil {
		return err
	}

	samlAssertion.TypedSpec().Value.Email = user.Identity

	ctx := actor.MarkContextAsInternalActor(r.Context())

	if err = sp.state.Create(ctx, samlAssertion); err != nil {
		return err
	}

	samlLabels, err := sp.ReadLabelsFromAssertion(ctx, assertion)
	if err != nil {
		return err
	}

	if err = sp.ensureUser(ctx, user.Identity, samlLabels); err != nil {
		return err
	}

	if nameID := assertion.Subject.NameID; nameID != nil {
		data := sloSessionData{
			NameID: nameID.Value,
			Format: nameID.Format,
		}

		for _, authnStatement := range assertion.AuthnStatements {
			if authnStatement.SessionIndex != "" {
				data.SessionIndex = authnStatement.SessionIndex

				break
			}
		}

		if cookieValue, marshalErr := json.Marshal(data); marshalErr != nil {
			sp.logger.Warn("failed to marshal SLO session cookie data", zap.Error(marshalErr))
		} else {
			http.SetCookie(w, &http.Cookie{
				Name:     NameIDCookieName,
				Value:    base64.URLEncoding.EncodeToString(cookieValue),
				MaxAge:   int(sloSessionCookieTTL.Seconds()),
				Path:     "/",
				HttpOnly: true,
				Secure:   true,
				SameSite: http.SameSiteLaxMode,
			})
		}
	}

	http.Redirect(w, r, fmt.Sprintf("/authenticate?%s", query.Encode()), http.StatusSeeOther)

	return nil
}

// DeleteSession shouldn't be used by SAML provider configured for Omni.
// Session cleanup is done on the first usage or by a timeout.
func (sp *SessionProvider) DeleteSession(http.ResponseWriter, *http.Request) error {
	return errors.New("not implemented")
}

// GetSession shouldn't be used by SAML provider configured for Omni.
func (sp *SessionProvider) GetSession(*http.Request) (samlsp.Session, error) {
	return nil, errors.New("not implemented")
}

// ReadLabelsFromAssertion extract user labels from the SAML assertion attributes.
func (sp *SessionProvider) ReadLabelsFromAssertion(ctx context.Context, assertion *saml.Assertion) (map[string]string, error) {
	knownFields := map[string]string{
		// Microsoft
		"http://schemas.microsoft.com/ws/2008/06/identity/claims/role": "role",
		// Google
		"Role": "role",
		// Samltest
		"role": "role",
	}

	config, err := safe.ReaderGet[*auth.Config](ctx, sp.state, auth.NewAuthConfig().Metadata())
	if err != nil {
		return nil, err
	}

	customFields := config.TypedSpec().Value.Saml.LabelRules

	get := func(attr saml.Attribute, from map[string]string) (string, []string) {
		dest, ok := from[attr.FriendlyName]
		if !ok {
			dest, ok = from[attr.Name]

			if !ok {
				return "", nil
			}
		}

		return dest, xslices.Map(attr.Values, func(val saml.AttributeValue) string {
			return val.Value
		})
	}

	samlLabels := make(map[string]string)

	for _, attributeStatement := range assertion.AttributeStatements {
		for _, attr := range attributeStatement.Attributes {
			if len(attr.Values) == 0 {
				continue
			}

			key, values := get(attr, customFields)
			if key == "" {
				key, values = get(attr, knownFields)
			}

			if key == "" {
				continue
			}

			for _, value := range values {
				samlLabels[fmt.Sprintf("%s%s/%s", auth.SAMLLabelPrefix, key, value)] = ""
			}
		}
	}

	return samlLabels, nil
}

func (sp *SessionProvider) ensureUser(ctx context.Context, email string, samlLabels map[string]string) error {
	users, err := sp.state.List(ctx, auth.NewUser("").Metadata())
	if err != nil {
		return err
	}

	var updateOnEachLogin bool

	r := role.Admin

	if len(users.Items) > 0 {
		var samlLabelRule *auth.SAMLLabelRule

		samlLabelRule, err = sp.getSAMLLabelRule(ctx, samlLabels)
		if err != nil {
			return err
		}

		r = role.None

		if samlLabelRule != nil {
			r = role.Role(getRoleFromRule(samlLabelRule))
			updateOnEachLogin = samlLabelRule.TypedSpec().Value.UpdateOnEachLogin
		}
	}

	if err = user.Ensure(ctx, sp.state, email, r, updateOnEachLogin); err != nil {
		return err
	}

	if err = sp.updateIdentityLabels(ctx, email, samlLabels); err != nil {
		return err
	}

	return err
}

func (sp *SessionProvider) updateIdentityLabels(ctx context.Context, identity string, samlLabels map[string]string) error {
	identityPtr := auth.NewIdentity(identity).Metadata()

	_, err := safe.StateUpdateWithConflicts[*auth.Identity](ctx, sp.state, identityPtr, func(r *auth.Identity) error {
		var toDelete []string

		for _, label := range r.Metadata().Labels().Keys() {
			if !strings.HasPrefix(label, auth.SAMLLabelPrefix) {
				continue
			}

			if _, ok := samlLabels[label]; !ok {
				toDelete = append(toDelete, label)
			}
		}

		r.Metadata().Labels().Do(func(temp kvutils.TempKV) {
			for k, v := range samlLabels {
				temp.Set(k, v)
			}

			for _, k := range toDelete {
				temp.Delete(k)
			}
		})

		return nil
	})

	return err
}

// getSAMLLabelRule returns the rule with highest role found in the SAML label rules.
//
// If there is no rule matching the labels, nil is returned.
func (sp *SessionProvider) getSAMLLabelRule(ctx context.Context, samlLabels map[string]string) (*auth.SAMLLabelRule, error) {
	labelRuleList, err := safe.ReaderListAll[*auth.SAMLLabelRule](ctx, sp.state)
	if err != nil {
		return nil, err
	}

	labelRules := slices.AppendSeq(make([]*auth.SAMLLabelRule, 0, labelRuleList.Len()), labelRuleList.All())

	return MatchSAMLLabelRule(labelRules, samlLabels, sp.logger), nil
}

// MatchSAMLLabelRule returns the SAMLLabelRules with the highest access level that matches the labels.
func MatchSAMLLabelRule(samlLabelRules []*auth.SAMLLabelRule, samlLabels map[string]string, logger *zap.Logger) *auth.SAMLLabelRule {
	samlRole := role.None

	resLabels := resource.Labels{}

	for key, value := range samlLabels {
		resLabels.Set(key, value)
	}

	var rule *auth.SAMLLabelRule

	for _, labelRule := range samlLabelRules {
		selectors, selectorsErr := labels.ParseSelectors(labelRule.TypedSpec().Value.GetMatchLabels())
		if selectorsErr != nil {
			logger.Warn("skip invalid match labels on identity label rule", zap.Error(selectorsErr))

			continue
		}

		if selectors.Matches(resLabels) {
			parsedRole, parseErr := role.Parse(getRoleFromRule(labelRule))
			if parseErr != nil {
				logger.Warn("skip invalid role on identity label rule", zap.Error(parseErr))

				continue
			}

			maxRole, maxErr := role.Max(parsedRole, samlRole)
			if maxErr != nil {
				logger.Warn("skip invalid role on identity label rule", zap.Error(maxErr))

				continue
			}

			if samlRole.Compare(maxRole) == -1 {
				rule = labelRule
			}

			samlRole = maxRole
		}
	}

	return rule
}

func getRoleFromRule(r *auth.SAMLLabelRule) string {
	if r.TypedSpec().Value.AssignRole != "" {
		return r.TypedSpec().Value.AssignRole
	}

	return r.TypedSpec().Value.AssignRoleOnRegistration //nolint:staticcheck
}

const (
	IdentityAttribute  = "identity"
	FullNameAttribute  = "fullName"
	FirstNameAttribute = "firstName"
	LastNameAttribute  = "lastName"
)

// LocateUserInfo searches for user email and fullname in the ACS response.
//
//nolint:gocognit
func LocateUserInfo(assertion *saml.Assertion, extraFields map[string]string) (UserInfo, error) {
	var (
		user      UserInfo
		givenName string
		surname   string
	)

	copyFields := map[string]string{
		// samltest.sp SAML.
		"urn:oasis:names:tc:SAML:attribute:subject-id": IdentityAttribute,
		"displayName": FullNameAttribute,
		// Microsoft SAML.
		"http://schemas.xmlsoap.org/ws/2005/05/identity/claims/name": IdentityAttribute,
		"http://schemas.microsoft.com/identity/claims/displayname":   FullNameAttribute,
		// keycloak X500 mapping
		"email":     IdentityAttribute,
		"givenName": FirstNameAttribute,
		"surname":   LastNameAttribute,
		// Zitadel SAML
		"UserName": IdentityAttribute,
		"FullName": FullNameAttribute,
		// Fusion SAML
		"full_name": FullNameAttribute,
	}

	maps.Copy(copyFields, extraFields)

	// Google SAML keeps that info in Subject.
	if strings.Contains(assertion.Issuer.Value, "google.com") {
		if sub := assertion.Subject; sub != nil {
			if nameID := sub.NameID; nameID != nil {
				user.Identity = nameID.Value
				user.Fullname = nameID.Value
			}
		}
	}

	for _, attributeStatement := range assertion.AttributeStatements {
		for _, attr := range attributeStatement.Attributes {
			param, ok := copyFields[attr.FriendlyName]
			if !ok {
				param, ok = copyFields[attr.Name]

				if !ok {
					continue
				}
			}

			if len(attr.Values) == 0 {
				continue
			}

			val := attr.Values[0].Value

			var dest *string

			switch param {
			case IdentityAttribute:
				_, err := mail.ParseAddress(val)
				if err == nil {
					dest = &user.Identity
				}
			case FullNameAttribute:
				dest = &user.Fullname
			case FirstNameAttribute:
				dest = &givenName
			case LastNameAttribute:
				dest = &surname
			}

			if dest == nil || *dest != "" {
				continue
			}

			*dest = val
		}
	}

	if givenName != "" && surname != "" {
		user.Fullname = givenName + " " + surname
	}

	if user.Identity == "" {
		return UserInfo{}, errors.New("unsupported SAML IDP: cannot find user identity in the assertion")
	}

	user.Identity = strings.ToLower(user.Identity)

	return user, nil
}
