package com.bonus.canteen.core.common.export.service;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.text.CharSequenceUtil;
import com.alibaba.excel.EasyExcelFactory;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.converters.Converter;
import com.alibaba.excel.converters.longconverter.LongStringConverter;
import com.alibaba.excel.write.builder.ExcelWriterBuilder;
import com.alibaba.excel.write.handler.WriteHandler;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.bonus.canteen.core.common.export.constant.ExportFileTypeEnum;
import com.bonus.canteen.core.common.export.constant.ExportTaskWaitPolicy;
import com.bonus.canteen.core.common.export.model.ExportRecord;
import com.bonus.canteen.core.common.export.po.ExcelExportParam;
import com.bonus.canteen.core.common.page.PageDTO;
import com.bonus.canteen.core.common.page.PageVO;
import com.bonus.canteen.core.common.utils.*;
import com.bonus.canteen.core.starter.threadpool.CleanTenantIdTaskDecorator;
import com.bonus.common.core.exception.ServiceException;
import com.bonus.common.houqin.constant.LeConstants;
import com.bonus.common.houqin.i18n.I18n;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.task.AsyncTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.io.File;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
import java.util.function.LongFunction;
import java.util.function.Supplier;
import java.util.stream.Collectors;

@Component
public class ExportAsyncManager {
   private static final Logger log = LoggerFactory.getLogger(ExportAsyncManager.class);
   protected static final long EXPORT_TOTAL_LIMIT = 1000000L;
   private static final long PAGE_SIZE = 1000L;
   @Resource
   private ExportRecordService exportRecordService;
   @Resource
   private ExportUploadManager exportUploadManager;
   private AsyncTaskExecutor asyncTaskExecutor;

   public <T> ExportRecord startExcelExportTask(ExcelExportParam<T> param) {
      if (ExportTaskWaitPolicy.DROP.equals(param.getWaitPolicy()) && this.exportRecordService.existsWaiting()) {
         throw new ServiceException(I18n.getMessage("common.export.another-export", new Object[0]));
      } else if (param.getTotalCount() != null && param.getTotalCount() > 1000000L) {
         throw new ServiceException(I18n.getMessage("common.export.max-limit", new Object[]{1000000L}));
      } else {
         ExportRecord exportRecord = this.exportRecordService.create(param.getFileName(), ExportFileTypeEnum.EXCEL, param.isIfPrivate(), param.getExpiredTime() != null ? param.getExpiredTime() : LocalDateTime.now().plusDays(7L));
         Long exportId = exportRecord.getExportId();
         String tenantIdStr = TenantContextHolder.getTenantId().toString();
         String traceId = LogUtil.getCurrentTraceId();

         try {
            this.asyncTaskExecutor.execute(() -> {
               this.doExecute(param, exportId, tenantIdStr, traceId);
            });
         } catch (Exception var7) {
            log.error("启动异步执行导出任务失败", var7);
            this.exportRecordService.didCanceled(exportId);
         }

         return exportRecord;
      }
   }

   public <T> ExportRecord startExcelExportTaskByPage(String fileName, String sheetName, Class<T> headClass, List<String> columnNames, PageDTO pageDTO, T statLine, Supplier<PageVO<T>> pageQueryFunc) {
      PageVO<?> totalPageData = this.queryTotalData(pageDTO, pageQueryFunc);
      long total = totalPageData.getTotal();
      long totalPages = this.getPages(total);
      AtomicLong currentPage = new AtomicLong(1L);
      return this.startExcelExportTask(ExcelExportParam.ofHeadClass(headClass).setFileName(fileName).setSheetName(sheetName).setColumnNames(columnNames).setTotalCount(total).setDataQueryFunc((currentTotal) -> {
         Collection<T> data = this.queryPageData(currentPage, totalPages, pageDTO, pageQueryFunc).getRecords();
         if (CollUtil.isEmpty(data)) {
            return data;
         } else {
            if (currentPage.get() == totalPages && statLine != null) {
               data.add(statLine);
            }

            currentPage.getAndIncrement();
            return data;
         }
      }));
   }

