observer/server/biz/dao/dataset_image.go

package dao

import (
	"context"

	"github.com/gogf/gf/v2/frame/g"

	"observer-server/biz/consts"
	"observer-server/biz/model/entity"
	"observer-server/common"
)

// DatasetImage 数据集图片表 DAO：文件名唯一（防重名），按数据集查询；表小不做缓存。
type datasetImageDao struct{}

var DatasetImage = &datasetImageDao{}

func init() {
	ctx := context.Background()
	_, err := g.DB().Exec(ctx, `CREATE TABLE IF NOT EXISTS dataset_image (
  id              INTEGER PRIMARY KEY AUTOINCREMENT,
  dataset_id      INTEGER NOT NULL,
  filename        TEXT NOT NULL,
  source          TEXT NOT NULL DEFAULT 'manual',
  prompt          TEXT,
  animal_count    INTEGER NOT NULL DEFAULT 0,
  labels_json     TEXT,
  created_at      TEXT NOT NULL,
  UNIQUE (dataset_id, filename)
)`)
	if err != nil {
		panic(err)
	}
	// 存量库迁移：标注列缺失时追加（新库建表已含列，幂等跳过）；candidates_json 已随 v10 删除
	// 注意 EnsureColumn 的 ddl 参数须自带列名（拼接为 ALTER TABLE ADD COLUMN <ddl>）
	common.EnsureColumn(ctx, consts.TableDatasetImage, "labels_json", "labels_json TEXT")
	common.EnsureColumn(ctx, consts.TableDatasetImage, "animal_count", "animal_count INTEGER NOT NULL DEFAULT 0")
}

// Insert 插入图片记录
func (d *datasetImageDao) Insert(ctx context.Context, m *entity.DatasetImage) (int64, error) {
	res, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Data(g.Map{
		"dataset_id":   m.DatasetId,
		"filename":     m.Filename,
		"source":       m.Source,
		"prompt":       m.Prompt,
		"animal_count": m.AnimalCount,
		"created_at":   m.CreatedAt,
	}).Insert()
	if err != nil {
		return 0, err
	}
	return res.LastInsertId()
}

// GetById 按主键取图片记录
func (d *datasetImageDao) GetById(ctx context.Context, id int64) (*entity.DatasetImage, error) {
	var one *entity.DatasetImage
	err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("id", id).Scan(&one)
	if err != nil {
		return nil, err
	}
	return one, nil
}

// ListByDataset 某数据集全部图片（按文件名编号倒序，详情页逐行展示用）。
// 文件名形如 <前缀>_<编号>.jpg（无前缀为时间戳数字名），须提取下划线后的数字按数值倒序：
// 字符串序会把 pheasant_101.jpg 排在 pheasant_2.jpg 之前（'1'<'9'），导致最大编号落到底部
func (d *datasetImageDao) ListByDataset(ctx context.Context, datasetId int64) ([]*entity.DatasetImage, error) {
	var list []*entity.DatasetImage
	err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
		Where("dataset_id", datasetId).
		Order("CAST(substr(filename, instr(filename, '_') + 1) AS INTEGER) DESC").Scan(&list)
	if err != nil {
		if common.IsNoRows(err) {
			return []*entity.DatasetImage{}, nil
		}
		return nil, err
	}
	return list, nil
}

// ListUnlabeledByDataset 未标注图片（labels_json 为空/null/'[]'；自动补标轮次用）
func (d *datasetImageDao) ListUnlabeledByDataset(ctx context.Context, datasetId int64) ([]*entity.DatasetImage, error) {
	var list []*entity.DatasetImage
	err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
		Where("dataset_id", datasetId).
		Where("labels_json IS NULL OR labels_json = '' OR labels_json = '[]'").
		OrderAsc("id").Scan(&list)
	if err != nil {
		if common.IsNoRows(err) {
			return []*entity.DatasetImage{}, nil
		}
		return nil, err
	}
	return list, nil
}

