package com.yeejoin.amos.boot.module.common.api.dao;


import com.yeejoin.amos.boot.module.common.api.context.EsRollbackContext;
import com.yeejoin.amos.boot.module.common.api.dto.ESEquipmentCategoryDto;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.elasticsearch.annotations.Query;
import org.springframework.data.repository.PagingAndSortingRepository;
import org.springframework.stereotype.Repository;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

@Repository
public interface ESEquipmentCategory extends PagingAndSortingRepository<ESEquipmentCategoryDto, String> {

    Logger log = LoggerFactory.getLogger(ESEquipmentCategory.class);

    List<ESEquipmentCategoryDto> findAllByProjectContraptionId(String PROJECT_CONTRAPTION_ID);

    @Query("{\"prefix\": {\"ORG_BRANCH_CODE.keyword\": \"?0\"}}")
    List<ESEquipmentCategoryDto> findByOrgBranchCodeKeywordStartingWith(String orgBranchCodePrefix);

    default Iterable<ESEquipmentCategoryDto> saveAllWithFallback(Iterable<ESEquipmentCategoryDto> dtos) {
        int batchSize = 2000;

        List<ESEquipmentCategoryDto> dtoList = StreamSupport.stream(dtos.spliterator(), false)
                .collect(Collectors.toList());

        if (dtoList.isEmpty()) {
            return Collections.emptyList();
        }

        log.info("saveAllWithCache开始处理ES数据，总数量: {}, 批次大小: {}", dtoList.size(), batchSize);

        List<ESEquipmentCategoryDto> allUpdateData = new ArrayList<>();
        List<ESEquipmentCategoryDto> allInsertData = new ArrayList<>();

        // 分批处理
        for (int i = 0; i < dtoList.size(); i += batchSize) {
            int endIndex = Math.min(i + batchSize, dtoList.size());
            List<ESEquipmentCategoryDto> batch = dtoList.subList(i, endIndex);

            long startTime = System.currentTimeMillis();

            List<String> batchIds = batch.stream()
                    .map(ESEquipmentCategoryDto::getSEQUENCE_NBR)
                    .collect(Collectors.toList());

            Iterable<ESEquipmentCategoryDto> batchBeforeData = this.findAllById(batchIds);
            Map<String, ESEquipmentCategoryDto> beforeDataMap = StreamSupport.stream(
                    batchBeforeData.spliterator(), false)
                    .collect(Collectors.toMap(
                        ESEquipmentCategoryDto::getSEQUENCE_NBR,
                        Function.identity(),
                        (existing, replacement) -> existing
                    ));

            // 快速分离更新和插入数据
            List<ESEquipmentCategoryDto> batchUpdateData = new ArrayList<>();
            List<ESEquipmentCategoryDto> batchInsertData = new ArrayList<>();

            for (ESEquipmentCategoryDto dto : batch) {
                if (beforeDataMap.containsKey(dto.getSEQUENCE_NBR())) {
                    batchUpdateData.add(beforeDataMap.get(dto.getSEQUENCE_NBR()));
                } else {
                    batchInsertData.add(dto);
                }
            }

            allUpdateData.addAll(batchUpdateData);
            allInsertData.addAll(batchInsertData);

            // 批量保存当前批次
            this.saveAll(batch);

            long endTime = System.currentTimeMillis();
            log.debug("处理批次 {}/{} 完成，耗时: {}ms, 更新: {}, 插入: {}",
                      (i/batchSize + 1),
                      (int)Math.ceil((double)dtoList.size()/batchSize),
                      (endTime - startTime),
                      batchUpdateData.size(),
                      batchInsertData.size());
        }

        // 一次性缓存所有数据（避免频繁的上下文操作）
        if (!allUpdateData.isEmpty()) {
            EsRollbackContext.putAllUpdateEsToMap(allUpdateData);
            log.info("es缓存更新数据: {} 条", allUpdateData.size());
        }
        if (!allInsertData.isEmpty()) {
            EsRollbackContext.putAllInsertEsToMap(allInsertData);
            log.info("es缓存插入数据: {} 条", allInsertData.size());
        }

        log.info("saveAllWithCache数据处理完成，总计: 更新{}条, 插入{}条",
                 allUpdateData.size(), allInsertData.size());

        return dtoList;
    }

    default ESEquipmentCategoryDto saveWithFallback(ESEquipmentCategoryDto dto) {
        // 将单个对象转换为列表进行处理
        List<ESEquipmentCategoryDto> dtoList = Collections.singletonList(dto);
        Iterable<ESEquipmentCategoryDto> result = saveAllWithFallback(dtoList);
        // 返回处理后的单个对象
        return result.iterator().next();
    }
}
