member-console/internal/entitlements/materialize.go

package entitlements

import (
	"context"
	"database/sql"
	"fmt"
)

// MaterializePoolEntitlements performs a full re-evaluation of all active provisions
// on a pool, producing numeric entitlements, contributions, and usage records.
// This function must be called within an existing transaction.
func MaterializePoolEntitlements(ctx context.Context, q *Queries, poolID string) error {
	// 1. Get all active provisions for this pool
	provisions, err := q.GetActivePoolProvisionsByPoolID(ctx, poolID)
	if err != nil {
		return fmt.Errorf("get active provisions: %w", err)
	}

	// 2. For each active provision, look up its entitlement set's rules
	// and compute contributions per resource key
	type contribution struct {
		provisionID      string
		contributedValue int64
		stackingPolicy   string
	}
	// resourceKey -> list of contributions
	contributionsByResource := make(map[string][]contribution)

	for _, prov := range provisions {
		rules, err := q.GetActiveRulesBySetID(ctx, prov.EntitlementSetID)
		if err != nil {
			return fmt.Errorf("get rules for entitlement set %s: %w", prov.EntitlementSetID, err)
		}

		for _, rule := range rules {
			if rule.RuleType != "limit" {
				continue // Only handle limit rules for this milestone
			}

			resourceKey := rule.ResourceKey.String
			baseValue := rule.ResourceValue.Int64

			// Apply per-unit multiplier
			value := baseValue
			if rule.ResourcePerUnit.Valid && rule.ResourcePerUnit.Bool {
				value = baseValue * int64(prov.Quantity)
			}

			policy := "additive"
			if rule.StackingPolicy.Valid {
				policy = rule.StackingPolicy.String
			}

			contributionsByResource[resourceKey] = append(
				contributionsByResource[resourceKey],
				contribution{
					provisionID:      prov.ProvisionID,
					contributedValue: value,
					stackingPolicy:   policy,
				},
			)
		}
	}

	// 3. Also process ended provisions: remove their contributions
	allProvisions, err := q.GetPoolProvisionsByPoolID(ctx, poolID)
	if err != nil {
		return fmt.Errorf("get all provisions: %w", err)
	}
	for _, prov := range allProvisions {
		if prov.Status == "ended" {
			if err := q.DeleteContributionsByProvisionID(ctx, prov.ProvisionID); err != nil {
				return fmt.Errorf("delete contributions for ended provision %s: %w", prov.ProvisionID, err)
			}
		}
	}

	// 4. For each resource key, ensure entitlement, contributions, and usage exist
	for resourceKey, contribs := range contributionsByResource {
		// Get or create the numeric entitlement
		ent, err := q.GetNumericEntitlementByPoolAndResource(ctx, GetNumericEntitlementByPoolAndResourceParams{
			PoolID:      poolID,
			ResourceKey: resourceKey,
		})
		if err == sql.ErrNoRows {
			ent, err = q.CreateNumericEntitlement(ctx, CreateNumericEntitlementParams{
				PoolID:          poolID,
				ResourceKey:     resourceKey,
				EntitlementType: "limit",
				ResourceLimit:   0,
			})
			if err != nil {
				return fmt.Errorf("create numeric entitlement for %s: %w", resourceKey, err)
			}
		} else if err != nil {
			return fmt.Errorf("get numeric entitlement for %s: %w", resourceKey, err)
		}

		// Delete existing contributions for this entitlement and re-create them
		existingContribs, err := q.ListContributionsByEntitlementID(ctx, ent.EntitlementID)
		if err != nil {
			return fmt.Errorf("list contributions for entitlement %s: %w", ent.EntitlementID, err)
		}
		for _, ec := range existingContribs {
			if err := q.DeleteContributionsByProvisionID(ctx, ec.ProvisionID); err != nil {
				return fmt.Errorf("delete contributions for provision %s: %w", ec.ProvisionID, err)
			}
		}

		// Create new contributions and compute effective limit
		var effectiveLimit int64
		for _, c := range contribs {
			_, err := q.CreateNumericEntitlementContribution(ctx, CreateNumericEntitlementContributionParams{
				EntitlementID:    ent.EntitlementID,
				ProvisionID:     c.provisionID,
				ContributedValue: c.contributedValue,
				StackingPolicy:  c.stackingPolicy,
			})
			if err != nil {
				return fmt.Errorf("create contribution: %w", err)
			}

			// Apply stacking policy (additive for this milestone)
			switch c.stackingPolicy {
			case "additive":
				effectiveLimit += c.contributedValue
			case "maximum":
				if c.contributedValue > effectiveLimit {
					effectiveLimit = c.contributedValue
				}
			}
		}

		// Update the entitlement's resource_limit
		_, err = q.UpdateNumericEntitlementLimit(ctx, UpdateNumericEntitlementLimitParams{
			EntitlementID: ent.EntitlementID,
			ResourceLimit: effectiveLimit,
		})
		if err != nil {
			return fmt.Errorf("update entitlement limit: %w", err)
		}

		// Ensure usage record exists
		_, err = q.GetUsageByPoolAndResource(ctx, GetUsageByPoolAndResourceParams{
			PoolID:      poolID,
			ResourceKey: resourceKey,
		})
		if err == sql.ErrNoRows {
			_, err = q.CreateNumericEntitlementUsage(ctx, CreateNumericEntitlementUsageParams{
				EntitlementID: ent.EntitlementID,
				PoolID:        poolID,
				ResourceKey:   resourceKey,
			})
			if err != nil {
				return fmt.Errorf("create usage record: %w", err)
			}
		} else if err != nil {
			return fmt.Errorf("get usage record: %w", err)
		}
	}

	// 5. Handle entitlements with no remaining contributions (all provisions ended)
	// For resource keys not in contributionsByResource but that have existing entitlements,
	// set their limit to 0 but keep the entitlement and usage records
	for _, prov := range allProvisions {
		if prov.Status != "ended" {
			continue
		}
		rules, err := q.GetActiveRulesBySetID(ctx, prov.EntitlementSetID)
		if err != nil {
			continue // Rules may have been deactivated
		}
		for _, rule := range rules {
			if rule.RuleType != "limit" || !rule.ResourceKey.Valid {
				continue
			}
			rk := rule.ResourceKey.String
			if _, hasActive := contributionsByResource[rk]; hasActive {
				continue // Already handled above
			}
			// This resource key has no active contributions — zero it out
			ent, err := q.GetNumericEntitlementByPoolAndResource(ctx, GetNumericEntitlementByPoolAndResourceParams{
				PoolID:      poolID,
				ResourceKey: rk,
			})
			if err != nil {
				continue // No entitlement to zero out
			}
			_, _ = q.UpdateNumericEntitlementLimit(ctx, UpdateNumericEntitlementLimitParams{
				EntitlementID: ent.EntitlementID,
				ResourceLimit: 0,
			})
		}
	}

	return nil
}