// GetByIds 按 id 批量取（删除图片定位文件用）
func (d *datasetImageDao) GetByIds(ctx context.Context, ids []int64) ([]*entity.DatasetImage, error) {
	var list []*entity.DatasetImage
	err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).WhereIn("id", ids).Scan(&list)
	if err != nil {
		if common.IsNoRows(err) {
			return []*entity.DatasetImage{}, nil
		}
		return nil, err
	}
	return list, nil
}

// GetByFilename 按文件名取（图片访问校验归属）
func (d *datasetImageDao) GetByFilename(ctx context.Context, datasetId int64, filename string) (*entity.DatasetImage, error) {
	var e entity.DatasetImage
	err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
		Where("dataset_id", datasetId).Where("filename", filename).Scan(&e)
	if err != nil {
		if common.IsNoRows(err) {
			return nil, nil
		}
		return nil, err
	}
	return &e, nil
}

// CountByDataset 某数据集图片数
func (d *datasetImageDao) CountByDataset(ctx context.Context, datasetId int64) (int64, error) {
	n, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("dataset_id", datasetId).Count()
	return int64(n), err
}

// UpdateLabels 覆写标注 JSON（整图粒度；空串=清空标注；AI 自动标注与人工保存共用）
func (d *datasetImageDao) UpdateLabels(ctx context.Context, id int64, labelsJson string) error {
	_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("id", id).
		Data(g.Map{"labels_json": labelsJson}).Update()
	return err
}

// CountLabeledByDataset 某数据集已标注图片数（labels_json 非空数组的图片行数）
func (d *datasetImageDao) CountLabeledByDataset(ctx context.Context, datasetId int64) (int64, error) {
	n, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
		Where("dataset_id", datasetId).
		Where("labels_json IS NOT NULL AND labels_json != '' AND labels_json != '[]'").
		Count()
	return int64(n), err
}

// CountStat 某数据集图片数/已标注数（列表实时统计：冗余计数列会漂移，展示一律以此为准）
type CountStat struct {
	Count   int64
	Labeled int64
}

// CountStatByDatasets 多数据集图片数/已标注数（一次 GROUP BY 查询；IN 按 ≤100 分批）
func (d *datasetImageDao) CountStatByDatasets(ctx context.Context, datasetIds []int64) (map[int64]CountStat, error) {
	stat := make(map[int64]CountStat)
	for start := 0; start < len(datasetIds); start += 100 {
		end := start + 100
		if end > len(datasetIds) {
			end = len(datasetIds)
		}
		var rows []struct {
			DatasetId int64 `orm:"dataset_id"`
			Cnt       int64 `orm:"cnt"`
			Labeled   int64 `orm:"labeled_cnt"`
		}
		err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).
			WhereIn("dataset_id", datasetIds[start:end]).
			Fields("dataset_id, COUNT(*) AS cnt, SUM(CASE WHEN labels_json IS NOT NULL AND labels_json != '' AND labels_json != '[]' THEN 1 ELSE 0 END) AS labeled_cnt").
			Group("dataset_id").
			Scan(&rows)
		if err != nil {
			return nil, err
		}
		for _, r := range rows {
			stat[r.DatasetId] = CountStat{Count: r.Cnt, Labeled: r.Labeled}
		}
	}
	return stat, nil
}

// DeleteByIds 按 id 批量删除（标注随行删除）
func (d *datasetImageDao) DeleteByIds(ctx context.Context, ids []int64) error {
	_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).WhereIn("id", ids).Delete()
	return err
}

// DeleteByDataset 删除某数据集全部图片记录（删数据集联动）
func (d *datasetImageDao) DeleteByDataset(ctx context.Context, datasetId int64) error {
	_, err := g.DB().Model(consts.TableDatasetImage).Ctx(ctx).Where("dataset_id", datasetId).Delete()
	return err
}