hi-server/scripts/test_invite_gift_days.go

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

const inviteGiftMarker = "codex-test-invite-gift-days"

type giftLog struct {
	Type        uint16 `json:"type"`
	OrderNo     string `json:"order_no"`
	SubscribeId int64  `json:"subscribe_id"`
	Amount      int64  `json:"amount"`
	Balance     int64  `json:"balance"`
	Remark      string `json:"remark,omitempty"`
	Timestamp   int64  `json:"timestamp"`
}

type commissionLog struct {
	Type      uint16 `json:"type"`
	Amount    int64  `json:"amount"`
	OrderNo   string `json:"order_no"`
	Timestamp int64  `json:"timestamp"`
}

type userSubscribe struct {
	ID         int64
	UserID     int64
	ExpireTime time.Time
}

func main() {
	var (
		dsn         = flag.String("dsn", "", "MySQL DSN, for example root:pass@tcp(host:3306)/ppanel?charset=utf8mb4&parseTime=true&loc=Asia%2FShanghai")
		writeDB     = flag.Bool("write-db", false, "create isolated rows, simulate invite gifts, and clean them up")
		keep        = flag.Bool("keep", false, "keep rows for manual inspection")
		cleanupOnly = flag.Bool("cleanup-only", false, "delete leftover rows created by this script and exit")
		giftDays    = flag.Int("gift-days", 3, "days to add to both invite users")
		commission  = flag.Int64("commission-percent", 10, "commission percent for commission-path simulation")
	)
	flag.Parse()

	if *dsn == "" {
		exitf("-dsn is required")
	}

	ctx := context.Background()
	db, err := sql.Open("mysql", *dsn)
	mustNoErr(err)
	defer db.Close()
	db.SetMaxIdleConns(1)
	db.SetMaxOpenConns(1)
	mustNoErr(db.PingContext(ctx))

	if *cleanupOnly {
		mustNoErr(cleanup(ctx, db))
		fmt.Println("cleanup done")
		return
	}

	if !*writeDB {
		fmt.Println("dry run only. Add -write-db to create isolated invite rows in the TEST database.")
		return
	}
	if *giftDays <= 0 {
		exitf("-gift-days must be positive")
	}

	mustNoErr(cleanup(ctx, db))
	if !*keep {
		defer func() {
			if err := cleanup(context.Background(), db); err != nil {
				fmt.Fprintf(os.Stderr, "cleanup failed: %v\n", err)
			}
		}()
	}

	planID := mustCreatePlan(ctx, db)
	runSelfInviteScenario(ctx, db, planID, *giftDays)
	runFamilyInviteScenario(ctx, db, planID, *giftDays)
	runCommissionScenario(ctx, db, planID, *giftDays, *commission)

	if *keep {
		fmt.Println("rows kept; cleanup with -cleanup-only. inviteGiftMarker:", inviteGiftMarker)
	}
}

func runSelfInviteScenario(ctx context.Context, db *sql.DB, planID int64, giftDays int) {
	refererID := mustCreateUser(ctx, db, "self-referer", 0)
	refereeID := mustCreateUser(ctx, db, "self-referee", refererID)

	baseExpire := time.Now().Add(10 * 24 * time.Hour).Truncate(time.Second)
	refererSubID := mustCreateUserSubscribe(ctx, db, refererID, planID, baseExpire)
	refereeSubID := mustCreateUserSubscribe(ctx, db, refereeID, planID, baseExpire)

	orderNo := fmt.Sprintf("%s-self-order-%d", inviteGiftMarker, time.Now().UnixNano())
	mustNoErr(simulateInviteGiftBoth(ctx, db, orderNo, refererID, refereeID, 0, giftDays))
	mustNoErr(simulateInviteGiftBoth(ctx, db, orderNo, refererID, refereeID, 0, giftDays))

	assertExpire(ctx, db, "referer", refererSubID, baseExpire, giftDays)
	assertExpire(ctx, db, "referee", refereeSubID, baseExpire, giftDays)

	logs := mustGiftLogCount(ctx, db, orderNo)
	if logs != 2 {
		exitf("gift log count mismatch after duplicate simulation: got=%d want=2", logs)
	}

	fmt.Printf("PASS self invite: referer=%d referee=%d order=%s gift_days=%d logs=%d\n", refererID, refereeID, orderNo, giftDays, logs)
}

