descrybe/apps/api/internal/billing/plan_features.go

package billing

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

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

var (
	ErrUnknownFeatureKey     = errors.New("unknown feature key")
	ErrUnknownFeatureSection = errors.New("unknown feature section")
	ErrInvalidFeatureGates   = errors.New("invalid feature gates payload")
	ErrFeatureDisabled       = errors.New("feature_disabled")
)

// FeatureGatesView is the admin global master-switch snapshot.
type FeatureGatesView struct {
	Sections map[string]bool `json:"sections"`
	Features map[string]bool `json:"features"`
}

// PlanFeaturesView is the admin per-plan feature editor payload.
type PlanFeaturesView struct {
	PlanID           int64           `json:"plan_id"`
	PlanName         string          `json:"plan_name"`
	IsCustom         bool            `json:"is_custom"`
	IsLegacy         bool            `json:"is_legacy"`
	Features         map[string]bool `json:"features"`
	ResolvedFeatures map[string]bool `json:"resolved_features"`
}

// Capabilities is the tenant-resolved plan ∩ global feature matrix.
type Capabilities struct {
	PlanID           int64           `json:"plan_id,omitempty"`
	PlanName         string          `json:"plan_name"`
	IsCustom         bool            `json:"is_custom"`
	IsLegacy         bool            `json:"is_legacy"`
	HasActivePlan    bool            `json:"has_active_plan"`
	Features         map[string]bool `json:"features"`
	Sections         map[string]bool `json:"sections"`
	DisabledFeatures []string        `json:"disabled_features"`
	FeatureETag      string          `json:"feature_etag"`
	Entitlements     Entitlements    `json:"entitlements"`
}

// FeatureGatesUpdate is the PUT /api/admin/feature-gates body.
type FeatureGatesUpdate struct {
	Sections map[string]bool `json:"sections"`
	Features map[string]bool `json:"features"`
}

// PlanFeaturesUpdate is the PUT /api/admin/plans/{id}/features body.
// Features replaces the stored overrides object (sparse map).
type PlanFeaturesUpdate struct {
	Features map[string]bool `json:"features"`
}

// SectionGateUpdate is the PUT /api/admin/feature-gates/sections/{section} body.
// Enabled is required (*bool) so omitting the field cannot silently disable a section.
type SectionGateUpdate struct {
	Enabled *bool `json:"enabled"`
}

// DefaultPlanFeatures returns the expanded default matrix for a plan name.
// Legacy (A1 / is_legacy patterns) uses the image-nav allow-list — not custom all-ON.
// Custom packages (isCustom, non-legacy) default all registry keys ON.
func DefaultPlanFeatures(planName string, isCustom bool) map[string]bool {
	return DefaultPlanFeaturesEx(planName, isCustom, IsLegacyPlanName(planName))
}

// DefaultPlanFeaturesEx is DefaultPlanFeatures with an explicit is_legacy flag.
func DefaultPlanFeaturesEx(planName string, isCustom, isLegacy bool) map[string]bool {
	out := make(map[string]bool, len(FeatureCatalogKeys))
	// Custom deals get enable-all, except A1* PAYG which keeps Stores + AI off.
	if IsCustomPackage(planName, isCustom) {
		if IsLegacyPlanName(planName) {
			return A1PaygPlanFeatures()
		}
		for _, k := range FeatureCatalogKeys {
			out[k] = true
		}
		return out
	}
	if IsLegacyPlan(planName, isLegacy) {
		for _, k := range FeatureCatalogKeys {
			out[k] = LegacyFeatureAllowed(k)
		}
		return out
	}
	norm := strings.ToLower(strings.TrimSpace(planName))
	for _, k := range FeatureCatalogKeys {
		allowed := true
		switch norm {
		case "", "free":
			allowed = !freePlanFeatureOff(k)
		case "starter", "plus":
			allowed = !starterPlanFeatureOff(k)
		default:
			// Growth / Business / Scale / named public ladder: all ON except unknown.
			allowed = true
		}
		out[k] = allowed
	}
	return out
}

