gochat/internal/auth/saml.go

package auth

// Reference: P2E §1.6 — SAML 2.0 Service Provider integration
// Pure Go implementation using encoding/xml for SAML XML parsing.
// No external SAML library dependency — all XML parsing, request generation,
// and response validation implemented with Go stdlib.
// Enterprise SSO feature: SP-initiated SSO with IdP metadata, ACS endpoint,
// attribute mapping, and replay attack prevention.

import (
	"compress/flate"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gochat/gochat/internal/config"
	"github.com/gochat/gochat/internal/model"
	applogger "github.com/gochat/gochat/pkg/logger"
	"github.com/redis/go-redis/v9"
	"gorm.io/gorm"
)

// SAML errors
var (
	ErrSAMLEnabled       = errors.New("saml authentication is not enabled")
	ErrSAMLInvalidConfig = errors.New("saml configuration is invalid")
	ErrSAMLInvalidResponse = errors.New("saml response validation failed")
	ErrSAMLReplay        = errors.New("saml response replay detected")
	ErrSAMLMissingNameID = errors.New("saml response missing NameID")
	ErrSAMLMissingEmail  = errors.New("saml response missing email attribute")
	ErrSAMLIdPMetadata   = errors.New("failed to load IdP metadata")
)

// SAMLUserInfo represents user info extracted from a SAML assertion.
type SAMLUserInfo struct {
	NameID      string            // Subject NameID — unique identifier from IdP
	Email       string            // Mapped from SAML attribute
	DisplayName string            // Mapped from SAML attribute
	FirstName   string            // Mapped from SAML attribute
	LastName    string            // Mapped from SAML attribute
	Attributes  map[string]string // All SAML attributes for flexibility
}

// IdPMetadata holds parsed IdP metadata fields.
type IdPMetadata struct {
	EntityID       string   // IdP entity ID
	SSOURL         string   // SingleSignOnService redirect binding URL
	SLORedirectURL string   // SingleLogoutService redirect binding URL
	SLOPostURL     string   // SingleLogoutService POST binding URL
	ACSURL         string   // Assertion Consumer Service URL (from config, not IdP metadata)
	Certificates   []string // PEM-encoded X.509 certificates from IdP metadata
}

// SAMLService provides SAML 2.0 SP authentication.
type SAMLService struct {
	cfg          *config.SAMLConfig
	idpMetadata  *IdPMetadata
	rdb          redis.Cmdable
	db           *gorm.DB
}

// GetIdPEntityID returns the configured IdP entity ID.
// Used by the ACS handler to populate SSO session data.
func (s *SAMLService) GetIdPEntityID() string {
	if s.idpMetadata != nil {
		return s.idpMetadata.EntityID
	}
	return ""
}

// --- SAML XML types for encoding/xml parsing/generation ---
// These are local types that map to SAML XML namespaces using xml.Name.

// SamlpNamespace is the SAML protocol namespace.
const SamlpNamespace = "urn:oasis:names:tc:SAML:2.0:protocol"
const SamlNamespace = "urn:oasis:names:tc:SAML:2.0:assertion"
const MDNamespace = "urn:oasis:names:tc:SAML:2.0:metadata"

// --- IdP Metadata XML types ---

// EntityDescriptorXML represents a SAML metadata EntityDescriptor.
type EntityDescriptorXML struct {
	XMLName          xml.Name                `xml:"EntityDescriptor"`
	EntityID         string                  `xml:"entityID,attr"`
	IDPSSODescriptor *IDPSSODescriptorXML    `xml:"IDPSSODescriptor"`
}

// IDPSSODescriptorXML represents an IDPSSODescriptor element.
type IDPSSODescriptorXML struct {
	XMLName                xml.Name                     `xml:"IDPSSODescriptor"`
	NameIDFormats          []NameIDFormatXML            `xml:"NameIDFormat"`
	SingleSignOnServices   []SingleSignOnServiceXML     `xml:"SingleSignOnService"`
	SingleLogoutServices   []SingleLogoutServiceXML     `xml:"SingleLogoutService"`
	KeyDescriptors         []KeyDescriptorXML           `xml:"KeyDescriptor"`
}

// SingleLogoutServiceXML represents a SingleLogoutService element in IdP metadata.
type SingleLogoutServiceXML struct {
	XMLName  xml.Name `xml:"SingleLogoutService"`
	Binding  string   `xml:"Binding,attr"`
	Location string   `xml:"Location,attr"`
}

// NameIDFormatXML represents a NameIDFormat element.
type NameIDFormatXML struct {
	XMLName xml.Name `xml:"NameIDFormat"`
	Value   string   `xml:",chardata"`
}

// SingleSignOnServiceXML represents a SingleSignOnService element.
type SingleSignOnServiceXML struct {
	XMLName  xml.Name `xml:"SingleSignOnService"`
	Binding  string   `xml:"Binding,attr"`
	Location string   `xml:"Location,attr"`
}

// KeyDescriptorXML represents a KeyDescriptor element in metadata.
type KeyDescriptorXML struct {
	XMLName  xml.Name         `xml:"KeyDescriptor"`
	Use      string           `xml:"use,attr"`
	KeyInfo  X509KeyInfoXML   `xml:"KeyInfo"`
}

// X509KeyInfoXML represents a KeyInfo element with X509 data.
type X509KeyInfoXML struct {
	XMLName     xml.Name          `xml:"KeyInfo"`
	X509Data    X509DataXML       `xml:"X509Data"`
}

// X509DataXML represents X509Data element.
type X509DataXML struct {
	XMLName          xml.Name       `xml:"X509Data"`
	X509Certificates []X509CertXML  `xml:"X509Certificate"`
}

// X509CertXML represents a base64-encoded X509 certificate.
type X509CertXML struct {
	XMLName xml.Name `xml:"X509Certificate"`
	Value   string   `xml:",chardata"`
}