func runFamilyInviteScenario(ctx context.Context, db *sql.DB, planID int64, giftDays int) {
	refererOwnerID := mustCreateUser(ctx, db, "family-referer-owner", 0)
	refererMemberID := mustCreateUser(ctx, db, "family-referer-member", 0)
	refereeOwnerID := mustCreateUser(ctx, db, "family-referee-owner", 0)
	refereeMemberID := mustCreateUser(ctx, db, "family-referee-member", refererMemberID)
	mustCreateFamily(ctx, db, refererOwnerID, refererMemberID)
	mustCreateFamily(ctx, db, refereeOwnerID, refereeMemberID)

	baseExpire := time.Now().Add(10 * 24 * time.Hour).Truncate(time.Second)
	refererOwnerSubID := mustCreateUserSubscribe(ctx, db, refererOwnerID, planID, baseExpire)
	refereeOwnerSubID := mustCreateUserSubscribe(ctx, db, refereeOwnerID, planID, baseExpire)
	refererMemberSubID := mustCreateUserSubscribe(ctx, db, refererMemberID, planID, baseExpire)
	refereeMemberSubID := mustCreateUserSubscribe(ctx, db, refereeMemberID, planID, baseExpire)

	orderNo := fmt.Sprintf("%s-family-order-%d", inviteGiftMarker, time.Now().UnixNano())
	mustNoErr(simulateInviteGiftBoth(ctx, db, orderNo, refererMemberID, refereeMemberID, refereeOwnerID, giftDays))
	mustNoErr(simulateInviteGiftBoth(ctx, db, orderNo, refererMemberID, refereeMemberID, refereeOwnerID, giftDays))

	assertExpire(ctx, db, "referer owner", refererOwnerSubID, baseExpire, giftDays)
	assertExpire(ctx, db, "referee owner", refereeOwnerSubID, baseExpire, giftDays)
	assertExpire(ctx, db, "referer member", refererMemberSubID, baseExpire, 0)
	assertExpire(ctx, db, "referee member", refereeMemberSubID, baseExpire, 0)

	logs := mustGiftLogCount(ctx, db, orderNo)
	if logs != 2 {
		exitf("family gift log count mismatch after duplicate simulation: got=%d want=2", logs)
	}
	fmt.Printf("PASS family invite: referer_member=%d->owner=%d referee_member=%d->owner=%d order=%s gift_days=%d logs=%d\n",
		refererMemberID, refererOwnerID, refereeMemberID, refereeOwnerID, orderNo, giftDays, logs)
}