   public <T, R> ExportRecord startExcelExportTaskByPage(String fileName, String sheetName, Class<R> headClass, List<String> columnNames, PageDTO pageDTO, T statLine, Supplier<PageVO<T>> pageQueryFunc, Function<T, R> mapFunc) {
      PageVO<?> totalPageData = this.queryTotalData(pageDTO, pageQueryFunc);
      long total = totalPageData.getTotal();
      long totalPages = this.getPages(total);
      AtomicLong currentPage = new AtomicLong(1L);
      return this.startExcelExportTask(ExcelExportParam.ofHeadClass(headClass).setFileName(fileName).setSheetName(sheetName).setColumnNames(columnNames).setTotalCount(total).setDataQueryFunc((currentTotal) -> {
         Collection<R> data = (Collection)this.queryPageData(currentPage, totalPages, pageDTO, pageQueryFunc).getRecords().stream().map(mapFunc).collect(Collectors.toList());
         if (CollUtil.isEmpty(data)) {
            return data;
         } else {
            if (currentPage.get() == totalPages && statLine != null) {
               data.add(mapFunc.apply(statLine));
            }

            currentPage.getAndIncrement();
            return data;
         }
      }));
   }

   public String getFileFullName(String fileName) {
      return CharSequenceUtil.endWith(fileName, ".xlsx") ? fileName : fileName + "_" + LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd_HH-mm-ss-SSS")) + ".xlsx";
   }

   protected <T> PageVO<T> queryTotalData(PageDTO pageDTO, Supplier<PageVO<T>> pageQueryFunc) {
      pageDTO.setCurrent(1L);
      pageDTO.setSize(-1L);
      pageDTO.setIfCount(LeConstants.COMMON_YES);
      pageDTO.setIfPageSizeZero(LeConstants.COMMON_NO);
      return (PageVO)pageQueryFunc.get();
   }

   protected <T> Page<T> queryTotalData(Page page, Supplier<Page<T>> pageQueryFunc) {
      page.setCurrent(1L);
      page.setSize(0L);
      return (Page)pageQueryFunc.get();
   }

   protected <T> PageVO<T> queryPageData(AtomicLong currentPage, long totalPages, PageDTO pageDTO, Supplier<PageVO<T>> pageQueryFunc) {
      if (currentPage.get() > totalPages) {
         return PageVO.empty();
      } else {
         pageDTO.setCurrent(currentPage.get());
         pageDTO.setSize(1000L);
         pageDTO.setIfCount(LeConstants.COMMON_NO);
         pageDTO.setIfPageSizeZero(LeConstants.COMMON_YES);
         return (PageVO)pageQueryFunc.get();
      }
   }

   protected <T> Page<T> queryPageData(AtomicLong currentPage, long totalPages, Page page, Supplier<Page<T>> pageQueryFunc) {
      if (currentPage.get() > totalPages) {
         return new Page();
      } else {
         page.setCurrent(currentPage.get());
         page.setSize(1000L);
         page.setSearchCount(false);
         return (Page)pageQueryFunc.get();
      }
   }

   protected long getPages(long total) {
      long pages = total / 1000L;
      if (total % 1000L != 0L) {
         ++pages;
      }

      return pages;
   }

   public <T, R> ExportRecord startExcelExportTaskByMPPage(String fileName, String sheetName, Class<R> headClass, List<String> columnNames, Page page, T statLine, Supplier<Page<T>> pageQueryFunc, Function<T, R> mapFunc) {
      Page<?> totalPageData = this.queryTotalData(page, pageQueryFunc);
      long total = totalPageData.getTotal();
      long totalPages = this.getPages(total);
      AtomicLong currentPage = new AtomicLong(1L);
      return this.startExcelExportTask(ExcelExportParam.ofHeadClass(headClass).setFileName(fileName).setSheetName(sheetName).setColumnNames(columnNames).setTotalCount(total).setDataQueryFunc((currentTotal) -> {
         Collection<R> data = (Collection)this.queryPageData(currentPage, totalPages, page, pageQueryFunc).getRecords().stream().map(mapFunc).collect(Collectors.toList());
         if (CollUtil.isEmpty(data)) {
            return data;
         } else {
            if (currentPage.get() == totalPages && statLine != null) {
               data.add(mapFunc.apply(statLine));
            }

            currentPage.getAndIncrement();
            return data;
         }
      }));
   }

