lablesys_backend/src/main/java/com/labelsys/backend/service/AnnotationTaskService.java

package com.labelsys.backend.service;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.labelsys.backend.common.ResultCode;
import com.labelsys.backend.common.exception.BusinessException;
import com.labelsys.backend.context.LoginUser;
import com.labelsys.backend.dto.common.PageResult;
import com.labelsys.backend.dto.request.AnnotationTaskPageQuery;
import com.labelsys.backend.dto.request.CreateAnnotationTaskRequest;
import com.labelsys.backend.dto.request.UpdateAnnotationTaskRequest;
import com.labelsys.backend.dto.response.AnnotationTaskResponse;
import com.labelsys.backend.entity.AnnotationTask;
import com.labelsys.backend.entity.AnnotationTaskResource;
import com.labelsys.backend.entity.SourceResource;
import com.labelsys.backend.enums.IndustryType;
import com.labelsys.backend.enums.TaskStatus;
import com.labelsys.backend.enums.TaskType;
import com.labelsys.backend.mapper.AnnotationTaskMapper;
import com.labelsys.backend.mapper.AnnotationTaskResourceMapper;
import com.labelsys.backend.mapper.SourceResourceMapper;
import com.labelsys.backend.util.IdGenerator;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

@Slf4j
@Service
@RequiredArgsConstructor
public class AnnotationTaskService {

    private final AnnotationTaskMapper         annotationTaskMapper;
    private final AnnotationTaskResourceMapper annotationTaskResourceMapper;
    private final SourceResourceMapper         sourceResourceMapper;
    private final DataPermissionService        dataPermissionService;

    @Transactional
    public AnnotationTaskResponse createTask(LoginUser currentUser, CreateAnnotationTaskRequest request) {
        try {
            List<SourceResource> resources = loadAndValidateResources(currentUser, request.resourceIds());

            AnnotationTask task = AnnotationTask.builder()
                    .id(IdGenerator.nextId())
                    .companyId(currentUser.companyId())
                    .creatorId(currentUser.userId())
                    .creatorRole(currentUser.role())
                    .taskName(request.taskName())
                    .industryType(defaultIndustryType(request.industryType()))
                    .taskType(defaultTaskType(request.taskType()))
                    .taskStatus(TaskStatus.PENDING.name())
                    .isDeleted(false)
                    .build();

            annotationTaskMapper.insert(task);
            saveTaskBindings(task.getId(), currentUser.companyId(), resources);

            log.info("created annotation task, companyId={}, userId={}, taskId={}, resourceCount={}",
                    currentUser.companyId(), currentUser.userId(), task.getId(), resources.size());

            return buildTaskResponse(task, resourceIds(resources));
        } catch (Exception e) {
            log.error("createTask failed, companyId={}, userId={}, taskName={}, error={}",
                    currentUser.companyId(), currentUser.userId(), request.taskName(), e.getMessage(), e);
            throw e;
        }
    }