func runCommissionScenario(ctx context.Context, db *sql.DB, planID int64, giftDays int, commissionPercent int64) {
	if commissionPercent <= 0 {
		fmt.Println("SKIP commission invite: commission-percent <= 0")
		return
	}
	const amount int64 = 599

	refererID := mustCreateUser(ctx, db, "commission-referer", 0)
	refereeID := mustCreateUser(ctx, db, "commission-referee", refererID)
	baseExpire := time.Now().Add(10 * 24 * time.Hour).Truncate(time.Second)
	refererSubID := mustCreateUserSubscribe(ctx, db, refererID, planID, baseExpire)
	refereeSubID := mustCreateUserSubscribe(ctx, db, refereeID, planID, baseExpire)

	orderNo := fmt.Sprintf("%s-commission-first-order-%d", inviteGiftMarker, time.Now().UnixNano())
	mustNoErr(simulateInviteCommission(ctx, db, orderNo, refererID, refereeID, 0, giftDays, amount, commissionPercent, true))
	mustNoErr(simulateInviteCommission(ctx, db, orderNo, refererID, refereeID, 0, giftDays, amount, commissionPercent, true))

	wantCommission := amount * commissionPercent / 100
	assertExpire(ctx, db, "commission referer", refererSubID, baseExpire, 0)
	assertExpire(ctx, db, "commission referee", refereeSubID, baseExpire, giftDays)
	assertCommission(ctx, db, refererID, wantCommission)
	assertLogCount(ctx, db, "commission first gift", 34, orderNo, 1)
	assertLogCount(ctx, db, "commission first commission", 33, orderNo, 1)

	nonFirstRefererID := mustCreateUser(ctx, db, "commission-nonfirst-referer", 0)
	nonFirstRefereeID := mustCreateUser(ctx, db, "commission-nonfirst-referee", nonFirstRefererID)
	nonFirstRefererSubID := mustCreateUserSubscribe(ctx, db, nonFirstRefererID, planID, baseExpire)
	nonFirstRefereeSubID := mustCreateUserSubscribe(ctx, db, nonFirstRefereeID, planID, baseExpire)
	nonFirstOrderNo := fmt.Sprintf("%s-commission-nonfirst-order-%d", inviteGiftMarker, time.Now().UnixNano())
	mustNoErr(simulateInviteCommission(ctx, db, nonFirstOrderNo, nonFirstRefererID, nonFirstRefereeID, 0, giftDays, amount, commissionPercent, false))
	mustNoErr(simulateInviteCommission(ctx, db, nonFirstOrderNo, nonFirstRefererID, nonFirstRefereeID, 0, giftDays, amount, commissionPercent, false))

	assertExpire(ctx, db, "commission non-first referer", nonFirstRefererSubID, baseExpire, 0)
	assertExpire(ctx, db, "commission non-first referee", nonFirstRefereeSubID, baseExpire, 0)
	assertCommission(ctx, db, nonFirstRefererID, wantCommission)
	assertLogCount(ctx, db, "commission non-first gift", 34, nonFirstOrderNo, 0)
	assertLogCount(ctx, db, "commission non-first commission", 33, nonFirstOrderNo, 1)

	fmt.Printf("PASS commission invite: percent=%d first_order_commission=%d non_first_commission=%d\n",
		commissionPercent, wantCommission, wantCommission)
}

func assertExpire(ctx context.Context, db *sql.DB, label string, subID int64, before time.Time, addedDays int) {
	got := mustExpire(ctx, db, subID)
	want := before.Add(time.Duration(addedDays) * 24 * time.Hour)
	if !got.Equal(want) {
		exitf("%s expire mismatch: got=%s want=%s", label, got, want)
	}
	fmt.Printf("PASS %s subscribe=%d expire %s -> %s\n", label, subID, before.Format(time.RFC3339), got.Format(time.RFC3339))
}

func simulateInviteGiftBoth(ctx context.Context, db *sql.DB, orderNo string, refererID, refereeID, forcedRefereeOwnerID int64, days int) error {
	refereeTargetID, err := resolveGiftTargetUser(ctx, db, refereeID, forcedRefereeOwnerID)
	if err != nil {
		return fmt.Errorf("resolve referee gift target: %w", err)
	}
	refererTargetID, err := resolveGiftTargetUser(ctx, db, refererID, 0)
	if err != nil {
		return fmt.Errorf("resolve referer gift target: %w", err)
	}
	if err := grantGiftDays(ctx, db, refereeTargetID, orderNo, days); err != nil {
		return fmt.Errorf("grant referee gift: %w", err)
	}
	if err := grantGiftDays(ctx, db, refererTargetID, orderNo, days); err != nil {
		return fmt.Errorf("grant referer gift: %w", err)
	}
	return nil
}

