package com.youchain.utils;

import cn.idev.excel.ExcelWriter;
import cn.idev.excel.FastExcel;
import cn.idev.excel.write.metadata.WriteSheet;
import cn.idev.excel.write.metadata.style.WriteCellStyle;
import cn.idev.excel.write.metadata.style.WriteFont;
import cn.idev.excel.write.style.HorizontalCellStyleStrategy;

import com.youchain.exception.BadRequestException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.poi.ss.usermodel.IndexedColors;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.MediaTypeFactory;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.Collection;
import java.util.List;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.ToLongFunction;
import java.util.stream.Collectors;

import static cn.hutool.core.util.CharsetUtil.UTF_8;


/**
 * excel 工具类
 */
@Slf4j
public class FastExcelUtil {

    /**
     * 设置下载文件消息头
     *
     * @param response 响应
     * @param fileName 文件名
     * @param fileSize 文件大小
     */
    public static void setDownloadFileHeader(HttpServletResponse response, String fileName, Long fileSize) {
        response.setCharacterEncoding(UTF_8);
        if (fileSize != null) {
            response.setHeader(HttpHeaders.CONTENT_LENGTH, String.valueOf(fileSize));
        }

        if (StringUtils.isNotEmpty(fileName)) {
            response.setHeader(HttpHeaders.CONTENT_TYPE, MediaTypeFactory.getMediaType(fileName).orElse(MediaType.APPLICATION_OCTET_STREAM) + ";charset=utf-8");
            try {
                response.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment;filename=" + URLEncoder.encode(fileName, UTF_8).replaceAll("\\+", "%20"));
            } catch (UnsupportedEncodingException e) {
                throw new RuntimeException(e);
            }
            response.setHeader(HttpHeaders.ACCESS_CONTROL_EXPOSE_HEADERS, HttpHeaders.CONTENT_DISPOSITION);
        }
    }

    /**
     * 获取样式
     *
     * @return horizontalCellStyleStrategy
     */
    public static HorizontalCellStyleStrategy getHorizontalCellStyleStrategy() {

        WriteCellStyle headWriteCellStyle = new WriteCellStyle();
        headWriteCellStyle.setFillForegroundColor(IndexedColors.SKY_BLUE.getIndex());

        WriteFont headWriteFont = new WriteFont();
        headWriteFont.setFontName("宋体");
        headWriteFont.setColor(IndexedColors.WHITE.getIndex());
        headWriteCellStyle.setWriteFont(headWriteFont);

        WriteCellStyle contentWriteCellStyle = new WriteCellStyle();
        contentWriteCellStyle.setFillForegroundColor(IndexedColors.BLACK.getIndex());

        return new HorizontalCellStyleStrategy(headWriteCellStyle, contentWriteCellStyle);
    }

    /**
     * 通用单sheet导出
     */
    public static void exportExcel(HttpServletResponse response, String fileName, String sheetName, Class head, Collection<?> data) throws IOException {
        // 设置下载消息头
        setDownloadFileHeader(response, fileName, null);
        //样式
        HorizontalCellStyleStrategy horizontalCellStyleStrategy = getHorizontalCellStyleStrategy();
        // 下载
        FastExcel.write(response.getOutputStream(), head)
                .autoCloseStream(Boolean.FALSE)
                .sheet(sheetName)
                .registerWriteHandler(horizontalCellStyleStrategy)
                .doWrite(data);
    }

    /**
     * 大数据导出
     * @param response 响应
     * @param fileName 文件名
     * @param sheetName sheet名称
     * @param head  头
     * @param pageQueryFunction 查询方法
     * @param convertFunction 转换方法
     * @param idExtractor id提取方法
     * @param pageSize 每页大小
     * @param <T> 泛型
     * @param <E> 泛型
     */
    public static <T, E> void batchExportExcel(HttpServletResponse response, String fileName, String sheetName, Class head, BiFunction<Long, Integer, List<T>> pageQueryFunction, Function<T, E> convertFunction, ToLongFunction<T> idExtractor, int pageSize) {
        long startTime = System.currentTimeMillis();
        // 设置下载消息头
        setDownloadFileHeader(response, fileName, null);
        HorizontalCellStyleStrategy horizontalCellStyleStrategy = getHorizontalCellStyleStrategy();
        try (ExcelWriter excelWriter = FastExcel.write(response.getOutputStream(), head).registerWriteHandler(horizontalCellStyleStrategy).build()) {
            WriteSheet writeSheet = FastExcel.writerSheet(sheetName).build();
            long lastId = 0;//最后一个id
            while (true) {
                List<T> batch = pageQueryFunction.apply(lastId, pageSize);
                if (batch.isEmpty()) {
                    break;
                }

                List<E> excelData = batch.stream()
                        .map(convertFunction)
                        .collect(Collectors.toList());
                excelWriter.write(excelData, writeSheet);

                lastId = batch.stream()
                        .mapToLong(idExtractor)
                        .max()
                        .orElse(lastId);
            }
            excelWriter.finish();
            log.info("导出总耗时: {}ms", System.currentTimeMillis() - startTime);
        } catch (Exception e) {
            throw new BadRequestException("导出失败");
        }
    }


}