breakpilot-compliance/ai-compliance-sdk/internal/sso/store.go

package sso

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Store handles SSO configuration and user data persistence.
type Store struct {
	pool *pgxpool.Pool
}

// NewStore creates a new SSO store.
func NewStore(pool *pgxpool.Pool) *Store {
	return &Store{pool: pool}
}

// ============================================================================
// SSO Configuration CRUD Operations
// ============================================================================

// CreateConfig creates a new SSO configuration for a tenant.
func (s *Store) CreateConfig(ctx context.Context, tenantID uuid.UUID, req *CreateSSOConfigRequest) (*SSOConfig, error) {
	now := time.Now().UTC()

	cfg := &SSOConfig{
		ID:               uuid.New(),
		TenantID:         tenantID,
		ProviderType:     req.ProviderType,
		Name:             req.Name,
		Enabled:          req.Enabled,
		OIDCIssuerURL:    req.OIDCIssuerURL,
		OIDCClientID:     req.OIDCClientID,
		OIDCClientSecret: req.OIDCClientSecret,
		OIDCRedirectURI:  req.OIDCRedirectURI,
		OIDCScopes:       req.OIDCScopes,
		RoleMapping:      req.RoleMapping,
		DefaultRoleID:    req.DefaultRoleID,
		AutoProvision:    req.AutoProvision,
		CreatedAt:        now,
		UpdatedAt:        now,
	}

	// Apply defaults
	if len(cfg.OIDCScopes) == 0 {
		cfg.OIDCScopes = []string{"openid", "profile", "email"}
	}
	if cfg.RoleMapping == nil {
		cfg.RoleMapping = map[string]string{}
	}

	roleMappingJSON, err := json.Marshal(cfg.RoleMapping)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal role_mapping: %w", err)
	}

	_, err = s.pool.Exec(ctx, `
		INSERT INTO sso_configurations (
			id, tenant_id, provider_type, name, enabled,
			oidc_issuer_url, oidc_client_id, oidc_client_secret, oidc_redirect_uri, oidc_scopes,
			saml_entity_id, saml_sso_url, saml_certificate, saml_acs_url,
			role_mapping, default_role_id, auto_provision,
			created_at, updated_at
		) VALUES (
			$1, $2, $3, $4, $5,
			$6, $7, $8, $9, $10,
			$11, $12, $13, $14,
			$15, $16, $17,
			$18, $19
		)
	`,
		cfg.ID, cfg.TenantID, string(cfg.ProviderType), cfg.Name, cfg.Enabled,
		cfg.OIDCIssuerURL, cfg.OIDCClientID, cfg.OIDCClientSecret, cfg.OIDCRedirectURI, cfg.OIDCScopes,
		cfg.SAMLEntityID, cfg.SAMLSSOURL, cfg.SAMLCertificate, cfg.SAMLACS_URL,
		roleMappingJSON, cfg.DefaultRoleID, cfg.AutoProvision,
		cfg.CreatedAt, cfg.UpdatedAt,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to insert sso configuration: %w", err)
	}

	return cfg, nil
}