// PlanAllowsFeature resolves plan_allows(key) without global gates.
func PlanAllowsFeature(planName string, isCustom bool, overrides map[string]bool, key string) bool {
	return PlanAllowsFeatureEx(planName, isCustom, IsLegacyPlanName(planName), overrides, key)
}

// PlanAllowsFeatureEx is PlanAllowsFeature with an explicit is_legacy flag.
func PlanAllowsFeatureEx(planName string, isCustom, isLegacy bool, overrides map[string]bool, key string) bool {
	// A1 PAYG deny list always wins — stale stored matrices must not re-enable
	// Stores / Marketing / Integrations after seed hygiene expands the deny set.
	if IsCustomPackage(planName, isCustom) && IsLegacyPlanName(planName) && A1PaygFeatureDenied(key) {
		return false
	}
	if overrides != nil {
		if v, ok := overrides[key]; ok {
			return v
		}
	}
	if IsCustomPackage(planName, isCustom) {
		if IsLegacyPlanName(planName) {
			return !A1PaygFeatureDenied(key)
		}
		return true
	}
	if IsLegacyPlan(planName, isLegacy) {
		defaults := DefaultPlanFeaturesEx(planName, isCustom, true)
		if v, ok := defaults[key]; ok {
			return v
		}
		return false
	}
	defaults := DefaultPlanFeaturesEx(planName, false, false)
	if v, ok := defaults[key]; ok {
		return v
	}
	return false
}

// ResolveEffectiveFeatures applies plan ∩ global section ∩ global feature.
func ResolveEffectiveFeatures(planName string, isCustom bool, overrides map[string]bool, gates FeatureGatesView) (features map[string]bool, sections map[string]bool, disabled []string) {
	return ResolveEffectiveFeaturesEx(planName, isCustom, IsLegacyPlanName(planName), overrides, gates)
}

// ResolveEffectiveFeaturesEx is ResolveEffectiveFeatures with an explicit is_legacy flag.
func ResolveEffectiveFeaturesEx(planName string, isCustom, isLegacy bool, overrides map[string]bool, gates FeatureGatesView) (features map[string]bool, sections map[string]bool, disabled []string) {
	sections = make(map[string]bool, len(FeatureSections))
	for _, s := range FeatureSections {
		enabled := true
		if gates.Sections != nil {
			if v, ok := gates.Sections[s]; ok {
				enabled = v
			}
		}
		sections[s] = enabled
	}
	features = make(map[string]bool, len(FeatureCatalogKeys))
	disabled = make([]string, 0)
	for _, key := range FeatureCatalogKeys {
		allowed := PlanAllowsFeatureEx(planName, isCustom, isLegacy, overrides, key)
		sec, _ := SectionOfFeature(key)
		if !sections[sec] {
			allowed = false
		}
		if gates.Features != nil {
			if v, ok := gates.Features[key]; ok && !v {
				allowed = false
			}
		}
		features[key] = allowed
		if !allowed {
			disabled = append(disabled, key)
		}
	}
	sort.Strings(disabled)
	return features, sections, disabled
}

func featureETag(features map[string]bool) string {
	keys := make([]string, 0, len(features))
	for k, v := range features {
		if v {
			keys = append(keys, k)
		}
	}
	sort.Strings(keys)
	sum := sha256.Sum256([]byte(strings.Join(keys, "\n")))
	return "sha256:" + hex.EncodeToString(sum[:])
}

// CapabilitiesResponseETag is a strong HTTP ETag for GET /api/billing/capabilities.
// It covers the feature map plus plan identity and remaining credits so conditional
// GETs do not skip wallet updates when only credits change.
func CapabilitiesResponseETag(c Capabilities) string {
	raw := fmt.Sprintf("%s|p%d|r%d|%t|%s", c.FeatureETag, c.PlanID, c.Entitlements.RemainingCredits, c.HasActivePlan, c.PlanName)
	sum := sha256.Sum256([]byte(raw))
	return `"` + "sha256:" + hex.EncodeToString(sum[:8]) + `"`
}

