pig-farm-controller/internal/infra/repository/area_controller_repository.go

package repository

import (
	"context"
	"fmt"
	"strconv"

	"git.huangwc.com/pig/pig-farm-controller/internal/infra/logs"
	"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"

	"gorm.io/gorm"
)

// AreaControllerRepository 定义了对 AreaController 模型的数据库操作接口
type AreaControllerRepository interface {
	FindByID(ctx context.Context, id uint32) (*models.AreaController, error)
	FindByNetworkID(ctx context.Context, networkID string) (*models.AreaController, error)
	Create(ctx context.Context, ac *models.AreaController) error
	ListAll(ctx context.Context) ([]*models.AreaController, error)
	Update(ctx context.Context, ac *models.AreaController) error
	// UpdateFirmwareVersion 更新指定ID的区域主控的固件版本号。
	UpdateFirmwareVersion(ctx context.Context, id uint32, version string) error
	Delete(ctx context.Context, id uint32) error
	// IsAreaControllerUsedByTasks 检查区域主控是否被特定任务类型使用，可以忽略指定任务类型
	IsAreaControllerUsedByTasks(ctx context.Context, areaControllerID uint32, ignoredTaskTypes []models.TaskType) (bool, error)
}

// gormAreaControllerRepository 是 AreaControllerRepository 的 GORM 实现。
type gormAreaControllerRepository struct {
	ctx context.Context
	db  *gorm.DB
}

// NewGormAreaControllerRepository 创建一个新的 AreaControllerRepository GORM 实现实例。
func NewGormAreaControllerRepository(ctx context.Context, db *gorm.DB) AreaControllerRepository {
	return &gormAreaControllerRepository{
		ctx: ctx,
		db:  db,
	}
}

// Create 创建一个新的 AreaController 记录。
func (r *gormAreaControllerRepository) Create(ctx context.Context, ac *models.AreaController) error {
	repoCtx := logs.AddFuncName(ctx, r.ctx, "Create")
	return r.db.WithContext(repoCtx).Create(ac).Error
}

// ListAll 返回所有 AreaController 的列表。
func (r *gormAreaControllerRepository) ListAll(ctx context.Context) ([]*models.AreaController, error) {
	repoCtx := logs.AddFuncName(ctx, r.ctx, "ListAll")
	var areaControllers []*models.AreaController
	if err := r.db.WithContext(repoCtx).Find(&areaControllers).Error; err != nil {
		return nil, err
	}
	return areaControllers, nil
}

// Update 更新一个已存在的 AreaController 记录。
func (r *gormAreaControllerRepository) Update(ctx context.Context, ac *models.AreaController) error {
	repoCtx := logs.AddFuncName(ctx, r.ctx, "Update")
	return r.db.WithContext(repoCtx).Save(ac).Error
}

// UpdateFirmwareVersion 使用 jsonb_set 函数原子性地更新 properties 字段中的固件版本号。
func (r *gormAreaControllerRepository) UpdateFirmwareVersion(ctx context.Context, id uint32, version string) error {
	repoCtx := logs.AddFuncName(ctx, r.ctx, "UpdateFirmwareVersion")

	// 使用 gorm.Expr 包装 PostgreSQL 的 jsonb_set 函数
	// jsonb_set(properties, '{firmware_version}', '"new_version"', true)
	// 注意：jsonb_set 的第三个参数需要是有效的 JSON 值，所以字符串需要被双引号包围。
	jsonbExpr := gorm.Expr(`jsonb_set(COALESCE(properties, '{}'::jsonb), '{firmware_version}', ?::jsonb)`, fmt.Sprintf(`"%s"`, version))

	result := r.db.WithContext(repoCtx).Model(&models.AreaController{}).Where("id = ?", id).Update("properties", jsonbExpr)
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return fmt.Errorf("更新固件版本失败：未找到ID为 %d 的区域主控", id)
	}
	return nil
}

// Delete 删除一个 AreaController 记录。
func (r *gormAreaControllerRepository) Delete(ctx context.Context, id uint32) error {
	repoCtx := logs.AddFuncName(ctx, r.ctx, "Delete")
	if err := r.db.WithContext(repoCtx).Delete(&models.AreaController{}, id).Error; err != nil {
		return fmt.Errorf("删除区域主控失败: %w", err)
	}
	return nil
}