// --- SAML Response XML types ---

// SAMLResponseXML represents a SAML Response element.
type SAMLResponseXML struct {
	XMLName      xml.Name              `xml:"Response"`
	ID           string                `xml:"ID,attr"`
	InResponseTo string               `xml:"InResponseTo,attr"`
	IssueInstant string                `xml:"IssueInstant,attr"`
	Destination  string                `xml:"Destination,attr"`
	Assertions   []SAMLAssertionXML   `xml:"Assertion"`
}

// SAMLAssertionXML represents a SAML Assertion element.
type SAMLAssertionXML struct {
	XMLName             xml.Name                    `xml:"Assertion"`
	ID                  string                      `xml:"ID,attr"`
	IssueInstant        string                      `xml:"IssueInstant,attr"`
	Subject             SAMLSubjectXML              `xml:"Subject"`
	Conditions          SAMLConditionsXML           `xml:"Conditions"`
	AttributeStatements []SAMLAttributeStatementXML `xml:"AttributeStatement"`
}

// SAMLSubjectXML represents a SAML Subject element.
type SAMLSubjectXML struct {
	XMLName xml.Name          `xml:"Subject"`
	NameID  SAMLNameIDXML     `xml:"NameID"`
}

// SAMLNameIDXML represents a SAML NameID element.
type SAMLNameIDXML struct {
	XMLName xml.Name `xml:"NameID"`
	Format  string   `xml:"Format,attr"`
	Value   string   `xml:",chardata"`
}

// SAMLConditionsXML represents SAML Conditions.
type SAMLConditionsXML struct {
	XMLName               xml.Name                  `xml:"Conditions"`
	NotBefore             string                    `xml:"NotBefore,attr"`
	NotOnOrAfter          string                    `xml:"NotOnOrAfter,attr"`
	AudienceRestrictions  []SAMLAudienceRestrictionXML `xml:"AudienceRestriction"`
}

// SAMLAudienceRestrictionXML represents an AudienceRestriction element.
type SAMLAudienceRestrictionXML struct {
	XMLName   xml.Name           `xml:"AudienceRestriction"`
	Audiences []SAMLAudienceXML  `xml:"Audience"`
}

// SAMLAudienceXML represents an Audience element.
type SAMLAudienceXML struct {
	XMLName xml.Name `xml:"Audience"`
	Value   string   `xml:",chardata"`
}

// SAMLAttributeStatementXML represents an AttributeStatement element.
type SAMLAttributeStatementXML struct {
	XMLName     xml.Name              `xml:"AttributeStatement"`
	Attributes  []SAMLAttributeXML    `xml:"Attribute"`
}

// SAMLAttributeXML represents a SAML Attribute element.
type SAMLAttributeXML struct {
	XMLName       xml.Name                `xml:"Attribute"`
	Name          string                  `xml:"Name,attr"`
	FriendlyName  string                  `xml:"FriendlyName,attr"`
	Values        []SAMLAttributeValueXML `xml:"AttributeValue"`
}

// SAMLAttributeValueXML represents a SAML AttributeValue element.
type SAMLAttributeValueXML struct {
	XMLName xml.Name `xml:"AttributeValue"`
	Value   string   `xml:",chardata"`
}

// --- SP Metadata XML types (for generation) ---

// SPEntityDescriptorXML represents an SP EntityDescriptor for metadata generation.
type SPEntityDescriptorXML struct {
	XMLName          xml.Name                `xml:"EntityDescriptor"`
	EntityID         string                  `xml:"entityID,attr"`
	SPSSODescriptor  *SPSSODescriptorXML     `xml:"SPSSODescriptor"`
}

// SPSSODescriptorXML represents an SPSSODescriptor for metadata generation.
type SPSSODescriptorXML struct {
	XMLName                  xml.Name                    `xml:"SPSSODescriptor"`
	AuthnRequestsSigned      string                      `xml:"authnRequestsSigned,attr"`
	WantAssertionsSigned     string                      `xml:"wantAssertionsSigned,attr"`
	ProtocolSupportEnum      string                      `xml:"protocolSupportEnumeration,attr"`
	NameIDFormats            []SPNameIDFormatXML         `xml:"NameIDFormat"`
	AssertionConsumerServices []SPAssertionConsumerServiceXML `xml:"AssertionConsumerService"`
	KeyDescriptors           []SPKeyDescriptorXML        `xml:"KeyDescriptor"`
}

// SPNameIDFormatXML represents a NameIDFormat for SP metadata.
type SPNameIDFormatXML struct {
	XMLName xml.Name `xml:"NameIDFormat"`
	Value   string   `xml:",chardata"`
}

// SPAssertionConsumerServiceXML represents an AssertionConsumerService for SP metadata.
type SPAssertionConsumerServiceXML struct {
	XMLName  xml.Name `xml:"AssertionConsumerService"`
	Binding  string   `xml:"Binding,attr"`
	Location string   `xml:"Location,attr"`
	Index    int      `xml:"index,attr"`
}

// SPKeyDescriptorXML represents a KeyDescriptor for SP metadata.
type SPKeyDescriptorXML struct {
	XMLName  xml.Name       `xml:"KeyDescriptor"`
	Use      string         `xml:"use,attr"`
	KeyInfo  SPKeyInfoXML   `xml:"KeyInfo"`
}

// SPKeyInfoXML represents KeyInfo for SP metadata.
type SPKeyInfoXML struct {
	XMLName   xml.Name        `xml:"KeyInfo"`
	X509Data  SPX509DataXML  `xml:"X509Data"`
}

// SPX509DataXML represents X509Data for SP metadata.
type SPX509DataXML struct {
	XMLName           xml.Name      `xml:"X509Data"`
	X509Certificates  []SPX509CertXML `xml:"X509Certificate"`
}