   protected <T> void doExecute(ExcelExportParam<T> param, Long exportId, String merchantIdString, String currentTraceId) {
      LogUtil.putLogTraceId(currentTraceId);
      File tempFile = null;

      try {
         tempFile = File.createTempFile(exportId.toString(), ".xlsx");
         ExcelWriterBuilder writerBuilder =EasyExcelFactory.write(tempFile).registerConverter(new LongStringConverter())
                 .registerConverter(new CustomLocalDateStringConverter())
                 .registerConverter(new CustomLocalDateConverter())
                 .includeColumnFiledNames(param.getColumnNames());
         List var10000;
         if (CollUtil.isNotEmpty(param.getAdditionalConvertors())) {
            List<Converter<?>> converters = (List<Converter<?>>) param.getAdditionalConvertors();
            Objects.requireNonNull(writerBuilder);
            converters.forEach(writerBuilder::registerConverter);
         }

         if (CollUtil.isNotEmpty(param.getAdditionalWriteHandlers())) {
            List<Object> handlers = Collections.singletonList(param.getAdditionalWriteHandlers());
            Objects.requireNonNull(writerBuilder);
            for (Object handler : handlers) {
               if (handler instanceof WriteHandler) {
                  writerBuilder.registerWriteHandler((WriteHandler) handler);
               } else {
                  throw new IllegalArgumentException("Handler must be an instance of WriteHandler");
               }
            }
         }


         ExcelWriter excelWriter = writerBuilder.build();
         Long totalCount = param.getTotalCount();
         LongFunction<Collection<T>> dataQueryFunc = (currentTotal) -> {
            log.info("任务：{} 导出进度： {} / {}", new Object[]{exportId, currentTotal, totalCount});
            if (this.exportRecordService.shouldStop(exportId)) {
               return CollUtil.newArrayList((T) new Object[0]);
            } else {
               Collection<T> data = (Collection)param.getDataQueryFunc().apply(currentTotal);
               int progress = (int)(totalCount != null && totalCount != 0L ? currentTotal * 100L / totalCount : 0L);
               this.exportRecordService.updateCreatingProgress(exportId, progress >= 100 ? 99 : progress, (String)null, (Long)null);
               return data;
            }
         };
         EasyExcelUtil.writeExcelByDownLoadIncludeWriteByPage(excelWriter, param.getSheetName(), param.getHeaderClass(), dataQueryFunc, param.getSheetNameFunc());
         if (this.exportRecordService.shouldStop(exportId)) {
            this.exportRecordService.didStop(exportId);
            return;
         }

         String fileUrl = this.exportUploadManager.upload(param.getFileName(), tempFile, merchantIdString);
         this.exportRecordService.updateCreatingProgress(exportId, 100, fileUrl, tempFile.length());
      } catch (Exception var20) {
         log.error("导出失败", var20);
         this.exportRecordService.didStop(exportId);
      } finally {
         try {
            if (tempFile != null && tempFile.exists() && tempFile.delete()) {
               log.info("删除临时文件： {} ", tempFile);
            }
         } catch (Exception var19) {
            log.error("删除文件失败", var19);
         }

      }

   }

   @PostConstruct
   protected void initTaskExecutor() {
      ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
      taskExecutor.setMaxPoolSize(1);
      taskExecutor.setCorePoolSize(1);
      taskExecutor.setThreadNamePrefix("async-yst-export-task");
      taskExecutor.setQueueCapacity(10);
      taskExecutor.setKeepAliveSeconds(3600);
      taskExecutor.setWaitForTasksToCompleteOnShutdown(true);
      taskExecutor.setAwaitTerminationSeconds(60);
      taskExecutor.setTaskDecorator(new CleanTenantIdTaskDecorator());
      taskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
      taskExecutor.initialize();
      this.asyncTaskExecutor = taskExecutor;
   }
}