func simulateInviteCommission(ctx context.Context, db *sql.DB, orderNo string, refererID, refereeID, forcedRefereeOwnerID int64, days int, amount int64, commissionPercent int64, isFirstOrder bool) error {
	if err := grantCommission(ctx, db, refererID, orderNo, amount, commissionPercent); err != nil {
		return fmt.Errorf("grant commission: %w", err)
	}
	if isFirstOrder {
		refereeTargetID, err := resolveGiftTargetUser(ctx, db, refereeID, forcedRefereeOwnerID)
		if err != nil {
			return fmt.Errorf("resolve referee gift target: %w", err)
		}
		if err := grantGiftDays(ctx, db, refereeTargetID, orderNo, days); err != nil {
			return fmt.Errorf("grant commission-path referee gift: %w", err)
		}
	}
	return nil
}

func resolveGiftTargetUser(ctx context.Context, db *sql.DB, userID int64, forcedOwnerID int64) (int64, error) {
	if forcedOwnerID > 0 {
		return forcedOwnerID, nil
	}
	var ownerID int64
	err := db.QueryRowContext(ctx, `
SELECT uf.owner_user_id
FROM user_family_member ufm
JOIN user_family uf ON uf.id = ufm.family_id AND uf.deleted_at IS NULL
WHERE ufm.user_id = ?
  AND ufm.deleted_at IS NULL
  AND ufm.status = 1
  AND ufm.role = 2
  AND uf.status = 1
ORDER BY ufm.role
LIMIT 1`, userID).Scan(&ownerID)
	if err == sql.ErrNoRows {
		return userID, nil
	}
	if err != nil {
		return 0, err
	}
	if ownerID > 0 && ownerID != userID {
		return ownerID, nil
	}
	return userID, nil
}

func grantCommission(ctx context.Context, db *sql.DB, refererID int64, orderNo string, amount int64, commissionPercent int64) error {
	var existing int64
	err := db.QueryRowContext(ctx,
		"SELECT COUNT(*) FROM system_logs WHERE type = 33 AND object_id = ? AND content LIKE ?",
		refererID, "%\""+orderNo+"\"%",
	).Scan(&existing)
	if err != nil {
		return err
	}
	if existing > 0 {
		return nil
	}

	commissionAmount := amount * commissionPercent / 100
	if _, err = db.ExecContext(ctx,
		"UPDATE `user` SET commission = commission + ?, updated_at = ? WHERE id = ?",
		commissionAmount, time.Now(), refererID,
	); err != nil {
		return err
	}

	content, err := json.Marshal(commissionLog{
		Type:      331,
		Amount:    commissionAmount,
		OrderNo:   orderNo,
		Timestamp: time.Now().UnixMilli(),
	})
	if err != nil {
		return err
	}
	_, err = db.ExecContext(ctx,
		"INSERT INTO system_logs (`type`, object_id, content, created_at, `date`) VALUES (33, ?, ?, ?, ?)",
		refererID, string(content), time.Now(), time.Now().Format("2006-01-02"),
	)
	return err
}

func grantGiftDays(ctx context.Context, db *sql.DB, userID int64, orderNo string, days int) error {
	var existing int64
	err := db.QueryRowContext(ctx,
		"SELECT COUNT(*) FROM system_logs WHERE type = 34 AND object_id = ? AND content LIKE ?",
		userID, "%\""+orderNo+"\"%",
	).Scan(&existing)
	if err != nil {
		return err
	}
	if existing > 0 {
		return nil
	}

	sub, err := findActiveSubscribe(ctx, db, userID)
	if err != nil {
		return err
	}
	nextExpire := sub.ExpireTime
	if !sub.ExpireTime.Equal(time.UnixMilli(0)) {
		nextExpire = sub.ExpireTime.Add(time.Duration(days) * 24 * time.Hour)
		if _, err = db.ExecContext(ctx,
			"UPDATE user_subscribe SET expire_time = ?, updated_at = ? WHERE id = ?",
			nextExpire, time.Now(), sub.ID,
		); err != nil {
			return err
		}
	}

	content, err := json.Marshal(giftLog{
		Type:        341,
		OrderNo:     orderNo,
		SubscribeId: sub.ID,
		Amount:      int64(days),
		Balance:     0,
		Remark:      "邀请赠送",
		Timestamp:   time.Now().UnixMilli(),
	})
	if err != nil {
		return err
	}
	_, err = db.ExecContext(ctx,
		"INSERT INTO system_logs (`type`, object_id, content, created_at, `date`) VALUES (34, ?, ?, ?, ?)",
		userID, string(content), time.Now(), time.Now().Format("2006-01-02"),
	)
	return err
}