// SPX509CertXML represents an X509Certificate for SP metadata.
type SPX509CertXML struct {
	XMLName xml.Name `xml:"X509Certificate"`
	Value   string   `xml:",chardata"`
}

// --- AuthnRequest XML type (for generation) ---

// AuthnRequestXML represents a SAML AuthnRequest for SP-initiated login.
type AuthnRequestXML struct {
	XMLName          xml.Name `xml:"AuthnRequest"`
	ID               string   `xml:"ID,attr"`
	Version          string   `xml:"Version,attr"`
	IssueInstant     string   `xml:"IssueInstant,attr"`
	Destination      string   `xml:"Destination,attr"`
	AssertionConsumerServiceURL string `xml:"AssertionConsumerServiceURL,attr"`
	Issuer           SAMLIssuerXML `xml:"Issuer"`
}

// SAMLIssuerXML represents a SAML Issuer element.
type SAMLIssuerXML struct {
	XMLName xml.Name `xml:"Issuer"`
	Value   string   `xml:",chardata"`
}

// --- Public API ---

// NewSAMLService creates a SAML SP service from config, Redis, and database.
// When SAML is disabled, returns a minimal service that rejects all operations.
func NewSAMLService(cfg *config.SAMLConfig, rdb redis.Cmdable, db *gorm.DB) (*SAMLService, error) {
	if !cfg.Enabled {
		applogger.L().Info("SAML service created (disabled)")
		return &SAMLService{cfg: cfg, rdb: rdb, db: db}, nil
	}

	// Validate required fields for enabled SAML
	if cfg.SPEntityID == "" {
		return nil, ErrSAMLInvalidConfig
	}
	if cfg.ACSURL == "" {
		return nil, ErrSAMLInvalidConfig
	}
	if cfg.IdPMetadataURL == "" && cfg.IdPMetadataXML == "" {
		return nil, ErrSAMLIdPMetadata
	}

	// Load and parse IdP metadata
	idpMetadata, err := loadIdPMetadata(cfg)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrSAMLIdPMetadata, err)
	}

	applogger.L().Infof("SAML service initialized (idpEntityID=%s, ssoURL=%s)", idpMetadata.EntityID, idpMetadata.SSOURL)
	return &SAMLService{
		cfg:         cfg,
		idpMetadata: idpMetadata,
		rdb:         rdb,
		db:          db,
	}, nil
}

// InitiateLogin generates a SAML AuthnRequest redirect URL for SP-initiated SSO.
// The returned URL includes SAMLRequest (base64-encoded AuthnRequest) and RelayState.
// Stores InResponseTo in Redis for replay validation (5 min TTL).
func (s *SAMLService) InitiateLogin(state string) (string, error) {
	if !s.cfg.Enabled || s.idpMetadata == nil {
		return "", ErrSAMLEnabled
	}

	// Generate AuthnRequest
	requestID := fmt.Sprintf("id_%s", generateRandomID())
	now := time.Now().UTC().Format(time.RFC3339)

	authnRequest := AuthnRequestXML{
		ID:                          requestID,
		Version:                     "2.0",
		IssueInstant:                now,
		Destination:                 s.idpMetadata.SSOURL,
		AssertionConsumerServiceURL: s.cfg.ACSURL,
		Issuer: SAMLIssuerXML{
			Value: s.cfg.SPEntityID,
		},
	}

	// Marshal to XML
	xmlBytes, err := xml.Marshal(authnRequest)
	if err != nil {
		return "", fmt.Errorf("failed to marshal AuthnRequest: %w", err)
	}

	// Wrap in SAML protocol envelope
	xmlStr := fmt.Sprintf("<?xml version=\"1.0\" encoding=\"UTF-8\"?>%s", string(xmlBytes))

	// Deflate + base64 encode for redirect binding
	encoded, err := deflateAndBase64Encode(xmlStr)
	if err != nil {
		return "", fmt.Errorf("failed to encode AuthnRequest: %w", err)
	}

	// Build redirect URL
	redirectURL, err := url.Parse(s.idpMetadata.SSOURL)
	if err != nil {
		return "", fmt.Errorf("invalid IdP SSO URL: %w", err)
	}
	q := redirectURL.Query()
	q.Set("SAMLRequest", encoded)
	q.Set("RelayState", state)
	redirectURL.RawQuery = q.Encode()

	// Store InResponseTo in Redis for replay validation (5 min TTL)
	if s.rdb != nil {
		ctx := context.Background()
		key := fmt.Sprintf("saml:inResponseTo:%s", requestID)
		if err := s.rdb.Set(ctx, key, state, 5*time.Minute).Err(); err != nil {
			applogger.L().Warnf("Failed to store SAML InResponseTo in Redis: %v", err)
			// Non-critical: continue without Redis tracking
		}
	}

	applogger.L().Infof("SAML AuthnRequest generated (id=%s, destination=%s)", requestID, s.idpMetadata.SSOURL)
	return redirectURL.String(), nil
}