// GetConfig retrieves an SSO configuration by ID and tenant.
func (s *Store) GetConfig(ctx context.Context, tenantID, configID uuid.UUID) (*SSOConfig, error) {
	var cfg SSOConfig
	var providerType string
	var roleMappingJSON []byte

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, tenant_id, provider_type, name, enabled,
			oidc_issuer_url, oidc_client_id, oidc_client_secret, oidc_redirect_uri, oidc_scopes,
			saml_entity_id, saml_sso_url, saml_certificate, saml_acs_url,
			role_mapping, default_role_id, auto_provision,
			created_at, updated_at
		FROM sso_configurations
		WHERE id = $1 AND tenant_id = $2
	`, configID, tenantID).Scan(
		&cfg.ID, &cfg.TenantID, &providerType, &cfg.Name, &cfg.Enabled,
		&cfg.OIDCIssuerURL, &cfg.OIDCClientID, &cfg.OIDCClientSecret, &cfg.OIDCRedirectURI, &cfg.OIDCScopes,
		&cfg.SAMLEntityID, &cfg.SAMLSSOURL, &cfg.SAMLCertificate, &cfg.SAMLACS_URL,
		&roleMappingJSON, &cfg.DefaultRoleID, &cfg.AutoProvision,
		&cfg.CreatedAt, &cfg.UpdatedAt,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get sso configuration: %w", err)
	}

	cfg.ProviderType = ProviderType(providerType)
	cfg.RoleMapping = unmarshalRoleMapping(roleMappingJSON)

	return &cfg, nil
}

// GetConfigByName retrieves an SSO configuration by name and tenant.
func (s *Store) GetConfigByName(ctx context.Context, tenantID uuid.UUID, name string) (*SSOConfig, error) {
	var cfg SSOConfig
	var providerType string
	var roleMappingJSON []byte

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, tenant_id, provider_type, name, enabled,
			oidc_issuer_url, oidc_client_id, oidc_client_secret, oidc_redirect_uri, oidc_scopes,
			saml_entity_id, saml_sso_url, saml_certificate, saml_acs_url,
			role_mapping, default_role_id, auto_provision,
			created_at, updated_at
		FROM sso_configurations
		WHERE tenant_id = $1 AND name = $2
	`, tenantID, name).Scan(
		&cfg.ID, &cfg.TenantID, &providerType, &cfg.Name, &cfg.Enabled,
		&cfg.OIDCIssuerURL, &cfg.OIDCClientID, &cfg.OIDCClientSecret, &cfg.OIDCRedirectURI, &cfg.OIDCScopes,
		&cfg.SAMLEntityID, &cfg.SAMLSSOURL, &cfg.SAMLCertificate, &cfg.SAMLACS_URL,
		&roleMappingJSON, &cfg.DefaultRoleID, &cfg.AutoProvision,
		&cfg.CreatedAt, &cfg.UpdatedAt,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get sso configuration by name: %w", err)
	}

	cfg.ProviderType = ProviderType(providerType)
	cfg.RoleMapping = unmarshalRoleMapping(roleMappingJSON)

	return &cfg, nil
}