    @Transactional
    public AnnotationTaskResponse updateTask(LoginUser currentUser, Long taskId, UpdateAnnotationTaskRequest request) {
        try {
            AnnotationTask task = annotationTaskMapper.findByIdAndCompanyId(taskId, currentUser.companyId());
            if (task == null) {
                throw new BusinessException(ResultCode.NOT_FOUND, "任务不存在");
            }
            assertTaskPermission(currentUser, task);

            boolean resourcesChanged = false;
            List<SourceResource> resources = null;

            if (request.resourceIds() != null && !request.resourceIds().isEmpty()) {
                List<Long> currentResourceIds = normalizeIds(
                        annotationTaskResourceMapper.listResourceIdsByTaskId(taskId));
                List<Long> targetResourceIds = normalizeIds(request.resourceIds());
                resourcesChanged = !currentResourceIds.equals(targetResourceIds);

                if (TaskStatus.RUNNING.name().equals(task.getTaskStatus()) && resourcesChanged) {
                    throw new BusinessException(ResultCode.CONFLICT, "运行中的任务不允许修改资源");
                }
                resources = loadAndValidateResources(currentUser, request.resourceIds());
            }

            if (request.industryType() != null) {
                task.setIndustryType(request.industryType());
            }

            if (request.taskType() != null) {
                task.setTaskType(request.taskType());
            }

            annotationTaskMapper.updateById(task);

            if (resourcesChanged && resources != null) {
                annotationTaskResourceMapper.deleteByTaskId(taskId);
                saveTaskBindings(taskId, currentUser.companyId(), resources);
            }

            log.info("updated annotation task, companyId={}, userId={}, taskId={}, resourcesChanged={}",
                    currentUser.companyId(), currentUser.userId(), taskId, resourcesChanged);

            List<Long> finalResourceIds = resources != null ? resourceIds(resources)
                    : normalizeIds(annotationTaskResourceMapper.listResourceIdsByTaskId(taskId));

            return buildTaskResponse(task, finalResourceIds);
        } catch (Exception e) {
            log.error("updateTask failed, companyId={}, userId={}, taskId={}, error={}",
                    currentUser.companyId(), currentUser.userId(), taskId, e.getMessage(), e);
            throw e;
        }
    }

    public AnnotationTaskResponse getTask(LoginUser currentUser, Long taskId) {
        try {
            AnnotationTask task = annotationTaskMapper.findByIdAndCompanyId(taskId, currentUser.companyId());
            if (task == null) {
                throw new BusinessException(ResultCode.NOT_FOUND, "任务不存在");
            }
            assertTaskPermission(currentUser, task);
            return buildTaskResponse(task, normalizeIds(annotationTaskResourceMapper.listResourceIdsByTaskId(taskId)));
        } catch (Exception e) {
            log.error("getTask failed, companyId={}, userId={}, taskId={}, error={}",
                    currentUser.companyId(), currentUser.userId(), taskId, e.getMessage(), e);
            throw e;
        }
    }

    public PageResult<AnnotationTaskResponse> pageTasks(LoginUser currentUser, AnnotationTaskPageQuery query) {
        try {
            List<String> allowedRoles = dataPermissionService.getAllowedRoles(currentUser);
            boolean shouldFilterByUserId = dataPermissionService.shouldFilterByUserId(currentUser);

            LambdaQueryWrapper<AnnotationTask> wrapper = new LambdaQueryWrapper<AnnotationTask>()
                    .eq(AnnotationTask::getCompanyId, currentUser.companyId())
                    .eq(query.taskType() != null, AnnotationTask::getTaskType, query.taskType())
                    .eq(StringUtils.hasText(query.taskStatus()), AnnotationTask::getTaskStatus, query.taskStatus())
                    .eq(query.isDeleted() != null, AnnotationTask::getIsDeleted, query.isDeleted())
                    .like(StringUtils.hasText(query.keyword()), AnnotationTask::getTaskName, query.keyword());

            if (shouldFilterByUserId) {
                wrapper.eq(AnnotationTask::getCreatorId, currentUser.userId());
            } else if (!allowedRoles.isEmpty()) {
                wrapper.in(AnnotationTask::getCreatorRole, allowedRoles);
            }

            wrapper.orderByDesc(AnnotationTask::getCreatedAt);

            Page<AnnotationTask> page = new Page<>(query.pageNo(), query.pageSize());
            Page<AnnotationTask> resultPage = annotationTaskMapper.selectPage(page, wrapper);

            List<AnnotationTaskResponse> records = resultPage.getRecords().stream()
                    .filter(task -> query.resourceId() == null
                            || annotationTaskResourceMapper.listResourceIdsByTaskId(task.getId())
                            .contains(query.resourceId()))
                    .map(task -> buildTaskResponse(task,
                            normalizeIds(annotationTaskResourceMapper.listResourceIdsByTaskId(task.getId()))))
                    .toList();

            return new PageResult<>(records, resultPage.getTotal(), (int) resultPage.getCurrent(),
                    (int) resultPage.getSize());
        } catch (Exception e) {
            log.error("pageTasks failed, companyId={}, userId={}, error={}",
                    currentUser.companyId(), currentUser.userId(), e.getMessage(), e);
            throw e;
        }
    }

