ai_site/platform/internal/invitestore/store.go

package invitestore

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"aijianzhan/platform/internal/authx"
)

type Invite struct {
	InviteID  int64      `json:"invite_id"`
	TenantID  int64      `json:"tenant_id"`
	Code      string     `json:"code"`
	Role      string     `json:"role"`
	OrgUnitID int64      `json:"org_unit_id,omitempty"`
	CreatedBy int64      `json:"created_by"`
	MaxUses   int        `json:"max_uses"`
	UsedCount int        `json:"used_count"`
	Status    string     `json:"status"`
	ExpiresAt *time.Time `json:"expires_at,omitempty"`
	CreatedAt time.Time  `json:"created_at"`
}

type CreateInput struct {
	Role      string
	OrgUnitID int64
	MaxUses   int
	ExpiresIn time.Duration // 0 = 不过期
}

type Store interface {
	Create(ctx context.Context, tenantID, createdBy int64, in CreateInput) (*Invite, error)
	List(ctx context.Context, tenantID int64) ([]Invite, error)
	GetByCode(ctx context.Context, code string) (*Invite, error)
	Consume(ctx context.Context, inviteID int64) error
	Revoke(ctx context.Context, tenantID, inviteID int64) error
}

type MemoryStore struct {
	mu   sync.Mutex
	byID map[int64]*Invite
	code map[string]int64
	seq  int64
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{byID: map[int64]*Invite{}, code: map[string]int64{}}
}

func (s *MemoryStore) Create(_ context.Context, tenantID, createdBy int64, in CreateInput) (*Invite, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	code, err := randomCode()
	if err != nil {
		return nil, err
	}
	s.seq++
	inv := &Invite{
		InviteID:  s.seq,
		TenantID:  tenantID,
		Code:      code,
		Role:      normalizeRole(in.Role),
		OrgUnitID: in.OrgUnitID,
		CreatedBy: createdBy,
		MaxUses:   maxUses(in.MaxUses),
		Status:    "active",
		CreatedAt: time.Now().UTC(),
	}
	if in.ExpiresIn > 0 {
		t := time.Now().UTC().Add(in.ExpiresIn)
		inv.ExpiresAt = &t
	}
	s.byID[inv.InviteID] = inv
	s.code[inv.Code] = inv.InviteID
	cp := *inv
	return &cp, nil
}

func (s *MemoryStore) List(_ context.Context, tenantID int64) ([]Invite, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]Invite, 0)
	for _, inv := range s.byID {
		if inv.TenantID == tenantID {
			out = append(out, *inv)
		}
	}
	return out, nil
}

func (s *MemoryStore) GetByCode(_ context.Context, code string) (*Invite, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	id, ok := s.code[strings.TrimSpace(code)]
	if !ok {
		return nil, fmt.Errorf("invite not found")
	}
	inv := s.byID[id]
	cp := *inv
	return &cp, nil
}

func (s *MemoryStore) Consume(_ context.Context, inviteID int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	inv, ok := s.byID[inviteID]
	if !ok {
		return fmt.Errorf("invite not found")
	}
	if err := usable(inv); err != nil {
		return err
	}
	inv.UsedCount++
	if inv.UsedCount >= inv.MaxUses {
		inv.Status = "exhausted"
	}
	return nil
}

func (s *MemoryStore) Revoke(_ context.Context, tenantID, inviteID int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	inv, ok := s.byID[inviteID]
	if !ok || inv.TenantID != tenantID {
		return fmt.Errorf("invite not found")
	}
	inv.Status = "revoked"
	return nil
}

type PostgresStore struct {
	DB *sql.DB
}

func NewPostgresStore(db *sql.DB) *PostgresStore {
	return &PostgresStore{DB: db}
}

func (s *PostgresStore) Create(ctx context.Context, tenantID, createdBy int64, in CreateInput) (*Invite, error) {
	code, err := randomCode()
	if err != nil {
		return nil, err
	}
	var expires any
	var expPtr *time.Time
	if in.ExpiresIn > 0 {
		t := time.Now().UTC().Add(in.ExpiresIn)
		expPtr = &t
		expires = t
	}
	inv := &Invite{
		TenantID:  tenantID,
		Code:      code,
		Role:      normalizeRole(in.Role),
		OrgUnitID: in.OrgUnitID,
		CreatedBy: createdBy,
		MaxUses:   maxUses(in.MaxUses),
		Status:    "active",
		ExpiresAt: expPtr,
	}
	err = s.DB.QueryRowContext(ctx, `
INSERT INTO platform_meta.tenant_invites(tenant_id, code, role, created_by, max_uses, expires_at, org_unit_id)
VALUES($1,$2,$3,$4,$5,$6,$7)
RETURNING invite_id, created_at`,
		tenantID, code, inv.Role, createdBy, inv.MaxUses, expires, in.OrgUnitID,
	).Scan(&inv.InviteID, &inv.CreatedAt)
	return inv, err
}