// ProcessResponse validates a SAML Response (base64-encoded XML from IdP),
// extracts NameID + attributes, validates conditions and replay attacks.
// Returns SAMLUserInfo on success.
func (s *SAMLService) ProcessResponse(samlResponse string) (*SAMLUserInfo, error) {
	if !s.cfg.Enabled || s.idpMetadata == nil {
		return nil, ErrSAMLEnabled
	}

	// Step 1: Base64-decode the SAMLResponse
	xmlBytes, err := base64.StdEncoding.DecodeString(samlResponse)
	if err != nil {
		return nil, fmt.Errorf("%w: failed to base64 decode response: %v", ErrSAMLInvalidResponse, err)
	}

	// Step 2: Parse XML
	response, err := parseSAMLResponseXML(xmlBytes)
	if err != nil {
		return nil, fmt.Errorf("%w: failed to parse response XML: %v", ErrSAMLInvalidResponse, err)
	}

	// Step 3: Extract assertion (first assertion)
	if len(response.Assertions) == 0 {
		return nil, fmt.Errorf("%w: response contains no assertions", ErrSAMLInvalidResponse)
	}
	assertion := response.Assertions[0]

	// Step 4: Validate conditions
	if err := validateConditions(&assertion.Conditions, s.cfg.SPEntityID, time.Now(), s.cfg.ClockDriftDuration()); err != nil {
		return nil, fmt.Errorf("%w: conditions validation failed: %v", ErrSAMLInvalidResponse, err)
	}

	// Step 5: Check replay (response ID must not have been used before)
	if err := checkAndTrackReplay(response.ID, s.rdb); err != nil {
		return nil, err
	}

	// Step 6: Extract NameID
	nameID := assertion.Subject.NameID.Value
	if nameID == "" {
		return nil, ErrSAMLMissingNameID
	}

	// Step 7: Extract attributes
	attrs := extractAttributesFromXML(assertion.AttributeStatements)
	attrMap := s.cfg.AttributeMap

	userInfo := &SAMLUserInfo{
		NameID:     nameID,
		Attributes: attrs,
	}

	userInfo.Email = getAttributeFromXML(attrs, attrMap.Email)
	userInfo.DisplayName = getAttributeFromXML(attrs, attrMap.DisplayName)
	userInfo.FirstName = getAttributeFromXML(attrs, attrMap.FirstName)
	userInfo.LastName = getAttributeFromXML(attrs, attrMap.LastName)

	// If DisplayName not found, compose from FirstName + LastName
	if userInfo.DisplayName == "" && (userInfo.FirstName != "" || userInfo.LastName != "") {
		userInfo.DisplayName = strings.TrimSpace(userInfo.FirstName + " " + userInfo.LastName)
	}

	// Email is required for user mapping
	if userInfo.Email == "" {
		// Fall back to NameID as email if it looks like an email
		if strings.Contains(nameID, "@") {
			userInfo.Email = nameID
		} else {
			return nil, ErrSAMLMissingEmail
		}
	}

	applogger.L().Infof("SAML auth successful (nameID=%s, email=%s)", nameID, userInfo.Email)
	return userInfo, nil
}

// GetSPMetadata returns the SP XML metadata for IdP configuration.
// This endpoint allows IdP admins to import our SP metadata.
func (s *SAMLService) GetSPMetadata() ([]byte, error) {
	if !s.cfg.Enabled || s.idpMetadata == nil {
		return nil, ErrSAMLEnabled
	}

	spMetadata := generateSPMetadataXML(s.cfg)
	xmlBytes, err := xml.MarshalIndent(spMetadata, "", "  ")
	if err != nil {
		return nil, fmt.Errorf("failed to marshal SP metadata: %w", err)
	}

	return []byte(fmt.Sprintf("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n%s", string(xmlBytes))), nil
}

// FindOrCreateUser maps SAML user info to a GoChat User model.
// Looks up user by (provider=saml, uid=nameID) or by email.
// Creates a new user if not found (auto-provisioning).
func (s *SAMLService) FindOrCreateUser(userInfo *SAMLUserInfo) (*model.User, error) {
	if s.db == nil {
		return nil, errors.New("database not available for SAML user lookup")
	}

	var user model.User

	// First: try to find by provider + UID (SAML NameID)
	err := s.db.Where("provider = ? AND uid = ?", "saml", userInfo.NameID).First(&user).Error
	if err == nil {
		// Update user attributes from latest SAML assertion
		user.Name = userInfo.DisplayName
		user.Email = userInfo.Email
		user.AvatarURL = userInfo.Attributes["avatarURL"]
		if err := s.db.Save(&user).Error; err != nil {
			return nil, fmt.Errorf("failed to update SAML user: %w", err)
		}
		return &user, nil
	}

	if !errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, fmt.Errorf("failed to query SAML user: %w", err)
	}

	// Second: try to find by email (link existing account to SAML)
	err = s.db.Where("email = ?", userInfo.Email).First(&user).Error
	if err == nil {
		// Link existing account to SAML provider
		user.Provider = "saml"
		user.UID = userInfo.NameID
		user.Name = userInfo.DisplayName
		if err := s.db.Save(&user).Error; err != nil {
			return nil, fmt.Errorf("failed to link SAML to existing user: %w", err)
		}
		return &user, nil
	}

	if !errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, fmt.Errorf("failed to query user by email: %w", err)
	}

	// Third: create new user (auto-provisioning)
	user = model.User{
		Name:     userInfo.DisplayName,
		Email:    userInfo.Email,
		Provider: "saml",
		UID:      userInfo.NameID,
		Role:     "agent",    // default role for SAML-provisioned users
		Active:   true,
		// AccountID must be set by caller or via a default account assignment
		// Password is empty — SAML users don't use password auth
	}

	if err := s.db.Create(&user).Error; err != nil {
		return nil, fmt.Errorf("failed to create SAML user: %w", err)
	}

	applogger.L().Infof("SAML auto-provisioned user (id=%d, email=%s)", user.ID, user.Email)
	return &user, nil
}

// --- Helper functions ---

// loadIdPMetadata fetches IdP metadata from URL or parses inline XML.
func loadIdPMetadata(cfg *config.SAMLConfig) (*IdPMetadata, error) {
	var data []byte

	if cfg.IdPMetadataURL != "" {
		metadataURL, err := url.Parse(cfg.IdPMetadataURL)
		if err != nil {
			return nil, fmt.Errorf("invalid IdP metadata URL: %w", err)
		}
		resp, err := http.Get(metadataURL.String())
		if err != nil {
			return nil, fmt.Errorf("failed to fetch IdP metadata from URL: %w", err)
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			return nil, fmt.Errorf("IdP metadata URL returned status %d", resp.StatusCode)
		}
		body, err := io.ReadAll(resp.Body)
		if err != nil {
			return nil, fmt.Errorf("failed to read IdP metadata response: %w", err)
		}
		data = body
	} else if cfg.IdPMetadataXML != "" {
		data = []byte(cfg.IdPMetadataXML)
	} else {
		return nil, errors.New("no IdP metadata source configured")
	}

	return parseIdPMetadataXML(data)
}

