go_gameshield/internal/repository/gatewayip.go

package repository

import (
    "context"
	"fmt"
	v1 "github.com/go-nunu/nunu-layout-advanced/api/v1"
	"github.com/go-nunu/nunu-layout-advanced/internal/model"
	"gorm.io/gorm"
)

type GatewayipRepository interface {
	GetGatewayip(ctx context.Context, id int64) (*model.Gatewayip, error)
	AddGatewayip(ctx context.Context, req model.Gatewayip) error
	EditGatewayip(ctx context.Context, req model.Gatewayip) error
	DeleteGatewayip(ctx context.Context, req model.Gatewayip) error
	GetGatewayipByHostIdFirst(ctx context.Context, hostId int64) (string, error)
	GetGatewayipByHostIdAll(ctx context.Context, hostId int64) (*model.Gatewayip, error)
	UpdateGatewayipByHostId(ctx context.Context, req model.Gatewayip) error
	DeleteGatewayipByHostId(ctx context.Context, hostId int64) error
	GetIpWhereHostIdNull(ctx context.Context,req v1.GlobalLimitRequireResponse) error
	CleanIPByHostId(ctx context.Context, hostId []int64) error
	GetGatewayipOnlyIpByHostIdAll(ctx context.Context, hostId int64) ([]string, error)
}

func NewGatewayipRepository(
	repository *Repository,
) GatewayipRepository {
	return &gatewayipRepository{
		Repository: repository,
	}
}

type gatewayipRepository struct {
	*Repository
}

func (r *gatewayipRepository) GetGatewayip(ctx context.Context, id int64) (*model.Gatewayip, error) {
	var req model.Gatewayip
	return &req, r.DB(ctx).Where("id = ?", id).First(&req).Error
}

func (r *gatewayipRepository) AddGatewayip(ctx context.Context, req model.Gatewayip) error {
	return r.DB(ctx).Create(&req).Error
}

func (r *gatewayipRepository) EditGatewayip(ctx context.Context, req model.Gatewayip) error {
	return r.DB(ctx).Model(&model.Gatewayip{}).Where("id = ?", req.Id).Updates(req).Error
}

func (r *gatewayipRepository) DeleteGatewayip(ctx context.Context, req model.Gatewayip) error {
	return r.DB(ctx).Model(&model.Gatewayip{}).Where("id = ?", req.Id).Delete(req).Error
}

func (r *gatewayipRepository) GetGatewayipByHostIdFirst(ctx context.Context, hostId int64) (string, error) {
	var req string
	return req, r.DB(ctx).Model(&model.Gatewayip{}).Where("host_id = ?", hostId).Pluck("ip", &req).Error
}

func (r *gatewayipRepository) GetGatewayipByHostIdAll(ctx context.Context, hostId int64) (*model.Gatewayip, error) {
	var req model.Gatewayip
	return &req, r.DB(ctx).Where("host_id = ?", hostId).Find(&req).Error
}

func (r *gatewayipRepository) UpdateGatewayipByHostId(ctx context.Context, req model.Gatewayip) error {
	return r.DB(ctx).Where("host_id = ?", req.HostId).Updates(&req).Error
}

func (r *gatewayipRepository) DeleteGatewayipByHostId(ctx context.Context, hostId int64) error {
	return r.DB(ctx).Model(&model.Gatewayip{}).Where("host_id = ?", hostId).Delete(&model.Gatewayip{}).Error
}


func (r *gatewayipRepository) GetIpWhereHostIdNull(ctx context.Context,req v1.GlobalLimitRequireResponse) error {
	if req.IpCount <= 0 {
		return fmt.Errorf("套餐IP数量错误, 请联系客服")
	}
	if req.HostId <= 0 {
		return fmt.Errorf("主机ID错误, 请联系客服")
	}

	var count int64
	err := r.DB(ctx).Model(&model.Gatewayip{}).Where("host_id = ?", req.HostId).Count(&count).Error
	if err != nil {
		return err
	}
	if count >= int64(req.IpCount) {
		return nil
	}

	req.IpCount = int(int64(req.IpCount) - count)

	// 使用事务保证操作的原子性
	return r.DB(ctx).Transaction(func(tx *gorm.DB) error {
		var idsToAssign []uint // 只需一个切片来接收ID

		// 步骤 1: 查询所需数量的可用IP ID。使用 Limit 可以提升性能，避免捞出所有可用IP。
		err := tx.Model(&model.Gatewayip{}).
			Where("operator = ?", req.Operator).
			Where("ban_udp = ?", req.IsBanUdp).
			Where("ban_overseas = ?", req.IsBanOverseas).
			Where("node_area = ?",req.NodeArea).
			Where("host_id IS NULL OR host_id = ?", 0).
			Order("id ASC").
			Limit(req.IpCount). // 优化点：直接用Limit限制查询数量
			Pluck("id", &idsToAssign).Error

		if err != nil {
			return err // 查询出错，事务回滚
		}

		// 步骤 2: 判断实际查到的数量是否足够
		if len(idsToAssign) < req.IpCount {
			return fmt.Errorf("库存不足, 请联系客服补充") // 数量不足，返回特定错误，事务回滚
		}

		// 步骤 3: 更新这些IP的 host_id
		// 注意：因为上面已经Limit了，所以idsToAssign的长度就是我们要更新的数量
		updateResult := tx.Model(&model.Gatewayip{}).
			Where("id IN ?", idsToAssign).
			Update("host_id", req.HostId)

		if updateResult.Error != nil {
			return updateResult.Error // 更新失败，事务回滚
		}

		// (可选) 健壮性检查
		if updateResult.RowsAffected != int64(req.IpCount) {
			return fmt.Errorf("IP分配异常: 期望更新 %d 条记录, 实际更新了 %d 条", req.IpCount, updateResult.RowsAffected)
		}

		// 返回 nil, GORM 会提交事务
		return nil
	})
}

func (r *gatewayipRepository) CleanIPByHostId(ctx context.Context, hostId []int64) error {
	return r.DB(ctx).Model(&model.Gatewayip{}).Where("host_id IN ?", hostId).Update("host_id", 0).Error
}

func (r *gatewayipRepository) GetGatewayipOnlyIpByHostIdAll(ctx context.Context, hostId int64) ([]string, error) {
	var req []string
	return req, r.DB(ctx).Model(&model.Gatewayip{}).Where("host_id = ?", hostId).Pluck("ip", &req).Error
}