func (s *PostgresStore) List(ctx context.Context, tenantID int64) ([]Invite, error) {
	rows, err := s.DB.QueryContext(ctx, `
SELECT invite_id, tenant_id, code, role, created_by, max_uses, used_count, status, expires_at, created_at, COALESCE(org_unit_id,0)
FROM platform_meta.tenant_invites WHERE tenant_id=$1 ORDER BY invite_id DESC`, tenantID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := make([]Invite, 0)
	for rows.Next() {
		var inv Invite
		var exp sql.NullTime
		if err := rows.Scan(&inv.InviteID, &inv.TenantID, &inv.Code, &inv.Role, &inv.CreatedBy,
			&inv.MaxUses, &inv.UsedCount, &inv.Status, &exp, &inv.CreatedAt, &inv.OrgUnitID); err != nil {
			return nil, err
		}
		if exp.Valid {
			t := exp.Time.UTC()
			inv.ExpiresAt = &t
		}
		out = append(out, inv)
	}
	return out, rows.Err()
}

func (s *PostgresStore) GetByCode(ctx context.Context, code string) (*Invite, error) {
	var inv Invite
	var exp sql.NullTime
	err := s.DB.QueryRowContext(ctx, `
SELECT invite_id, tenant_id, code, role, created_by, max_uses, used_count, status, expires_at, created_at, COALESCE(org_unit_id,0)
FROM platform_meta.tenant_invites WHERE code=$1`, strings.TrimSpace(code),
	).Scan(&inv.InviteID, &inv.TenantID, &inv.Code, &inv.Role, &inv.CreatedBy,
		&inv.MaxUses, &inv.UsedCount, &inv.Status, &exp, &inv.CreatedAt, &inv.OrgUnitID)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, fmt.Errorf("invite not found")
	}
	if err != nil {
		return nil, err
	}
	if exp.Valid {
		t := exp.Time.UTC()
		inv.ExpiresAt = &t
	}
	return &inv, nil
}

func (s *PostgresStore) Consume(ctx context.Context, inviteID int64) error {
	res, err := s.DB.ExecContext(ctx, `
UPDATE platform_meta.tenant_invites
SET used_count = used_count + 1,
    status = CASE WHEN used_count + 1 >= max_uses THEN 'exhausted' ELSE status END
WHERE invite_id=$1 AND status='active'
  AND (expires_at IS NULL OR expires_at > now())
  AND used_count < max_uses`, inviteID)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return fmt.Errorf("invite not usable")
	}
	return nil
}

func (s *PostgresStore) Revoke(ctx context.Context, tenantID, inviteID int64) error {
	res, err := s.DB.ExecContext(ctx, `
UPDATE platform_meta.tenant_invites SET status='revoked'
WHERE invite_id=$1 AND tenant_id=$2`, inviteID, tenantID)
	if err != nil {
		return err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return fmt.Errorf("invite not found")
	}
	return nil
}

func randomCode() (string, error) {
	b := make([]byte, 12)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

func normalizeRole(role string) string {
	n := authx.NormalizeRole(role)
	if authx.ValidPlatformRole(n) {
		return n
	}
	return authx.Role编辑
}

func maxUses(n int) int {
	if n <= 0 {
		return 1
	}
	if n > 1000 {
		return 1000
	}
	return n
}

func usable(inv *Invite) error {
	if inv.Status != "active" {
		return fmt.Errorf("invite not usable")
	}
	if inv.ExpiresAt != nil && time.Now().UTC().After(*inv.ExpiresAt) {
		return fmt.Errorf("invite expired")
	}
	if inv.UsedCount >= inv.MaxUses {
		return fmt.Errorf("invite exhausted")
	}
	return nil
}

// ValidateUsable 供业务层在 Consume 前检查。
func ValidateUsable(inv *Invite) error {
	return usable(inv)
}