func findActiveSubscribe(ctx context.Context, db *sql.DB, userID int64) (*userSubscribe, error) {
	var row userSubscribe
	err := db.QueryRowContext(ctx, `
SELECT id, user_id, expire_time
FROM user_subscribe
WHERE user_id = ?
  AND status IN (0, 1)
  AND (expire_time > ? OR expire_time = '1970-01-01 08:00:00')
ORDER BY expire_time DESC, id DESC
LIMIT 1`, userID, time.Now()).Scan(&row.ID, &row.UserID, &row.ExpireTime)
	if err != nil {
		return nil, err
	}
	return &row, nil
}

func mustCreatePlan(ctx context.Context, db *sql.DB) int64 {
	var sort int64
	mustNoErr(db.QueryRowContext(ctx, "SELECT COALESCE(MAX(sort), 0) + 1 FROM subscribe").Scan(&sort))
	res, err := db.ExecContext(ctx, `
INSERT INTO subscribe
(name, language, description, unit_price, unit_time, discount, replacement, inventory, traffic, speed_limit, device_limit, quota, new_user_only, nodes, node_tags, node_group_ids, node_group_id, traffic_limit, `+"`show`"+`, sell, sort, deduction_ratio, allow_deduction, reset_cycle, renewal_reset, show_original_price, created_at, updated_at)
VALUES (?, 'en', '', 599, 'Month', '', 0, -1, 1073741824, 0, 0, 0, false, '', '', '[]', 0, '', false, false, ?, 0, true, 0, false, true, ?, ?)`,
		inviteGiftMarker+"-plan", sort, time.Now(), time.Now())
	mustNoErr(err)
	id, err := res.LastInsertId()
	mustNoErr(err)
	return id
}

func mustCreateUser(ctx context.Context, db *sql.DB, role string, refererID int64) int64 {
	res, err := db.ExecContext(ctx, `
INSERT INTO `+"`user`"+`
(password, algo, avatar, balance, refer_code, referer_id, commission, referral_percentage, only_first_purchase, gift_amount, enable, is_admin, enable_balance_notify, enable_login_notify, enable_subscribe_notify, enable_trade_notify, rules, member_status, remark, created_at, updated_at, salt)
VALUES (?, 'default', '', 0, '', ?, 0, 0, true, 0, true, false, true, true, true, true, '', '', ?, ?, ?, 'default')`,
		inviteGiftMarker, refererID, inviteGiftMarker+"-"+role, time.Now(), time.Now())
	mustNoErr(err)
	id, err := res.LastInsertId()
	mustNoErr(err)
	_, err = db.ExecContext(ctx, "UPDATE `user` SET refer_code = ?, updated_at = ? WHERE id = ?", fmt.Sprintf("codex%d", id), time.Now(), id)
	mustNoErr(err)
	return id
}

func mustCreateFamily(ctx context.Context, db *sql.DB, ownerID, memberID int64) int64 {
	res, err := db.ExecContext(ctx, `
INSERT INTO user_family
(owner_user_id, max_members, status, created_at, updated_at)
VALUES (?, 3, 1, ?, ?)`, ownerID, time.Now(), time.Now())
	mustNoErr(err)
	familyID, err := res.LastInsertId()
	mustNoErr(err)

	now := time.Now()
	_, err = db.ExecContext(ctx, `
INSERT INTO user_family_member
(family_id, user_id, role, status, join_source, joined_at, created_at, updated_at)
VALUES
(?, ?, 1, 1, ?, ?, ?, ?),
(?, ?, 2, 1, ?, ?, ?, ?)`,
		familyID, ownerID, inviteGiftMarker, now, now, now,
		familyID, memberID, inviteGiftMarker, now, now, now)
	mustNoErr(err)
	return familyID
}