// ListConfigs lists all SSO configurations for a tenant.
func (s *Store) ListConfigs(ctx context.Context, tenantID uuid.UUID) ([]SSOConfig, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT
			id, tenant_id, provider_type, name, enabled,
			oidc_issuer_url, oidc_client_id, oidc_client_secret, oidc_redirect_uri, oidc_scopes,
			saml_entity_id, saml_sso_url, saml_certificate, saml_acs_url,
			role_mapping, default_role_id, auto_provision,
			created_at, updated_at
		FROM sso_configurations
		WHERE tenant_id = $1
		ORDER BY name ASC
	`, tenantID)
	if err != nil {
		return nil, fmt.Errorf("failed to list sso configurations: %w", err)
	}
	defer rows.Close()

	var configs []SSOConfig
	for rows.Next() {
		cfg, err := scanSSOConfig(rows)
		if err != nil {
			return nil, err
		}
		configs = append(configs, *cfg)
	}

	return configs, nil
}

// UpdateConfig updates an existing SSO configuration with partial updates.
func (s *Store) UpdateConfig(ctx context.Context, tenantID, configID uuid.UUID, req *UpdateSSOConfigRequest) (*SSOConfig, error) {
	cfg, err := s.GetConfig(ctx, tenantID, configID)
	if err != nil {
		return nil, err
	}
	if cfg == nil {
		return nil, fmt.Errorf("sso configuration not found")
	}

	// Apply partial updates
	if req.Name != nil {
		cfg.Name = *req.Name
	}
	if req.Enabled != nil {
		cfg.Enabled = *req.Enabled
	}
	if req.OIDCIssuerURL != nil {
		cfg.OIDCIssuerURL = *req.OIDCIssuerURL
	}
	if req.OIDCClientID != nil {
		cfg.OIDCClientID = *req.OIDCClientID
	}
	if req.OIDCClientSecret != nil {
		cfg.OIDCClientSecret = *req.OIDCClientSecret
	}
	if req.OIDCRedirectURI != nil {
		cfg.OIDCRedirectURI = *req.OIDCRedirectURI
	}
	if req.OIDCScopes != nil {
		cfg.OIDCScopes = req.OIDCScopes
	}
	if req.RoleMapping != nil {
		cfg.RoleMapping = req.RoleMapping
	}
	if req.DefaultRoleID != nil {
		cfg.DefaultRoleID = req.DefaultRoleID
	}
	if req.AutoProvision != nil {
		cfg.AutoProvision = *req.AutoProvision
	}

	cfg.UpdatedAt = time.Now().UTC()

	roleMappingJSON, err := json.Marshal(cfg.RoleMapping)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal role_mapping: %w", err)
	}

	_, err = s.pool.Exec(ctx, `
		UPDATE sso_configurations SET
			name = $3, enabled = $4,
			oidc_issuer_url = $5, oidc_client_id = $6, oidc_client_secret = $7,
			oidc_redirect_uri = $8, oidc_scopes = $9,
			saml_entity_id = $10, saml_sso_url = $11, saml_certificate = $12, saml_acs_url = $13,
			role_mapping = $14, default_role_id = $15, auto_provision = $16,
			updated_at = $17
		WHERE id = $1 AND tenant_id = $2
	`,
		cfg.ID, cfg.TenantID,
		cfg.Name, cfg.Enabled,
		cfg.OIDCIssuerURL, cfg.OIDCClientID, cfg.OIDCClientSecret,
		cfg.OIDCRedirectURI, cfg.OIDCScopes,
		cfg.SAMLEntityID, cfg.SAMLSSOURL, cfg.SAMLCertificate, cfg.SAMLACS_URL,
		roleMappingJSON, cfg.DefaultRoleID, cfg.AutoProvision,
		cfg.UpdatedAt,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to update sso configuration: %w", err)
	}

	return cfg, nil
}

// DeleteConfig deletes an SSO configuration by ID and tenant.
func (s *Store) DeleteConfig(ctx context.Context, tenantID, configID uuid.UUID) error {
	_, err := s.pool.Exec(ctx,
		"DELETE FROM sso_configurations WHERE id = $1 AND tenant_id = $2",
		configID, tenantID,
	)
	if err != nil {
		return fmt.Errorf("failed to delete sso configuration: %w", err)
	}
	return nil
}

// GetEnabledConfig retrieves the active/enabled SSO configuration for a tenant.
func (s *Store) GetEnabledConfig(ctx context.Context, tenantID uuid.UUID) (*SSOConfig, error) {
	var cfg SSOConfig
	var providerType string
	var roleMappingJSON []byte

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, tenant_id, provider_type, name, enabled,
			oidc_issuer_url, oidc_client_id, oidc_client_secret, oidc_redirect_uri, oidc_scopes,
			saml_entity_id, saml_sso_url, saml_certificate, saml_acs_url,
			role_mapping, default_role_id, auto_provision,
			created_at, updated_at
		FROM sso_configurations
		WHERE tenant_id = $1 AND enabled = true
		LIMIT 1
	`, tenantID).Scan(
		&cfg.ID, &cfg.TenantID, &providerType, &cfg.Name, &cfg.Enabled,
		&cfg.OIDCIssuerURL, &cfg.OIDCClientID, &cfg.OIDCClientSecret, &cfg.OIDCRedirectURI, &cfg.OIDCScopes,
		&cfg.SAMLEntityID, &cfg.SAMLSSOURL, &cfg.SAMLCertificate, &cfg.SAMLACS_URL,
		&roleMappingJSON, &cfg.DefaultRoleID, &cfg.AutoProvision,
		&cfg.CreatedAt, &cfg.UpdatedAt,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get enabled sso configuration: %w", err)
	}

	cfg.ProviderType = ProviderType(providerType)
	cfg.RoleMapping = unmarshalRoleMapping(roleMappingJSON)

	return &cfg, nil
}

// ============================================================================
// SSO User Operations
// ============================================================================

// UpsertUser inserts or updates an SSO user via JIT provisioning.
// On conflict (tenant_id, sso_config_id, external_id), the user's email,
// display name, groups, and last login timestamp are updated.
func (s *Store) UpsertUser(ctx context.Context, tenantID, ssoConfigID uuid.UUID, externalID, email, displayName string, groups []string) (*SSOUser, error) {
	now := time.Now().UTC()
	id := uuid.New()

	var user SSOUser
	err := s.pool.QueryRow(ctx, `
		INSERT INTO sso_users (
			id, tenant_id, sso_config_id,
			external_id, email, display_name, groups,
			last_login, is_active,
			created_at, updated_at
		) VALUES (
			$1, $2, $3,
			$4, $5, $6, $7,
			$8, true,
			$8, $8
		)
		ON CONFLICT (tenant_id, sso_config_id, external_id) DO UPDATE SET
			email = EXCLUDED.email,
			display_name = EXCLUDED.display_name,
			groups = EXCLUDED.groups,
			last_login = EXCLUDED.last_login,
			is_active = true,
			updated_at = EXCLUDED.updated_at
		RETURNING
			id, tenant_id, sso_config_id,
			external_id, email, display_name, groups,
			last_login, is_active,
			created_at, updated_at
	`,
		id, tenantID, ssoConfigID,
		externalID, email, displayName, groups,
		now,
	).Scan(
		&user.ID, &user.TenantID, &user.SSOConfigID,
		&user.ExternalID, &user.Email, &user.DisplayName, &user.Groups,
		&user.LastLogin, &user.IsActive,
		&user.CreatedAt, &user.UpdatedAt,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to upsert sso user: %w", err)
	}

	return &user, nil
}

