package cn.iocoder.yudao.module.llm.utils;

import cn.iocoder.yudao.framework.common.util.collection.CollectionUtils;
import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.DatasetAnswerRespVO;
import cn.iocoder.yudao.module.llm.controller.admin.modelassesstaskmanual.vo.ManualAssessmentPageRespVO;
import com.alibaba.excel.util.StringUtils;
import org.apache.poi.hssf.usermodel.HSSFCellStyle;
import org.apache.poi.hssf.usermodel.HSSFRow;
import org.apache.poi.hssf.usermodel.HSSFSheet;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.apache.poi.ss.usermodel.HorizontalAlignment;
import org.apache.poi.ss.usermodel.VerticalAlignment;
import org.apache.poi.ss.util.CellRangeAddress;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.stream.Collectors;

public class ManalExcelExportUtils {

    public static <T> void write(HttpServletResponse response, String filename, String sheetName,
                                 Class<T> head, List<ManualAssessmentPageRespVO> data) throws IOException {
        HSSFWorkbook template = new HSSFWorkbook();
        HSSFSheet sheet = template.createSheet(sheetName);

        HSSFCellStyle cellStyle = template.createCellStyle();
        cellStyle.setVerticalAlignment(VerticalAlignment.CENTER);
        cellStyle.setAlignment(HorizontalAlignment.CENTER);
        // 创建样式并设置垂直居中
        int count = 0;
        HSSFRow row = sheet.createRow(count);
        row.createCell(0).setCellValue("模型名称");
        row.createCell(1).setCellValue("结果集名称");
        row.createCell(2).setCellValue("System(人设信息)");
        row.createCell(3).setCellValue("Prompt");
        row.createCell(4).setCellValue("Response(参考回答)");
        row.createCell(5).setCellValue("Completion(模型回答)");

        List<ManualAssessmentPageRespVO> collect1 = data.stream().filter(item -> item.getStatus() == 2).collect(Collectors.toList());
        HashMap<String, Integer> rowMap = new HashMap<>();
        if (collect1.size() != 0) {
            int num = 6;
            List<HashMap<String, Object>> reqRespVos = collect1.get(0).getReqRespVos();
            for (HashMap<String, Object> map : reqRespVos) {
                row.createCell(num).setCellValue(String.valueOf(map.get("label")));
                rowMap.put(String.valueOf(map.get("label")), Integer.valueOf(num));
                num++;
            }
        }
        count++;
        List<Integer> id = new ArrayList<>();
        for (ManualAssessmentPageRespVO item : data) {
            List<DatasetAnswerRespVO> datasetAnswerRespVO = item.getAnswers();
            List<String> collect = datasetAnswerRespVO.stream().map(DatasetAnswerRespVO::getAnswer).collect(Collectors.toList());
            int start = count;
            row = sheet.createRow(count);
            row.createCell(0).setCellValue(item.getModelServiceName() == null ? "" : item.getModelServiceName());
            row.getCell(0).setCellStyle(cellStyle);
            row.createCell(1).setCellValue(item.getDatasetName() == null ? "" : item.getDatasetName());
            row.getCell(1).setCellStyle(cellStyle);
            row.createCell(2).setCellValue(StringUtils.isBlank(item.getSystem())?"-":item.getSystem());
            row.getCell(2).setCellStyle(cellStyle);
            row.createCell(3).setCellValue(item.getQuestion() == null ? "" : item.getQuestion());
            row.getCell(3).setCellStyle(cellStyle);
            row.createCell(4).setCellValue(collect.get(0) == null ? "" : collect.get(0));
            row.getCell(4).setCellStyle(cellStyle);
            row.createCell(5).setCellValue(item.getModelAnswer() == null ? "" : item.getModelAnswer());
            row.getCell(5).setCellStyle(cellStyle);
            if (!CollectionUtils.isAnyEmpty(item.getReqRespVos())) {
                for (HashMap<String, Object> map : item.getReqRespVos()) {
                    int num = rowMap.get(String.valueOf(map.get("label")));
                    row.createCell(num).setCellValue(String.valueOf(map.get("score")) == null ? "" : String.valueOf(map.get("score")));
                    row.getCell(num).setCellStyle(cellStyle);
                }
            }
            count++;
            if (collect.size() > 1) {
                for (int i = 1; i < collect.size(); i++) {
                    row = sheet.createRow(count);
                    row.createCell(4).setCellValue(collect.get(i));
                    row.getCell(4).setCellStyle(cellStyle);
                    count++;
                }
            }
            int end = count - 1;
            if (start != end) {
                sheet.addMergedRegion(new CellRangeAddress(start, end, 0, 0));
                sheet.addMergedRegion(new CellRangeAddress(start, end, 1, 1));
                sheet.addMergedRegion(new CellRangeAddress(start, end, 2, 2));
                sheet.addMergedRegion(new CellRangeAddress(start, end, 3, 3));
                sheet.addMergedRegion(new CellRangeAddress(start, end, 5, 5));
                if (!CollectionUtils.isAnyEmpty(item.getReqRespVos())) {
                    for (HashMap<String, Object> map : item.getReqRespVos()) {
                        int num = rowMap.get(String.valueOf(map.get("label")));
                        sheet.addMergedRegion(new CellRangeAddress(start, end, num, num));
                    }
                }
            }

        }
        // 导出 Excel
        try {
            response.setCharacterEncoding("UTF-8");
            response.addHeader("Content-Disposition", "attachment;filename=" + URLEncoder.encode(filename, StandardCharsets.UTF_8.name()));
            response.setContentType("application/vnd.ms-excel;charset=UTF-8");
            template.write(response.getOutputStream());
            response.getOutputStream().close();
            template.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

/*        // 设置 header 和 contentType。写在最后的原因是，避免报错时，响应 contentType 已经被修改了
        response.addHeader("Content-Disposition", "attachment;filename=" + URLEncoder.encode(filename, StandardCharsets.UTF_8.name()));
        response.setContentType("application/vnd.ms-excel;charset=UTF-8");*/
    }
}