func mustCreateUserSubscribe(ctx context.Context, db *sql.DB, userID, planID int64, expire time.Time) int64 {
	token := fmt.Sprintf("%s-token-%d-%d", inviteGiftMarker, userID, time.Now().UnixNano())
	uuid := fmt.Sprintf("%08d-0000-4000-8000-%012d", userID, time.Now().UnixNano()%1_000_000_000_000)
	res, err := db.ExecContext(ctx, `
INSERT INTO user_subscribe
(user_id, order_id, subscribe_id, node_group_id, group_locked, traffic, download, upload, expired_download, expired_upload, token, uuid, status, note, created_at, updated_at, start_time, expire_time)
VALUES (?, 0, ?, 0, false, 1073741824, 0, 0, 0, 0, ?, ?, 1, ?, ?, ?, ?, ?)`,
		userID, planID, token, uuid, inviteGiftMarker, time.Now(), time.Now(), time.Now().Add(-time.Hour), expire)
	mustNoErr(err)
	id, err := res.LastInsertId()
	mustNoErr(err)
	return id
}

func mustExpire(ctx context.Context, db *sql.DB, subID int64) time.Time {
	var expire time.Time
	mustNoErr(db.QueryRowContext(ctx, "SELECT expire_time FROM user_subscribe WHERE id = ?", subID).Scan(&expire))
	return expire
}

func mustGiftLogCount(ctx context.Context, db *sql.DB, orderNo string) int64 {
	var count int64
	mustNoErr(db.QueryRowContext(ctx, "SELECT COUNT(*) FROM system_logs WHERE type = 34 AND content LIKE ?", "%"+orderNo+"%").Scan(&count))
	return count
}

func assertCommission(ctx context.Context, db *sql.DB, userID int64, want int64) {
	var got int64
	mustNoErr(db.QueryRowContext(ctx, "SELECT commission FROM `user` WHERE id = ?", userID).Scan(&got))
	if got != want {
		exitf("commission mismatch: user=%d got=%d want=%d", userID, got, want)
	}
	fmt.Printf("PASS commission user=%d amount=%d\n", userID, got)
}

func assertLogCount(ctx context.Context, db *sql.DB, label string, logType uint8, orderNo string, want int64) {
	var got int64
	mustNoErr(db.QueryRowContext(ctx, "SELECT COUNT(*) FROM system_logs WHERE type = ? AND content LIKE ?", logType, "%"+orderNo+"%").Scan(&got))
	if got != want {
		exitf("%s log count mismatch: got=%d want=%d", label, got, want)
	}
	fmt.Printf("PASS %s logs=%d\n", label, got)
}

func cleanup(ctx context.Context, db *sql.DB) error {
	stmts := []string{
		"DELETE FROM user_family_member WHERE join_source = '" + inviteGiftMarker + "'",
		"DELETE FROM user_family WHERE owner_user_id IN (SELECT id FROM `user` WHERE remark LIKE '" + inviteGiftMarker + "%')",
		"DELETE FROM system_logs WHERE type IN (33, 34) AND content LIKE '%" + inviteGiftMarker + "%'",
		"DELETE FROM user_subscribe WHERE note = '" + inviteGiftMarker + "' OR token LIKE '" + inviteGiftMarker + "%'",
		"DELETE FROM subscribe WHERE name LIKE '" + inviteGiftMarker + "%'",
		"DELETE FROM `user` WHERE remark LIKE '" + inviteGiftMarker + "%'",
	}
	for _, stmt := range stmts {
		if _, err := db.ExecContext(ctx, stmt); err != nil {
			return fmt.Errorf("%s: %w", stmt, err)
		}
	}
	return nil
}

func mustNoErr(err error) {
	if err != nil {
		exitf("%v", err)
	}
}

func exitf(format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	fmt.Fprintln(os.Stderr, "FAIL:", strings.TrimSpace(msg))
	os.Exit(1)
}