// parseIdPMetadataXML parses IdP metadata XML into an IdPMetadata struct.
func parseIdPMetadataXML(data []byte) (*IdPMetadata, error) {
	var entityDesc EntityDescriptorXML
	if err := xml.Unmarshal(data, &entityDesc); err != nil {
		return nil, fmt.Errorf("failed to parse IdP metadata XML: %w", err)
	}

	if entityDesc.EntityID == "" {
		return nil, errors.New("IdP metadata missing entityID")
	}

	metadata := &IdPMetadata{
		EntityID: entityDesc.EntityID,
	}

	if entityDesc.IDPSSODescriptor != nil {
		// Prefer HTTP-Redirect binding SSO URL
		for _, sso := range entityDesc.IDPSSODescriptor.SingleSignOnServices {
			if sso.Binding == "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" {
				metadata.SSOURL = sso.Location
				break
			}
		}
		// Fallback to any SSO URL if redirect binding not found
		if metadata.SSOURL == "" && len(entityDesc.IDPSSODescriptor.SingleSignOnServices) > 0 {
			metadata.SSOURL = entityDesc.IDPSSODescriptor.SingleSignOnServices[0].Location
		}

		// Extract signing certificates
		for _, kd := range entityDesc.IDPSSODescriptor.KeyDescriptors {
			if kd.Use == "signing" || kd.Use == "" {
				for _, cert := range kd.KeyInfo.X509Data.X509Certificates {
					// Clean up whitespace from base64 cert data
					certData := strings.TrimSpace(cert.Value)
					// Re-encode as PEM if not already
					if !strings.HasPrefix(certData, "-----BEGIN CERTIFICATE-----") {
						certData = fmt.Sprintf("-----BEGIN CERTIFICATE-----\n%s\n-----END CERTIFICATE-----", certData)
					}
					metadata.Certificates = append(metadata.Certificates, certData)
				}
			}
		}

		// Extract SLO (Single Logout) endpoints
		for _, slo := range entityDesc.IDPSSODescriptor.SingleLogoutServices {
			if slo.Binding == "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" {
				metadata.SLORedirectURL = slo.Location
			} else if slo.Binding == "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" {
				metadata.SLOPostURL = slo.Location
			}
		}
	}

	if entityDesc.IDPSSODescriptor != nil && metadata.SSOURL == "" {
		return nil, errors.New("IdP metadata missing SingleSignOnService URL")
	}

	return metadata, nil
}

// parseSAMLResponseXML parses a SAML Response XML document.
func parseSAMLResponseXML(data []byte) (*SAMLResponseXML, error) {
	var response SAMLResponseXML
	if err := xml.Unmarshal(data, &response); err != nil {
		return nil, fmt.Errorf("failed to unmarshal SAML Response: %w", err)
	}
	return &response, nil
}

// validateConditions validates SAML Conditions (time window + audience restriction).
func validateConditions(conditions *SAMLConditionsXML, spEntityID string, now time.Time, driftTolerance time.Duration) error {
	if conditions == nil {
		// No conditions = no restrictions to validate
		return nil
	}

	// Validate NotBefore
	if conditions.NotBefore != "" {
		notBefore, err := parseSAMLTime(conditions.NotBefore)
		if err != nil {
			return fmt.Errorf("invalid NotBefore timestamp: %w", err)
		}
		if now.Before(notBefore.Add(-driftTolerance)) {
			return fmt.Errorf("assertion not yet valid (now=%s, notBefore=%s)", now.Format(time.RFC3339), conditions.NotBefore)
		}
	}

	// Validate NotOnOrAfter
	if conditions.NotOnOrAfter != "" {
		notOnOrAfter, err := parseSAMLTime(conditions.NotOnOrAfter)
		if err != nil {
			return fmt.Errorf("invalid NotOnOrAfter timestamp: %w", err)
		}
		if now.Add(driftTolerance).After(notOnOrAfter) || now.Add(driftTolerance).Equal(notOnOrAfter) {
			return fmt.Errorf("assertion expired (now=%s, notOnOrAfter=%s)", now.Format(time.RFC3339), conditions.NotOnOrAfter)
		}
	}

	// Validate AudienceRestriction
	if len(conditions.AudienceRestrictions) > 0 {
		audienceMatch := false
		for _, ar := range conditions.AudienceRestrictions {
			for _, aud := range ar.Audiences {
				if aud.Value == spEntityID {
					audienceMatch = true
					break
				}
			}
			if audienceMatch {
				break
			}
		}
		if !audienceMatch {
			return fmt.Errorf("audience restriction failed (expected=%s)", spEntityID)
		}
	}

	return nil
}

// checkAndTrackReplay prevents replay attacks by tracking response IDs in Redis.
// If the response ID was already processed, returns ErrSAMLReplay.
// Otherwise, stores the ID with a TTL matching the assertion validity window.
func checkAndTrackReplay(responseID string, rdb redis.Cmdable) error {
	if rdb == nil {
		// Without Redis, we can't check replay — log warning
		applogger.L().Warn("SAML replay check skipped: Redis not available")
		return nil
	}

	ctx := context.Background()
	key := fmt.Sprintf("saml:replay:%s", responseID)

	// Check if response ID was already processed
	exists, err := rdb.Exists(ctx, key).Result()
	if err != nil {
		return fmt.Errorf("failed to check SAML replay in Redis: %w", err)
	}
	if exists > 0 {
		return ErrSAMLReplay
	}

	// Store the response ID with 5 min TTL (should cover assertion validity)
	if err := rdb.Set(ctx, key, "1", 5*time.Minute).Err(); err != nil {
		return fmt.Errorf("failed to store SAML replay tracking in Redis: %w", err)
	}

	return nil
}