    @Transactional
    public void deleteTask(LoginUser currentUser, Long taskId) {
        try {
            AnnotationTask task = annotationTaskMapper.findByIdAndCompanyId(taskId, currentUser.companyId());
            if (task == null) {
                throw new BusinessException(ResultCode.NOT_FOUND, "任务不存在");
            }
            assertTaskPermission(currentUser, task);
            if (TaskStatus.RUNNING.name().equals(task.getTaskStatus())) {
                throw new BusinessException(ResultCode.CONFLICT, "运行中的任务不允许删除");
            }
            task.setIsDeleted(true);
            annotationTaskMapper.updateById(task);
            log.info("deleted annotation task logically, companyId={}, userId={}, taskId={}", currentUser.companyId(),
                    currentUser.userId(), taskId);
        } catch (Exception e) {
            log.error("deleteTask failed, companyId={}, userId={}, taskId={}, error={}",
                    currentUser.companyId(), currentUser.userId(), taskId, e.getMessage(), e);
            throw e;
        }
    }

    private List<SourceResource> loadAndValidateResources(LoginUser currentUser, List<Long> resourceIds) {
        if (resourceIds == null || resourceIds.isEmpty()) {
            throw new BusinessException(ResultCode.BAD_REQUEST, "任务资源不能为空");
        }
        List<Long> normalizedIds = normalizeIds(resourceIds);
        List<SourceResource> resources = sourceResourceMapper.selectByCompanyIdAndIds(currentUser.companyId(),
                normalizedIds);
        if (resources.size() != normalizedIds.size()) {
            throw new BusinessException(ResultCode.BAD_REQUEST, "存在无效资源");
        }
        for (SourceResource resource : resources) {
            if (!dataPermissionService.canAccessCreator(currentUser, resource.getCreatorId(),
                    resource.getCreatorRole())) {
                throw new BusinessException(ResultCode.FORBIDDEN, "无权访问资源");
            }
        }
        resources.sort(Comparator.comparing(SourceResource::getId));
        return resources;
    }

    private void saveTaskBindings(Long taskId, Long companyId, List<SourceResource> resources) {
        for (SourceResource resource : resources) {
            annotationTaskResourceMapper.insert(AnnotationTaskResource.builder()
                    .id(IdGenerator.nextId())
                    .companyId(companyId)
                    .taskId(taskId)
                    .resourceId(resource.getId())
                    .build());
        }
    }

    private AnnotationTaskResponse buildTaskResponse(AnnotationTask task, List<Long> resourceIds) {
        return new AnnotationTaskResponse(
                task.getId(),
                task.getTaskName(),
                task.getIndustryType(),
                task.getTaskType(),
                task.getTaskStatus(),
                resourceIds,
                task.getCreatedAt(),
                task.getUpdatedAt()
        );
    }

    private List<Long> resourceIds(List<SourceResource> resources) {
        return resources.stream().map(SourceResource::getId).sorted().toList();
    }

    private List<Long> normalizeIds(List<Long> resourceIds) {
        Set<Long> uniqueIds = new HashSet<>(resourceIds);
        List<Long> sortedIds = new ArrayList<>(uniqueIds);
        sortedIds.sort(Long::compareTo);
        return sortedIds;
    }

    private void assertTaskPermission(LoginUser currentUser, AnnotationTask task) {
        if (!dataPermissionService.canAccessCreator(currentUser, task.getCreatorId(), task.getCreatorRole())) {
            throw new BusinessException(ResultCode.FORBIDDEN, "无权操作任务");
        }
    }

    private IndustryType defaultIndustryType(IndustryType industryType) {
        return industryType != null ? industryType : IndustryType.TRANSPORT;
    }

    private TaskType defaultTaskType(TaskType taskType) {
        return taskType != null ? taskType : TaskType.EXTRACT_QA;
    }
}