geo/server/internal/tenant/repository/tenant_plan_repo.go

package repository

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"

	sharedconfig "github.com/geo-platform/tenant-api/internal/shared/config"
	"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
)

const (
	tenantPlanBlockedReasonRequired = "subscription_required"
	tenantPlanBlockedReasonExpired  = "trial_plan_expired"
	tenantPlanBlockedReasonInactive = "subscription_inactive"
	defaultArticleQuotaFallback     = 0
)

type TenantPlanAccess struct {
	SubscriptionID     int64
	PlanID             int64
	PlanCode           string
	PlanName           string
	PlanStatus         string
	SubscriptionStatus string
	StartAt            time.Time
	EndAt              time.Time
	Policy             sharedconfig.PlanQuotaPolicy
}

func (a *TenantPlanAccess) HasActiveAccess(now time.Time) bool {
	if a == nil {
		return false
	}
	if !strings.EqualFold(strings.TrimSpace(a.PlanStatus), "active") {
		return false
	}
	if !strings.EqualFold(strings.TrimSpace(a.SubscriptionStatus), "active") {
		return false
	}
	nowUTC := now.UTC()
	if !a.StartAt.IsZero() && a.StartAt.After(nowUTC) {
		return false
	}
	if a.EndAt.IsZero() {
		return true
	}
	return a.EndAt.After(nowUTC)
}

func (a *TenantPlanAccess) BlockedReason(now time.Time) string {
	if a == nil {
		return tenantPlanBlockedReasonRequired
	}
	if a.HasActiveAccess(now) {
		return ""
	}
	if strings.EqualFold(strings.TrimSpace(a.PlanCode), "free") && !a.EndAt.IsZero() && !a.EndAt.After(now.UTC()) {
		return tenantPlanBlockedReasonExpired
	}
	return tenantPlanBlockedReasonInactive
}

func (a *TenantPlanAccess) ArticleGenerationLimit() int {
	if a == nil {
		return defaultArticleQuotaFallback
	}
	return maxInt(a.Policy.ArticleGeneration, 0)
}

func (a *TenantPlanAccess) BrandLimit() int {
	if a == nil {
		return 0
	}
	return maxInt(a.Policy.BrandLimit, 0)
}

func (a *TenantPlanAccess) ImageStorageBytes() int64 {
	if a == nil {
		return 0
	}
	return maxInt64(a.Policy.ImageStorageBytes, 0)
}

func (a *TenantPlanAccess) ArticleQuotaCycle() string {
	if a == nil {
		return sharedconfig.ArticleQuotaCycleLifetime
	}
	switch strings.ToLower(strings.TrimSpace(a.Policy.ArticleQuotaCycle)) {
	case sharedconfig.ArticleQuotaCycleMonthly:
		return sharedconfig.ArticleQuotaCycleMonthly
	default:
		return sharedconfig.ArticleQuotaCycleLifetime
	}
}

func (a *TenantPlanAccess) AIPointsMonthlyLimit() int {
	if a == nil {
		return 0
	}
	return maxInt(a.Policy.AIPointsMonthly, 0)
}

func (a *TenantPlanAccess) AIPointBaseChars() int {
	if a == nil || a.Policy.AIPointBaseChars <= 0 {
		return 1000
	}
	return a.Policy.AIPointBaseChars
}

func (a *TenantPlanAccess) AIPointsWindow(now time.Time) (time.Time, time.Time, *time.Time) {
	if a == nil {
		return time.Time{}, time.Time{}, nil
	}

	windowStart := a.StartAt.UTC()
	windowEnd := a.EndAt.UTC()
	cycleStart, cycleEnd := monthlyCycleBounds(windowStart, now.UTC())
	if cycleStart.Before(windowStart) {
		cycleStart = windowStart
	}
	if !windowEnd.IsZero() && cycleEnd.After(windowEnd) {
		cycleEnd = windowEnd
	}
	resetAt := cycleEnd
	return cycleStart, cycleEnd, &resetAt
}

func (a *TenantPlanAccess) ArticleQuotaWindow(now time.Time) (time.Time, time.Time, *time.Time) {
	if a == nil {
		return time.Time{}, time.Time{}, nil
	}

	windowStart := a.StartAt.UTC()
	windowEnd := a.EndAt.UTC()
	if a.ArticleQuotaCycle() != sharedconfig.ArticleQuotaCycleMonthly {
		return windowStart, windowEnd, nil
	}

	cycleStart, cycleEnd := monthlyCycleBounds(windowStart, now.UTC())
	if cycleStart.Before(windowStart) {
		cycleStart = windowStart
	}
	if !windowEnd.IsZero() && cycleEnd.After(windowEnd) {
		cycleEnd = windowEnd
	}
	resetAt := cycleEnd
	return cycleStart, cycleEnd, &resetAt
}

type TenantPlanRepository interface {
	UpsertConfiguredPlans(ctx context.Context, membership sharedconfig.MembershipConfig) error
	GetTenantPlanAccess(ctx context.Context, tenantID int64, now time.Time) (*TenantPlanAccess, error)
}

type tenantPlanRepository struct {
	db generated.DBTX
}

func NewTenantPlanRepository(db generated.DBTX) TenantPlanRepository {
	return &tenantPlanRepository{db: db}
}