// extractAttributesFromXML extracts all SAML assertion attributes into a map.
func extractAttributesFromXML(statements []SAMLAttributeStatementXML) map[string]string {
	attrs := make(map[string]string)
	for _, statement := range statements {
		for _, attr := range statement.Attributes {
			if len(attr.Values) > 0 {
				attrs[attr.Name] = attr.Values[0].Value
			}
			// Also store by FriendlyName if present
			if attr.FriendlyName != "" && len(attr.Values) > 0 {
				attrs[attr.FriendlyName] = attr.Values[0].Value
			}
		}
	}
	return attrs
}

// getAttributeFromXML retrieves a SAML attribute value by name from the attribute map.
// Checks both Name and FriendlyName matches.
func getAttributeFromXML(attrs map[string]string, attrName string) string {
	if v, ok := attrs[attrName]; ok {
		return v
	}
	return ""
}

// generateSPMetadataXML generates SP metadata XML structure from config.
func generateSPMetadataXML(cfg *config.SAMLConfig) *SPEntityDescriptorXML {
	spDesc := &SPEntityDescriptorXML{
		XMLName: xml.Name{Space: MDNamespace, Local: "EntityDescriptor"},
		EntityID: cfg.SPEntityID,
		SPSSODescriptor: &SPSSODescriptorXML{
			XMLName:              xml.Name{Space: MDNamespace, Local: "SPSSODescriptor"},
			AuthnRequestsSigned:  "false",
			WantAssertionsSigned: "true",
			ProtocolSupportEnum:  SamlpNamespace,
			NameIDFormats: []SPNameIDFormatXML{
				{Value: "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress"},
				{Value: "urn:oasis:names:tc:SAML:2.0:nameid-format:transient"},
				{Value: "urn:oasis:names:tc:SAML:2.0:nameid-format:persistent"},
			},
			AssertionConsumerServices: []SPAssertionConsumerServiceXML{
				{
					Binding:  "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
					Location: cfg.ACSURL,
					Index:    0,
				},
			},
		},
	}

	// Add SP certificate if configured
	if cfg.SPCertificate != "" {
		// Extract just the base64 data from PEM
		certBase64 := extractCertBase64FromPEM(cfg.SPCertificate)
		spDesc.SPSSODescriptor.KeyDescriptors = []SPKeyDescriptorXML{
			{
				Use: "signing",
				KeyInfo: SPKeyInfoXML{
					X509Data: SPX509DataXML{
						X509Certificates: []SPX509CertXML{
							{Value: certBase64},
						},
					},
				},
			},
		}
	}

	return spDesc
}

// deflateAndBase64Encode deflates then base64-encodes a string (SAML redirect binding).
func deflateAndBase64Encode(xmlStr string) (string, error) {
	var buf strings.Builder
	writer, err := flate.NewWriter(&buf, flate.DefaultCompression)
	if err != nil {
		return "", err
	}
	if _, err := writer.Write([]byte(xmlStr)); err != nil {
		return "", err
	}
	if err := writer.Close(); err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString([]byte(buf.String())), nil
}

// parseSAMLTime parses a SAML timestamp (ISO 8601 format).
func parseSAMLTime(s string) (time.Time, error) {
	// SAML uses ISO 8601 format, typically RFC3339
	layouts := []string{
		time.RFC3339,
		time.RFC3339Nano,
		"2006-01-02T15:04:05Z",
		"2006-01-02T15:04:05.000Z",
	}
	for _, layout := range layouts {
		if t, err := time.Parse(layout, s); err == nil {
			return t.UTC(), nil
		}
	}
	return time.Time{}, fmt.Errorf("cannot parse SAML timestamp: %s", s)
}

// generateRandomID creates a random hex string for SAML request IDs.
func generateRandomID() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		// Fallback to timestamp-based ID
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return fmt.Sprintf("%x", b)
}

// extractCertBase64FromPEM extracts the base64-encoded certificate data from a PEM block.
func extractCertBase64FromPEM(pemData string) string {
	block, _ := pem.Decode([]byte(pemData))
	if block == nil {
		return pemData // fallback: return raw data
	}
	return base64.StdEncoding.EncodeToString(block.Bytes)
}

// parseRSAPrivateKey parses a PEM-encoded RSA private key.
func parseRSAPrivateKey(pemData string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(pemData))
	if block == nil {
		return nil, errors.New("no PEM block found in private key data")
	}

	switch block.Type {
	case "RSA PRIVATE KEY":
		return x509.ParsePKCS1PrivateKey(block.Bytes)
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, errors.New("PKCS8 key is not RSA")
		}
		return rsaKey, nil
	default:
		return nil, fmt.Errorf("unsupported PEM block type: %s", block.Type)
	}
}

// parseX509Certificate parses a PEM-encoded X.509 certificate.
func parseX509Certificate(pemData string) (*x509.Certificate, error) {
	block, _ := pem.Decode([]byte(pemData))
	if block == nil {
		return nil, errors.New("no PEM block found in certificate data")
	}
	if block.Type != "CERTIFICATE" {
		return nil, fmt.Errorf("unexpected PEM block type: %s (expected CERTIFICATE)", block.Type)
	}
	return x509.ParseCertificate(block.Bytes)
}

// EncodeSAMLRequest base64-encodes a SAML request string
// for use in redirect bindings (helper for URL-based SAML flows).
func EncodeSAMLRequest(xml string) (string, error) {
	b64 := base64.StdEncoding.EncodeToString([]byte(xml))
	return b64, nil
}