func validateFeatureOverrides(features map[string]bool) error {
	if features == nil {
		return nil
	}
	for k := range features {
		if !IsKnownFeatureKey(k) {
			return fmt.Errorf("%w: %s", ErrUnknownFeatureKey, k)
		}
	}
	return nil
}

func validateGatesUpdate(sections, features map[string]bool) error {
	for s := range sections {
		if !IsKnownFeatureSection(s) {
			return fmt.Errorf("%w: %s", ErrUnknownFeatureSection, s)
		}
	}
	for k := range features {
		if !IsKnownFeatureKey(k) {
			return fmt.Errorf("%w: %s", ErrUnknownFeatureKey, k)
		}
	}
	return nil
}

func decodeFeaturesJSON(raw []byte) (map[string]bool, error) {
	if len(raw) == 0 {
		return map[string]bool{}, nil
	}
	var m map[string]bool
	if err := json.Unmarshal(raw, &m); err != nil {
		return nil, err
	}
	if m == nil {
		m = map[string]bool{}
	}
	return m, nil
}

func encodeFeaturesJSON(m map[string]bool) ([]byte, error) {
	if m == nil {
		m = map[string]bool{}
	}
	return json.Marshal(m)
}

func emptyGatesView() FeatureGatesView {
	sections := make(map[string]bool, len(FeatureSections))
	for _, s := range FeatureSections {
		sections[s] = true
	}
	return FeatureGatesView{
		Sections: sections,
		Features: map[string]bool{},
	}
}

func cloneGatesView(v FeatureGatesView) FeatureGatesView {
	out := FeatureGatesView{
		Sections: make(map[string]bool, len(v.Sections)),
		Features: make(map[string]bool, len(v.Features)),
	}
	for k, enabled := range v.Sections {
		out.Sections[k] = enabled
	}
	for k, enabled := range v.Features {
		out.Features[k] = enabled
	}
	return out
}

func (s *Service) invalidateFeatureGatesCache() {
	if s == nil {
		return
	}
	s.gatesMu.Lock()
	s.gatesCache = nil
	s.gatesCachedAt = time.Time{}
	s.gatesMu.Unlock()
}

func (s *Service) storeFeatureGatesCache(view FeatureGatesView) {
	if s == nil {
		return
	}
	copied := cloneGatesView(view)
	s.gatesMu.Lock()
	s.gatesCache = &copied
	s.gatesCachedAt = time.Now()
	s.gatesMu.Unlock()
}

// GetFeatureGates returns global section/feature master switches (missing => enabled).
func (s *Service) GetFeatureGates(ctx context.Context) (FeatureGatesView, error) {
	if s == nil || s.Pool == nil {
		return emptyGatesView(), nil
	}
	s.gatesMu.RLock()
	if s.gatesCache != nil && time.Since(s.gatesCachedAt) < featureGatesCacheTTL {
		cached := cloneGatesView(*s.gatesCache)
		s.gatesMu.RUnlock()
		return cached, nil
	}
	s.gatesMu.RUnlock()

	view, err := s.loadFeatureGates(ctx)
	if err != nil {
		return FeatureGatesView{}, err
	}
	s.storeFeatureGatesCache(view)
	return cloneGatesView(view), nil
}

func (s *Service) loadFeatureGates(ctx context.Context) (FeatureGatesView, error) {
	view := emptyGatesView()
	rows, err := s.Pool.Query(ctx, `
		SELECT gate_key, kind, enabled FROM platform_feature_gates`)
	if err != nil {
		// Table may not exist yet (migration pending).
		if isUndefinedRelation(err) {
			return view, nil
		}
		return FeatureGatesView{}, err
	}
	defer rows.Close()
	for rows.Next() {
		var key, kind string
		var enabled bool
		if err := rows.Scan(&key, &kind, &enabled); err != nil {
			return FeatureGatesView{}, err
		}
		switch kind {
		case "section":
			view.Sections[key] = enabled
		case "feature":
			view.Features[key] = enabled
		}
	}
	if err := rows.Err(); err != nil {
		return FeatureGatesView{}, err
	}
	return view, nil
}

