package com.ycwl.basic.face.pipeline.stages;

import com.ycwl.basic.face.pipeline.core.FaceMatchingContext;
import com.ycwl.basic.mapper.FaceSampleAssociationMapper;
import com.ycwl.basic.model.task.resp.SearchFaceRespVo;
import com.ycwl.basic.pipeline.annotation.StageConfig;
import com.ycwl.basic.pipeline.core.AbstractPipelineStage;
import com.ycwl.basic.pipeline.core.StageResult;
import com.ycwl.basic.pipeline.enums.StageOptionalMode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * 根据人脸样本关联分组扩展匹配结果
 * <p>
 * 当匹配结果中包含某分组的任一 faceSampleId 时，
 * 将该分组内所有 faceSampleId 加入匹配结果。
 * 关联扩展的样本ID跳过时间范围和设备限制筛选（通过Stage执行顺序保证）。
 */
@Slf4j
@Component
@StageConfig(
    stageId = "expand_sample_association",
    optionalMode = StageOptionalMode.SUPPORT,
    description = "根据人脸样本关联分组扩展匹配结果",
    defaultEnabled = true
)
public class ExpandSampleAssociationStage extends AbstractPipelineStage<FaceMatchingContext> {

    @Autowired
    private FaceSampleAssociationMapper faceSampleAssociationMapper;

    @Override
    public String getName() {
        return "ExpandSampleAssociation";
    }

    @Override
    protected boolean shouldExecuteByBusinessLogic(FaceMatchingContext context) {
        SearchFaceRespVo searchResult = context.getSearchResult();
        if (searchResult == null || searchResult.getSampleListIds() == null
                || searchResult.getSampleListIds().isEmpty()) {
            return false;
        }
        return context.getFace() != null && context.getFace().getScenicId() != null;
    }

    @Override
    protected StageResult<FaceMatchingContext> doExecute(FaceMatchingContext context) {
        // 防御性检查
        SearchFaceRespVo searchResult = context.getSearchResult();
        if (searchResult == null || searchResult.getSampleListIds() == null
                || searchResult.getSampleListIds().isEmpty()) {
            return StageResult.skipped("searchResult或sampleListIds为空");
        }

        Long scenicId = context.getFace().getScenicId();
        List<Long> originalMatchedIds = searchResult.getSampleListIds();

        try {
            List<Long> associatedIds = faceSampleAssociationMapper
                    .findAssociatedSampleIds(scenicId, originalMatchedIds);

            if (associatedIds == null || associatedIds.isEmpty()) {
                log.debug("未找到关联样本, faceId={}, scenicId={}", context.getFaceId(), scenicId);
                return StageResult.success("无关联样本");
            }

            // 获取当前的 sampleListIds（可能已经过筛选，也可能与 searchResult 相同）
            List<Long> currentSampleIds = context.getSampleListIds();
            if (currentSampleIds == null || currentSampleIds.isEmpty()) {
                currentSampleIds = new ArrayList<>(originalMatchedIds);
            }

            // 计算净新增的关联ID
            Set<Long> currentSet = new HashSet<>(currentSampleIds);
            List<Long> netNewIds = associatedIds.stream()
                    .filter(id -> !currentSet.contains(id))
                    .distinct()
                    .collect(Collectors.toList());

            if (netNewIds.isEmpty()) {
                log.debug("关联样本均已存在于当前结果中, faceId={}", context.getFaceId());
                return StageResult.success("关联样本均已存在");
            }

            // 合并：当前筛选后的ID + 净新增关联ID
            List<Long> expandedIds = Stream.concat(
                    currentSampleIds.stream(),
                    netNewIds.stream()
            ).distinct().collect(Collectors.toList());

            context.setSampleListIds(expandedIds);
            context.setAssociatedSampleIds(netNewIds);

            log.info("关联样本扩展完成: faceId={}, scenicId={}, 原始匹配数={}, 扩展前={}, 净新增={}, 扩展后={}",
                    context.getFaceId(), scenicId,
                    originalMatchedIds.size(), currentSampleIds.size(),
                    netNewIds.size(), expandedIds.size());

            return StageResult.success(String.format("关联扩展: +%d个样本", netNewIds.size()));

        } catch (Exception e) {
            log.error("关联样本扩展失败, faceId={}, scenicId={}", context.getFaceId(), scenicId, e);
            return StageResult.degraded("关联样本扩展失败: " + e.getMessage());
        }
    }
}