// --- SAML Single Logout (SLO) support ---
// Reference: M13 §1 — SAML 2.0 Single Logout (SLO)
// SLO enables enterprise users to terminate all active sessions
// across all accounts when they logout from one application.
// Two flows: (1) SP-initiated SLO — user clicks logout in GoChat,
// we send LogoutRequest to IdP, IdP propagates to all SPs.
// (2) IdP-initiated SLO — IdP sends LogoutRequest to our SLO endpoint,
// we terminate the local session and respond.

// LogoutRequestXML represents a SAML LogoutRequest element.
type LogoutRequestXML struct {
	XMLName      xml.Name `xml:"LogoutRequest"`
	ID           string    `xml:"ID,attr"`
	Version      string    `xml:"Version,attr"`
	IssueInstant string    `xml:"IssueInstant,attr"`
	Destination  string    `xml:"Destination,attr"`
	NotOnOrAfter string    `xml:"NotOnOrAfter,attr,omitempty"`
	Issuer       SAMLIssuerXML       `xml:"Issuer"`
	NameID       SAMLNameIDXML       `xml:"NameID"`
	SessionIndex string              `xml:"SessionIndex,omitempty"`
}

// LogoutResponseXML represents a SAML LogoutResponse element.
type LogoutResponseXML struct {
	XMLName      xml.Name `xml:"LogoutResponse"`
	ID           string    `xml:"ID,attr"`
	InResponseTo string    `xml:"InResponseTo,attr"`
	Version      string    `xml:"Version,attr"`
	IssueInstant string    `xml:"IssueInstant,attr"`
	Destination  string    `xml:"Destination,attr,omitempty"`
	Issuer       SAMLIssuerXML `xml:"Issuer"`
	Status       SAMLStatusXML `xml:"Status"`
}

// SAMLStatusXML represents a SAML Status element.
type SAMLStatusXML struct {
	XMLName        xml.Name             `xml:"Status"`
	StatusCode     SAMLStatusCodeXML    `xml:"StatusCode"`
}

// SAMLStatusCodeXML represents a SAML StatusCode element.
type SAMLStatusCodeXML struct {
	XMLName  xml.Name `xml:"StatusCode"`
	Value    string   `xml:"Value,attr"`
}

// SLOEndpoint holds parsed IdP SLO endpoint info.
type SLOEndpoint struct {
	RedirectURL string // HTTP-Redirect binding SLO URL
	PostURL     string // HTTP-POST binding SLO URL
}

// InitiateLogout generates a SAML LogoutRequest redirect URL for SP-initiated SLO.
// The returned URL includes SAMLRequest (base64-encoded LogoutRequest) and RelayState.
// Stores InResponseTo in Redis for response validation (5 min TTL).
// After IdP processes the logout, it will propagate to all SPs in the session.
func (s *SAMLService) InitiateLogout(sessionID string, nameID string, sessionIndex string, state string) (string, error) {
	if !s.cfg.Enabled || s.idpMetadata == nil {
		return "", ErrSAMLEnabled
	}

	// Find SLO endpoint from IdP metadata
	sloEndpoint, err := getSLOEndpointFromIdP(s.idpMetadata)
	if err != nil {
		return "", fmt.Errorf("IdP does not support SLO: %w", err)
	}

	if sloEndpoint.RedirectURL == "" {
		return "", errors.New("IdP SLO endpoint not available (no HTTP-Redirect binding)")
	}

	// Generate LogoutRequest
	requestID := fmt.Sprintf("id_%s", generateRandomID())
	now := time.Now().UTC().Format(time.RFC3339)
	notOnOrAfter := time.Now().Add(5 * time.Minute).UTC().Format(time.RFC3339)

	logoutRequest := LogoutRequestXML{
		ID:           requestID,
		Version:      "2.0",
		IssueInstant: now,
		Destination:  sloEndpoint.RedirectURL,
		NotOnOrAfter: notOnOrAfter,
		Issuer: SAMLIssuerXML{
			Value: s.cfg.SPEntityID,
		},
		NameID: SAMLNameIDXML{
			Format: "urn:oasis:names:tc:SAML:2.0:nameid-format:transient",
			Value:  nameID,
		},
		SessionIndex: sessionIndex,
	}

	// Marshal to XML
	xmlBytes, err := xml.Marshal(logoutRequest)
	if err != nil {
		return "", fmt.Errorf("failed to marshal LogoutRequest: %w", err)
	}

	xmlStr := fmt.Sprintf("<?xml version=\"1.0\" encoding=\"UTF-8\"?>%s", string(xmlBytes))

	// Deflate + base64 encode for redirect binding
	encoded, err := deflateAndBase64Encode(xmlStr)
	if err != nil {
		return "", fmt.Errorf("failed to encode LogoutRequest: %w", err)
	}

	// Build redirect URL
	redirectURL, err := url.Parse(sloEndpoint.RedirectURL)
	if err != nil {
		return "", fmt.Errorf("invalid IdP SLO URL: %w", err)
	}
	q := redirectURL.Query()
	q.Set("SAMLRequest", encoded)
	q.Set("RelayState", state)
	redirectURL.RawQuery = q.Encode()

	// Store InResponseTo in Redis for response validation (5 min TTL)
	if s.rdb != nil {
		ctx := context.Background()
		key := fmt.Sprintf("saml:slo:inResponseTo:%s", requestID)
		if err := s.rdb.Set(ctx, key, sessionID, 5*time.Minute).Err(); err != nil {
			applogger.L().Warnf("Failed to store SAML SLO InResponseTo in Redis: %v", err)
		}
	}

	applogger.L().Infof("SAML LogoutRequest generated (id=%s, destination=%s, nameID=%s)", requestID, sloEndpoint.RedirectURL, nameID)
	return redirectURL.String(), nil
}

