label_backend/src/main/java/com/label/service/ExportService.java

package com.label.service;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.label.common.exception.BusinessException;
import com.label.common.result.PageResult;
import com.label.common.shiro.TokenPrincipal;
import com.label.common.storage.RustFsClient;
import com.label.entity.TrainingDataset;
import com.label.mapper.TrainingDatasetMapper;
import com.label.entity.ExportBatch;
import com.label.mapper.ExportBatchMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.io.ByteArrayInputStream;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

/**
 * 训练数据导出服务。
 *
 * createBatch() 步骤：
 * 1. 校验 sampleIds 非空（EMPTY_SAMPLES 400）
 * 2. 查询 training_dataset，校验全部为 APPROVED（INVALID_SAMPLES 400）
 * 3. 生成 JSONL（每行一个 glm_format_json）
 * 4. 上传 RustFS（bucket: finetune-export, key: export/{batchUuid}.jsonl）
 * 5. 插入 export_batch 记录
 * 6. 批量更新 training_dataset.export_batch_id + exported_at
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class ExportService {

    private static final String EXPORT_BUCKET = "finetune-export";

    private final ExportBatchMapper exportBatchMapper;
    private final TrainingDatasetMapper datasetMapper;
    private final RustFsClient rustFsClient;

    // ------------------------------------------------------------------ 创建批次 --

    /**
     * 创建导出批次。
     *
     * @param sampleIds 待导出的 training_dataset ID 列表
     * @param principal 当前用户
     * @return 新建的 ExportBatch
     */
    @Transactional
    public ExportBatch createBatch(List<Long> sampleIds, TokenPrincipal principal) {
        if (sampleIds == null || sampleIds.isEmpty()) {
            throw new BusinessException("EMPTY_SAMPLES", "导出样本 ID 列表不能为空", HttpStatus.BAD_REQUEST);
        }

        // 查询样本
        List<TrainingDataset> samples = datasetMapper.selectList(
                new LambdaQueryWrapper<TrainingDataset>()
                        .in(TrainingDataset::getId, sampleIds)
                        .eq(TrainingDataset::getCompanyId, principal.getCompanyId()));

        // 校验全部已审批
        boolean hasNonApproved = samples.stream()
                .anyMatch(s -> !"APPROVED".equals(s.getStatus()));
        if (hasNonApproved || samples.size() != sampleIds.size()) {
            throw new BusinessException("INVALID_SAMPLES",
                    "部分样本不处于 APPROVED 状态或不属于当前租户", HttpStatus.BAD_REQUEST);
        }

        // 生成 JSONL（每行一个 JSON 对象）
        String jsonl = samples.stream()
                .map(TrainingDataset::getGlmFormatJson)
                .collect(Collectors.joining("\n"));
        byte[] jsonlBytes = jsonl.getBytes(StandardCharsets.UTF_8);

        // 生成唯一批次 UUID，上传 RustFS
        UUID batchUuid = UUID.randomUUID();
        String filePath = "export/" + batchUuid + ".jsonl";

        rustFsClient.upload(EXPORT_BUCKET, filePath,
                new ByteArrayInputStream(jsonlBytes), jsonlBytes.length,
                "application/jsonl");

        // 插入 export_batch 记录（若 DB 写入失败，尝试清理 RustFS 孤儿文件）
        ExportBatch batch = new ExportBatch();
        batch.setCompanyId(principal.getCompanyId());
        batch.setBatchUuid(batchUuid);
        batch.setSampleCount(samples.size());
        batch.setDatasetFilePath(filePath);
        batch.setFinetuneStatus("NOT_STARTED");
        try {
            exportBatchMapper.insert(batch);
        } catch (Exception e) {
            // DB 插入失败：尝试删除已上传的 RustFS 文件，防止产生孤儿文件
            try {
                rustFsClient.delete(EXPORT_BUCKET, filePath);
            } catch (Exception deleteEx) {
                log.error("DB 写入失败后清理 RustFS 文件亦失败，孤儿文件: {}/{}", EXPORT_BUCKET, filePath, deleteEx);
            }
            throw e;
        }

        // 批量更新 training_dataset.export_batch_id + exported_at
        datasetMapper.update(null, new LambdaUpdateWrapper<TrainingDataset>()
                .in(TrainingDataset::getId, sampleIds)
                .set(TrainingDataset::getExportBatchId, batch.getId())
                .set(TrainingDataset::getExportedAt, LocalDateTime.now())
                .set(TrainingDataset::getUpdatedAt, LocalDateTime.now()));

        log.info("导出批次已创建: batchId={}, sampleCount={}, path={}",
                batch.getId(), samples.size(), filePath);
        return batch;
    }

    // ------------------------------------------------------------------ 查询样本 --

    /**
     * 分页查询已审批、可导出的训练样本。
     */
    public PageResult<TrainingDataset> listSamples(int page, int pageSize,
                                                    String sampleType, Boolean exported,
                                                    TokenPrincipal principal) {
        pageSize = Math.min(pageSize, 100);
        LambdaQueryWrapper<TrainingDataset> wrapper = new LambdaQueryWrapper<TrainingDataset>()
                .eq(TrainingDataset::getStatus, "APPROVED")
                .eq(TrainingDataset::getCompanyId, principal.getCompanyId())
                .orderByDesc(TrainingDataset::getCreatedAt);

        if (sampleType != null && !sampleType.isBlank()) {
            wrapper.eq(TrainingDataset::getSampleType, sampleType);
        }
        if (exported != null) {
            if (exported) {
                wrapper.isNotNull(TrainingDataset::getExportBatchId);
            } else {
                wrapper.isNull(TrainingDataset::getExportBatchId);
            }
        }

        Page<TrainingDataset> result = datasetMapper.selectPage(new Page<>(page, pageSize), wrapper);
        return PageResult.of(result.getRecords(), result.getTotal(), page, pageSize);
    }

    // ------------------------------------------------------------------ 查询批次列表 --

    /**
     * 分页查询导出批次。
     */
    public PageResult<ExportBatch> listBatches(int page, int pageSize, TokenPrincipal principal) {
        pageSize = Math.min(pageSize, 100);
        Page<ExportBatch> result = exportBatchMapper.selectPage(
                new Page<>(page, pageSize),
                new LambdaQueryWrapper<ExportBatch>()
                        .eq(ExportBatch::getCompanyId, principal.getCompanyId())
                        .orderByDesc(ExportBatch::getCreatedAt));
        return PageResult.of(result.getRecords(), result.getTotal(), page, pageSize);
    }

    // ------------------------------------------------------------------ 查询批次 --

    public ExportBatch getById(Long batchId, TokenPrincipal principal) {
        ExportBatch batch = exportBatchMapper.selectById(batchId);
        if (batch == null || !batch.getCompanyId().equals(principal.getCompanyId())) {
            throw new BusinessException("NOT_FOUND", "导出批次不存在: " + batchId, HttpStatus.NOT_FOUND);
        }
        return batch;
    }
}