// SetFeatureGates upserts provided section/feature gates (partial). Omitted maps are left unchanged.
func (s *Service) SetFeatureGates(ctx context.Context, sections, features map[string]bool, updatedBy *uuid.UUID) (FeatureGatesView, error) {
	if err := validateGatesUpdate(sections, features); err != nil {
		return FeatureGatesView{}, err
	}
	if s == nil || s.Pool == nil {
		return FeatureGatesView{}, errors.New("billing service unavailable")
	}
	tx, err := s.Pool.Begin(ctx)
	if err != nil {
		return FeatureGatesView{}, err
	}
	defer tx.Rollback(ctx)

	upsert := func(key, kind string, enabled bool) error {
		_, err := tx.Exec(ctx, `
			INSERT INTO platform_feature_gates (gate_key, kind, enabled, updated_at, updated_by)
			VALUES ($1, $2, $3, now(), $4)
			ON CONFLICT (gate_key) DO UPDATE SET
				kind = EXCLUDED.kind,
				enabled = EXCLUDED.enabled,
				updated_at = now(),
				updated_by = EXCLUDED.updated_by`, key, kind, enabled, updatedBy)
		return err
	}
	for k, v := range sections {
		if err := upsert(k, "section", v); err != nil {
			return FeatureGatesView{}, err
		}
	}
	for k, v := range features {
		if err := upsert(k, "feature", v); err != nil {
			return FeatureGatesView{}, err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return FeatureGatesView{}, err
	}
	s.invalidateFeatureGatesCache()
	return s.GetFeatureGates(ctx)
}

// SetSectionGate enables/disables one section for ALL plans (global master switch).
func (s *Service) SetSectionGate(ctx context.Context, section string, enabled bool, updatedBy *uuid.UUID) (FeatureGatesView, error) {
	section = strings.TrimSpace(section)
	if !IsKnownFeatureSection(section) {
		return FeatureGatesView{}, fmt.Errorf("%w: %s", ErrUnknownFeatureSection, section)
	}
	return s.SetFeatureGates(ctx, map[string]bool{section: enabled}, nil, updatedBy)
}

func (s *Service) loadPlanFeaturesRow(ctx context.Context, planID int64) (name string, isCustom bool, isLegacy bool, overrides map[string]bool, err error) {
	if s == nil || s.Pool == nil {
		return "", false, false, nil, errors.New("billing service unavailable")
	}
	var raw []byte
	err = s.Pool.QueryRow(ctx, `
		SELECT name, is_custom, COALESCE(is_legacy, false), COALESCE(features, '{}'::jsonb)
		FROM plans WHERE id = $1`, planID).Scan(&name, &isCustom, &isLegacy, &raw)
	if errors.Is(err, pgx.ErrNoRows) {
		return "", false, false, nil, ErrPlanNotFound
	}
	if err != nil {
		if isUndefinedColumn(err) {
			// Pre-migration: fall back without is_legacy and/or features.
			err = s.Pool.QueryRow(ctx, `
				SELECT name, is_custom, COALESCE(features, '{}'::jsonb)
				FROM plans WHERE id = $1`, planID).Scan(&name, &isCustom, &raw)
			if errors.Is(err, pgx.ErrNoRows) {
				return "", false, false, nil, ErrPlanNotFound
			}
			if err != nil {
				if isUndefinedColumn(err) {
					err = s.Pool.QueryRow(ctx, `SELECT name, is_custom FROM plans WHERE id = $1`, planID).
						Scan(&name, &isCustom)
					if errors.Is(err, pgx.ErrNoRows) {
						return "", false, false, nil, ErrPlanNotFound
					}
					if err != nil {
						return "", false, false, nil, err
					}
					return name, isCustom, IsLegacyPlanName(name), map[string]bool{}, nil
				}
				return "", false, false, nil, err
			}
			overrides, err = decodeFeaturesJSON(raw)
			if err != nil {
				return "", false, false, nil, err
			}
			return name, isCustom, IsLegacyPlanName(name), overrides, nil
		}
		return "", false, false, nil, err
	}
	overrides, err = decodeFeaturesJSON(raw)
	if err != nil {
		return "", false, false, nil, err
	}
	if !isLegacy {
		isLegacy = IsLegacyPlanName(name)
	}
	return name, isCustom, isLegacy, overrides, nil
}

func planFeaturesView(planID int64, name string, isCustom, isLegacy bool, overrides map[string]bool) PlanFeaturesView {
	if overrides == nil {
		overrides = map[string]bool{}
	}
	resolved := make(map[string]bool, len(FeatureCatalogKeys))
	for _, k := range FeatureCatalogKeys {
		resolved[k] = PlanAllowsFeatureEx(name, isCustom, isLegacy, overrides, k)
	}
	return PlanFeaturesView{
		PlanID:           planID,
		PlanName:         name,
		IsCustom:         isCustom,
		IsLegacy:         IsLegacyPlan(name, isLegacy),
		Features:         overrides,
		ResolvedFeatures: resolved,
	}
}

// GetPlanFeatures returns stored overrides + plan_allows resolved matrix (globals ignored).
func (s *Service) GetPlanFeatures(ctx context.Context, planID int64) (PlanFeaturesView, error) {
	name, isCustom, isLegacy, overrides, err := s.loadPlanFeaturesRow(ctx, planID)
	if err != nil {
		return PlanFeaturesView{}, err
	}
	return planFeaturesView(planID, name, isCustom, isLegacy, overrides), nil
}

// SetPlanFeatures replaces the plan's features override object.
func (s *Service) SetPlanFeatures(ctx context.Context, planID int64, features map[string]bool) (PlanFeaturesView, error) {
	if s == nil || s.Pool == nil {
		return PlanFeaturesView{}, errors.New("billing service unavailable")
	}
	if features == nil {
		features = map[string]bool{}
	}
	if err := validateFeatureOverrides(features); err != nil {
		return PlanFeaturesView{}, err
	}
	raw, err := encodeFeaturesJSON(features)
	if err != nil {
		return PlanFeaturesView{}, err
	}
	tag, err := s.Pool.Exec(ctx, `
		UPDATE plans SET features = $2::jsonb, updated_at = now() WHERE id = $1`, planID, raw)
	if err != nil {
		if isUndefinedColumn(err) {
			return PlanFeaturesView{}, errors.New("plans.features column missing — run migration 026_plan_features")
		}
		return PlanFeaturesView{}, err
	}
	if tag.RowsAffected() == 0 {
		return PlanFeaturesView{}, ErrPlanNotFound
	}
	return s.GetPlanFeatures(ctx, planID)
}

// EnableAllPlanFeatures sets every registry key to true on the plan (custom packages helper).
func (s *Service) EnableAllPlanFeatures(ctx context.Context, planID int64) (PlanFeaturesView, error) {
	return s.SetPlanFeatures(ctx, planID, AllRegistryFeatures(true))
}

// DisableAllPlanFeatures sets every registry key to false on the plan.
func (s *Service) DisableAllPlanFeatures(ctx context.Context, planID int64) (PlanFeaturesView, error) {
	return s.SetPlanFeatures(ctx, planID, AllRegistryFeatures(false))
}

// CapabilitiesForCompany returns effective features for the company's active plan ∩ globals.
func (s *Service) CapabilitiesForCompany(ctx context.Context, companyID uuid.UUID) (Capabilities, error) {
	if s == nil || s.Pool == nil {
		return Capabilities{}, errors.New("billing service unavailable")
	}
	gates, err := s.GetFeatureGates(ctx)
	if err != nil {
		return Capabilities{}, err
	}

	var planID int64
	var planName string
	var isCustom bool
	var isLegacy bool
	var monthly *int
	var isTrial bool
	var raw []byte
	hasPlan := false

	err = s.Pool.QueryRow(ctx, `
		SELECT p.id, p.name, p.is_custom, COALESCE(p.is_legacy, false), p.monthly_credits, cp.is_trial, COALESCE(p.features, '{}'::jsonb)
		FROM company_plans cp
		JOIN plans p ON p.id = cp.plan_id
		WHERE cp.company_id = $1 AND cp.is_active = true
		ORDER BY cp.created_at DESC LIMIT 1`, companyID).
		Scan(&planID, &planName, &isCustom, &isLegacy, &monthly, &isTrial, &raw)
	if err == nil {
		hasPlan = true
	} else if errors.Is(err, pgx.ErrNoRows) {
		planName = "Free"
		raw = []byte("{}")
	} else if isUndefinedColumn(err) {
		err = s.Pool.QueryRow(ctx, `
			SELECT p.id, p.name, p.is_custom, p.monthly_credits, cp.is_trial, COALESCE(p.features, '{}'::jsonb)
			FROM company_plans cp
			JOIN plans p ON p.id = cp.plan_id
			WHERE cp.company_id = $1 AND cp.is_active = true
			ORDER BY cp.created_at DESC LIMIT 1`, companyID).
			Scan(&planID, &planName, &isCustom, &monthly, &isTrial, &raw)
		if err == nil {
			hasPlan = true
			isLegacy = IsLegacyPlanName(planName)
		} else if errors.Is(err, pgx.ErrNoRows) {
			planName = "Free"
			raw = []byte("{}")
		} else if isUndefinedColumn(err) {
			err = s.Pool.QueryRow(ctx, `
				SELECT p.id, p.name, p.is_custom, p.monthly_credits, cp.is_trial
				FROM company_plans cp
				JOIN plans p ON p.id = cp.plan_id
				WHERE cp.company_id = $1 AND cp.is_active = true
				ORDER BY cp.created_at DESC LIMIT 1`, companyID).
				Scan(&planID, &planName, &isCustom, &monthly, &isTrial)
			if err == nil {
				hasPlan = true
				raw = []byte("{}")
				isLegacy = IsLegacyPlanName(planName)
			} else if errors.Is(err, pgx.ErrNoRows) {
				planName = "Free"
				raw = []byte("{}")
			} else {
				return Capabilities{}, err
			}
		} else {
			return Capabilities{}, err
		}
	} else {
		return Capabilities{}, err
	}

	overrides, err := decodeFeaturesJSON(raw)
	if err != nil {
		return Capabilities{}, err
	}

	var total, used int
	_ = s.Pool.QueryRow(ctx, `SELECT total_credits, used_credits FROM credit_balances WHERE company_id = $1`, companyID).
		Scan(&total, &used)
	remaining := RemainingCreditsClamped(total, used)
	monthlyVal := 0
	if monthly != nil {
		monthlyVal = *monthly
	}
	ent := ComputeEntitlements(planName, monthlyVal, remaining, isTrial)

	isLegacy = IsLegacyPlan(planName, isLegacy)
	features, sections, disabled := ResolveEffectiveFeaturesEx(planName, isCustom, isLegacy, overrides, gates)
	out := Capabilities{
		PlanName:         planName,
		IsCustom:         isCustom,
		IsLegacy:         isLegacy,
		HasActivePlan:    hasPlan,
		Features:         features,
		Sections:         sections,
		DisabledFeatures: disabled,
		FeatureETag:      featureETag(features),
		Entitlements:     ent,
	}
	if hasPlan {
		out.PlanID = planID
	}
	return out, nil
}

func isUndefinedRelation(err error) bool {
	if err == nil {
		return false
	}
	msg := strings.ToLower(err.Error())
	return strings.Contains(msg, "does not exist") && strings.Contains(msg, "platform_feature_gates")
}

func isUndefinedColumn(err error) bool {
	if err == nil {
		return false
	}
	msg := strings.ToLower(err.Error())
	missing := strings.Contains(msg, "does not exist") || strings.Contains(msg, "undefined column") || strings.Contains(msg, "undefined_column")
	if !missing {
		return false
	}
	// Postgres: column "x" of relation "y" does not exist — features or is_legacy pre-migration.
	return strings.Contains(msg, "column") || strings.Contains(msg, "features") || strings.Contains(msg, "is_legacy")
}