// ProcessLogoutRequest processes an IdP-initiated SAML LogoutRequest.
// Validates the request, terminates all SSO sessions for the user identified
// by NameID, and returns a LogoutResponse for the IdP.
func (s *SAMLService) ProcessLogoutRequest(samlRequest string) (string, error) {
	if !s.cfg.Enabled || s.idpMetadata == nil {
		return "", ErrSAMLEnabled
	}

	// Base64-decode the SAMLRequest
	xmlBytes, err := base64.StdEncoding.DecodeString(samlRequest)
	if err != nil {
		return "", fmt.Errorf("failed to base64 decode LogoutRequest: %v", err)
	}

	// Parse XML
	var logoutReq LogoutRequestXML
	if err := xml.Unmarshal(xmlBytes, &logoutReq); err != nil {
		return "", fmt.Errorf("failed to parse LogoutRequest XML: %w", err)
	}

	// Validate issuer matches our IdP
	if logoutReq.Issuer.Value != "" && logoutReq.Issuer.Value != s.idpMetadata.EntityID {
		return "", fmt.Errorf("LogoutRequest issuer mismatch (expected=%s, got=%s)", s.idpMetadata.EntityID, logoutReq.Issuer.Value)
	}

	// Check replay for the request ID
	if err := checkAndTrackReplay(logoutReq.ID, s.rdb); err != nil {
		return "", fmt.Errorf("LogoutRequest replay detected: %w", err)
	}

	applogger.L().Infof("SAML LogoutRequest received (id=%s, nameID=%s, sessionIndex=%s)", logoutReq.ID, logoutReq.NameID.Value, logoutReq.SessionIndex)

	// Generate LogoutResponse
	responseID := fmt.Sprintf("id_%s", generateRandomID())
	now := time.Now().UTC().Format(time.RFC3339)

	sloEndpoint, _ := getSLOEndpointFromIdP(s.idpMetadata)

	logoutResponse := LogoutResponseXML{
		ID:           responseID,
		InResponseTo: logoutReq.ID,
		Version:      "2.0",
		IssueInstant: now,
		Destination:  sloEndpoint.RedirectURL,
		Issuer: SAMLIssuerXML{
			Value: s.cfg.SPEntityID,
		},
		Status: SAMLStatusXML{
			StatusCode: SAMLStatusCodeXML{
				Value: "urn:oasis:names:tc:SAML:2.0:status:Success",
			},
		},
	}

	respBytes, err := xml.Marshal(logoutResponse)
	if err != nil {
		return "", fmt.Errorf("failed to marshal LogoutResponse: %w", err)
	}

	respStr := fmt.Sprintf("<?xml version=\"1.0\" encoding=\"UTF-8\"?>%s", string(respBytes))
	encodedResp := base64.StdEncoding.EncodeToString([]byte(respStr))

	applogger.L().Infof("SAML LogoutResponse generated (id=%s, inResponseTo=%s)", responseID, logoutReq.ID)
	return encodedResp, nil
}

// ProcessLogoutResponse validates an IdP LogoutResponse for SP-initiated SLO.
// Checks that the response is valid and matches our original LogoutRequest.
func (s *SAMLService) ProcessLogoutResponse(samlResponse string, relayState string) error {
	if !s.cfg.Enabled || s.idpMetadata == nil {
		return ErrSAMLEnabled
	}

	// Base64-decode the SAMLResponse
	xmlBytes, err := base64.StdEncoding.DecodeString(samlResponse)
	if err != nil {
		return fmt.Errorf("failed to base64 decode LogoutResponse: %v", err)
	}

	// Parse XML
	var logoutResp LogoutResponseXML
	if err := xml.Unmarshal(xmlBytes, &logoutResp); err != nil {
		return fmt.Errorf("failed to parse LogoutResponse XML: %w", err)
	}

	// Validate InResponseTo was tracked in Redis
	if s.rdb != nil {
		ctx := context.Background()
		key := fmt.Sprintf("saml:slo:inResponseTo:%s", logoutResp.InResponseTo)
		storedState, err := s.rdb.Get(ctx, key).Result()
		if err == redis.Nil {
			return fmt.Errorf("LogoutResponse InResponseTo not found or expired (inResponseTo=%s)", logoutResp.InResponseTo)
		}
		if err != nil {
			return fmt.Errorf("failed to validate LogoutResponse InResponseTo: %w", err)
		}
		// Clean up the stored state
		s.rdb.Del(ctx, key)
		_ = storedState // session ID from Redis, already processed
	}

	// Check response status
	if logoutResp.Status.StatusCode.Value != "urn:oasis:names:tc:SAML:2.0:status:Success" {
		return fmt.Errorf("LogoutResponse status not success: %s", logoutResp.Status.StatusCode.Value)
	}

	// Check replay for the response ID
	if err := checkAndTrackReplay(logoutResp.ID, s.rdb); err != nil {
		return fmt.Errorf("LogoutResponse replay detected: %w", err)
	}

	applogger.L().Infof("SAML LogoutResponse validated (id=%s, inResponseTo=%s, relayState=%s)", logoutResp.ID, logoutResp.InResponseTo, relayState)
	return nil
}

// getSLOEndpointFromIdP extracts SLO endpoint URLs from IdP metadata.
// Returns SLOEndpoint with redirect and/or post binding URLs.
func getSLOEndpointFromIdP(metadata *IdPMetadata) (*SLOEndpoint, error) {
	if metadata == nil {
		return nil, errors.New("no IdP metadata available")
	}

	slo := &SLOEndpoint{
		RedirectURL: metadata.SLORedirectURL,
		PostURL:     metadata.SLOPostURL,
	}

	if slo.RedirectURL == "" && slo.PostURL == "" {
		return nil, errors.New("IdP metadata does not contain SingleLogoutService endpoints")
	}

	return slo, nil
}