func (r *tenantPlanRepository) UpsertConfiguredPlans(ctx context.Context, membership sharedconfig.MembershipConfig) error {
	for _, plan := range membership.Plans {
		policyJSON, err := plan.QuotaPolicyJSON()
		if err != nil {
			return err
		}
		if _, err := r.db.Exec(ctx, `
INSERT INTO plans (plan_code, name, quota_policy_json, status)
VALUES ($1, $2, $3::jsonb, 'active')
ON CONFLICT (plan_code) WHERE deleted_at IS NULL
DO UPDATE SET
	name = EXCLUDED.name,
	quota_policy_json = EXCLUDED.quota_policy_json,
	status = 'active',
	updated_at = NOW()
`, plan.Code, plan.Name, policyJSON); err != nil {
			return err
		}
	}
	return nil
}

func (r *tenantPlanRepository) GetTenantPlanAccess(ctx context.Context, tenantID int64, now time.Time) (*TenantPlanAccess, error) {
	return loadTenantPlanAccess(ctx, r.db, tenantID, now)
}

func loadTenantPlanAccess(ctx context.Context, db generated.DBTX, tenantID int64, now time.Time) (*TenantPlanAccess, error) {
	row := db.QueryRow(ctx, `
SELECT
	s.id,
	s.plan_id,
	p.plan_code,
	p.name,
	p.status,
	s.status,
	s.start_at,
	s.end_at,
	p.quota_policy_json
FROM tenant_plan_subscriptions s
JOIN plans p ON p.id = s.plan_id
WHERE s.tenant_id = $1
  AND s.deleted_at IS NULL
  AND p.deleted_at IS NULL
ORDER BY
	CASE
		WHEN s.status = 'active' AND p.status = 'active' AND s.start_at <= $2 AND s.end_at > $2 THEN 0
		WHEN s.status = 'active' AND s.end_at <= $2 THEN 1
		ELSE 2
	END,
	s.start_at DESC,
	s.id DESC
LIMIT 1
`, tenantID, now.UTC())

	var access TenantPlanAccess
	var rawPolicy []byte
	if err := row.Scan(
		&access.SubscriptionID,
		&access.PlanID,
		&access.PlanCode,
		&access.PlanName,
		&access.PlanStatus,
		&access.SubscriptionStatus,
		&access.StartAt,
		&access.EndAt,
		&rawPolicy,
	); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}

	access.StartAt = access.StartAt.UTC()
	access.EndAt = access.EndAt.UTC()
	access.Policy = sharedconfig.PlanQuotaPolicy{}
	if len(rawPolicy) > 0 {
		if err := json.Unmarshal(rawPolicy, &access.Policy); err != nil {
			return nil, err
		}
	}
	access.Policy.ArticleQuotaCycle = normalizeQuotaCycle(access.Policy.ArticleQuotaCycle, access.PlanCode)
	return &access, nil
}

func CountEffectiveQuotaReservations(
	ctx context.Context,
	db generated.DBTX,
	tenantID int64,
	quotaType string,
	startAt, endAt time.Time,
) (int, error) {
	row := db.QueryRow(ctx, `
SELECT COALESCE(
	SUM(
		CASE
			WHEN status IN ('pending', 'confirmed') THEN reserved_amount
			ELSE 0
		END
	),
	0
)::INT
FROM quota_reservations
WHERE tenant_id = $1
  AND quota_type = $2
  AND created_at >= $3
  AND created_at < $4
`, tenantID, quotaType, startAt.UTC(), endAt.UTC())

	var used int
	if err := row.Scan(&used); err != nil {
		return 0, err
	}
	return used, nil
}

func monthlyCycleBounds(anchor, now time.Time) (time.Time, time.Time) {
	start := anchor.UTC()
	current := now.UTC()
	if current.Before(start) {
		return start, addMonthsPreserveAnchor(start, 1)
	}

	months := (current.Year()-start.Year())*12 + int(current.Month()-start.Month())
	cycleStart := addMonthsPreserveAnchor(start, months)
	for cycleStart.After(current) {
		months--
		cycleStart = addMonthsPreserveAnchor(start, months)
	}

	cycleEnd := addMonthsPreserveAnchor(start, months+1)
	for !cycleEnd.After(current) {
		months++
		cycleStart = addMonthsPreserveAnchor(start, months)
		cycleEnd = addMonthsPreserveAnchor(start, months+1)
	}

	return cycleStart, cycleEnd
}

// addMonthsPreserveAnchor advances anchor by n months while clamping the
// resulting day to the target month's last day. This avoids Go's time.AddDate
// drift where e.g. Jan 31 + 1 month becomes Mar 3 instead of Feb 28.
func addMonthsPreserveAnchor(anchor time.Time, n int) time.Time {
	if n == 0 {
		return anchor
	}
	y := anchor.Year()
	m := int(anchor.Month()) + n
	for m > 12 {
		y++
		m -= 12
	}
	for m < 1 {
		y--
		m += 12
	}
	// Last day of the target month: day 0 of the following month.
	lastDay := time.Date(y, time.Month(m)+1, 0, 0, 0, 0, 0, anchor.Location()).Day()
	day := anchor.Day()
	if day > lastDay {
		day = lastDay
	}
	return time.Date(y, time.Month(m), day, anchor.Hour(), anchor.Minute(), anchor.Second(), anchor.Nanosecond(), anchor.Location())
}

func normalizeQuotaCycle(value, planCode string) string {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case sharedconfig.ArticleQuotaCycleMonthly:
		return sharedconfig.ArticleQuotaCycleMonthly
	default:
		if strings.EqualFold(strings.TrimSpace(planCode), "free") {
			return sharedconfig.ArticleQuotaCycleLifetime
		}
		return sharedconfig.ArticleQuotaCycleMonthly
	}
}

func maxInt(value, fallback int) int {
	if value < fallback {
		return fallback
	}
	return value
}

func maxInt64(value, fallback int64) int64 {
	if value < fallback {
		return fallback
	}
	return value
}