// FindByID 通过 ID 查找一个 AreaController。
func (r *gormAreaControllerRepository) FindByID(ctx context.Context, id uint32) (*models.AreaController, error) {
	repoCtx := logs.AddFuncName(ctx, r.ctx, "FindByID")
	var areaController models.AreaController
	if err := r.db.WithContext(repoCtx).First(&areaController, id).Error; err != nil {
		return nil, err
	}
	return &areaController, nil
}

// FindByNetworkID 通过 NetworkID 查找一个 AreaController。
func (r *gormAreaControllerRepository) FindByNetworkID(ctx context.Context, networkID string) (*models.AreaController, error) {
	repoCtx := logs.AddFuncName(ctx, r.ctx, "FindByNetworkID")
	var areaController models.AreaController
	if err := r.db.WithContext(repoCtx).Where("network_id = ?", networkID).First(&areaController).Error; err != nil {
		return nil, err
	}
	return &areaController, nil
}

// IsAreaControllerUsedByTasks 检查区域主控是否被特定任务类型使用，可以忽略指定任务类型
func (r *gormAreaControllerRepository) IsAreaControllerUsedByTasks(ctx context.Context, areaControllerID uint32, ignoredTaskTypes []models.TaskType) (bool, error) {
	repoCtx, logger := logs.Trace(ctx, r.ctx, "IsAreaControllerUsedByTasks")

	// 将 ignoredTaskTypes 转换为 map，以便高效查找
	ignoredMap := make(map[models.TaskType]struct{})
	for _, tt := range ignoredTaskTypes {
		ignoredMap[tt] = struct{}{}
	}

	areaControllerIDStr := strconv.FormatUint(uint64(areaControllerID), 10)

	// 定义所有可能与 AreaControllerID 相关的任务类型列表
	// 方便未来扩展，如果新增任务类型与区域主控关联，只需在此处添加
	relevantTaskTypes := []models.TaskType{
		models.TaskTypeAreaCollectorThresholdCheck,
		// TODO: 如果未来有其他任务类型通过 parameters 关联 AreaControllerID，请在此处添加
		// 例如: models.TaskTypeAnotherAreaControllerTask,
	}

	for _, taskType := range relevantTaskTypes {
		// 如果当前任务类型在忽略列表中，则跳过检查
		if _, ok := ignoredMap[taskType]; ok {
			continue
		}

		var taskCount int64
		var query *gorm.DB

		// 根据任务类型构建不同的查询条件
		switch taskType {
		case models.TaskTypeAreaCollectorThresholdCheck:
			// TaskTypeAreaCollectorThresholdCheck 任务的 AreaControllerID 存储在 parameters->>'AreaControllerID'
			query = r.db.WithContext(repoCtx).
				Model(&models.Task{}).
				Where("type = ?", models.TaskTypeAreaCollectorThresholdCheck).
				Where("parameters->>'AreaControllerID' = ?", areaControllerIDStr)
		// TODO: 如果未来有其他任务类型通过不同的 parameters 字段关联 AreaControllerID，请在此处添加 case
		// case models.TaskTypeAnotherAreaControllerTask:
		//     query = r.db.WithContext(repoCtx).
		//         Model(&models.Task{}).
		//         Where("type = ?", models.TaskTypeAnotherAreaControllerTask).
		//         Where("parameters->>'AnotherFieldForAreaControllerID' = ?", areaControllerIDStr)
		default:
			// 对于未明确处理的 relevantTaskTypes，可以记录警告或直接跳过
			logger.Warnf(fmt.Sprintf("IsAreaControllerUsedByTasks: 未处理的区域主控相关任务类型: %s", taskType))
			continue
		}

		if query != nil {
			err := query.Count(&taskCount).Error
			if err != nil {
				return false, fmt.Errorf("查询区域主控任务使用情况失败 (任务类型: %s): %w", taskType, err)
			}
			if taskCount > 0 {
				return true, nil // 发现有未被忽略的任务正在使用此区域主控
			}
		}
	}

	return false, nil // 没有发现任何未被忽略的任务正在使用此区域主控
}