// GetUserByExternalID looks up an SSO user by their external identity provider ID.
func (s *Store) GetUserByExternalID(ctx context.Context, tenantID, ssoConfigID uuid.UUID, externalID string) (*SSOUser, error) {
	var user SSOUser

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, tenant_id, sso_config_id,
			external_id, email, display_name, groups,
			last_login, is_active,
			created_at, updated_at
		FROM sso_users
		WHERE tenant_id = $1 AND sso_config_id = $2 AND external_id = $3
	`, tenantID, ssoConfigID, externalID).Scan(
		&user.ID, &user.TenantID, &user.SSOConfigID,
		&user.ExternalID, &user.Email, &user.DisplayName, &user.Groups,
		&user.LastLogin, &user.IsActive,
		&user.CreatedAt, &user.UpdatedAt,
	)
	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get sso user by external id: %w", err)
	}

	return &user, nil
}

// ListUsers lists all SSO-provisioned users for a tenant.
func (s *Store) ListUsers(ctx context.Context, tenantID uuid.UUID) ([]SSOUser, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT
			id, tenant_id, sso_config_id,
			external_id, email, display_name, groups,
			last_login, is_active,
			created_at, updated_at
		FROM sso_users
		WHERE tenant_id = $1
		ORDER BY display_name ASC
	`, tenantID)
	if err != nil {
		return nil, fmt.Errorf("failed to list sso users: %w", err)
	}
	defer rows.Close()

	var users []SSOUser
	for rows.Next() {
		user, err := scanSSOUser(rows)
		if err != nil {
			return nil, err
		}
		users = append(users, *user)
	}

	return users, nil
}

// ============================================================================
// Row Scanning Helpers
// ============================================================================

// scanSSOConfig scans an SSO configuration row from pgx.Rows.
func scanSSOConfig(rows pgx.Rows) (*SSOConfig, error) {
	var cfg SSOConfig
	var providerType string
	var roleMappingJSON []byte

	err := rows.Scan(
		&cfg.ID, &cfg.TenantID, &providerType, &cfg.Name, &cfg.Enabled,
		&cfg.OIDCIssuerURL, &cfg.OIDCClientID, &cfg.OIDCClientSecret, &cfg.OIDCRedirectURI, &cfg.OIDCScopes,
		&cfg.SAMLEntityID, &cfg.SAMLSSOURL, &cfg.SAMLCertificate, &cfg.SAMLACS_URL,
		&roleMappingJSON, &cfg.DefaultRoleID, &cfg.AutoProvision,
		&cfg.CreatedAt, &cfg.UpdatedAt,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to scan sso configuration: %w", err)
	}

	cfg.ProviderType = ProviderType(providerType)
	cfg.RoleMapping = unmarshalRoleMapping(roleMappingJSON)

	return &cfg, nil
}

// scanSSOUser scans an SSO user row from pgx.Rows.
func scanSSOUser(rows pgx.Rows) (*SSOUser, error) {
	var user SSOUser

	err := rows.Scan(
		&user.ID, &user.TenantID, &user.SSOConfigID,
		&user.ExternalID, &user.Email, &user.DisplayName, &user.Groups,
		&user.LastLogin, &user.IsActive,
		&user.CreatedAt, &user.UpdatedAt,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to scan sso user: %w", err)
	}

	return &user, nil
}

// unmarshalRoleMapping safely unmarshals JSONB role_mapping bytes into a map.
func unmarshalRoleMapping(data []byte) map[string]string {
	if data == nil {
		return map[string]string{}
	}
	var m map[string]string
	if err := json.Unmarshal(data, &m); err != nil {
		return map[string]string{}
	}
	return m
}