package com.doumee.core.annotation.excel;

import org.apache.poi.hssf.usermodel.*;
import org.apache.poi.ooxml.POIXMLDocumentPart;
import org.apache.poi.ss.usermodel.*;
import org.apache.poi.xssf.usermodel.*;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;

/**
 * 读取 Excel 工作表中按行列锚定的内嵌图片
 */
public class ExcelEmbeddedImageReader {

    private ExcelEmbeddedImageReader() {
    }

    public static Map<Integer, RowEmbeddedImages> read(String filePath, int sheetIndex, int headerNum,
                                                       int mainImageCol, int multiImageCol) throws IOException {
        try (InputStream in = new FileInputStream(filePath)) {
            return read(in, sheetIndex, headerNum, mainImageCol, multiImageCol);
        }
    }

    public static Map<Integer, RowEmbeddedImages> read(InputStream in, int sheetIndex, int headerNum,
                                                     int mainImageCol, int multiImageCol) throws IOException {
        try (Workbook workbook = WorkbookFactory.create(in)) {
            Sheet sheet = workbook.getSheetAt(sheetIndex);
            if (sheet instanceof XSSFSheet) {
                return readXssf((XSSFSheet) sheet, headerNum, mainImageCol, multiImageCol);
            }
            if (sheet instanceof HSSFSheet) {
                return readHssf((HSSFSheet) sheet, headerNum, mainImageCol, multiImageCol);
            }
            return Collections.emptyMap();
        }
    }

    private static Map<Integer, RowEmbeddedImages> readXssf(XSSFSheet sheet, int headerNum,
                                                            int mainImageCol, int multiImageCol) {
        Map<Integer, List<CellPicture>> cellPictures = new HashMap<>();
        collectXssfDrawings(sheet, cellPictures);
        return buildRowImages(cellPictures, headerNum, mainImageCol, multiImageCol);
    }

    private static void collectXssfDrawings(XSSFSheet sheet, Map<Integer, List<CellPicture>> cellPictures) {
        XSSFDrawing drawing = sheet.getDrawingPatriarch();
        if (drawing != null) {
            appendXssfPictures(drawing, cellPictures);
        }
        for (POIXMLDocumentPart part : sheet.getRelations()) {
            if (part instanceof XSSFDrawing && part != drawing) {
                appendXssfPictures((XSSFDrawing) part, cellPictures);
            }
        }
    }

    private static void appendXssfPictures(XSSFDrawing drawing, Map<Integer, List<CellPicture>> cellPictures) {
        if (drawing == null) {
            return;
        }
        for (XSSFShape shape : drawing.getShapes()) {
            if (!(shape instanceof XSSFPicture)) {
                continue;
            }
            XSSFPicture picture = (XSSFPicture) shape;
            ClientAnchor anchor = picture.getClientAnchor();
            if (anchor == null) {
                continue;
            }
            PictureData pictureData = picture.getPictureData();
            if (pictureData == null || pictureData.getData() == null || pictureData.getData().length == 0) {
                continue;
            }
            addCellPicture(cellPictures, anchor, pictureData);
        }
    }

    private static Map<Integer, RowEmbeddedImages> readHssf(HSSFSheet sheet, int headerNum,
                                                            int mainImageCol, int multiImageCol) {
        Map<Integer, List<CellPicture>> cellPictures = new HashMap<>();
        HSSFPatriarch patriarch = sheet.getDrawingPatriarch();
        if (patriarch != null) {
            for (HSSFShape shape : patriarch.getChildren()) {
                if (!(shape instanceof HSSFPicture)) {
                    continue;
                }
                HSSFPicture picture = (HSSFPicture) shape;
                ClientAnchor anchor = picture.getClientAnchor();
                if (anchor == null) {
                    continue;
                }
                PictureData pictureData = picture.getPictureData();
                if (pictureData == null || pictureData.getData() == null || pictureData.getData().length == 0) {
                    continue;
                }
                addCellPicture(cellPictures, anchor, pictureData);
            }
        }
        return buildRowImages(cellPictures, headerNum, mainImageCol, multiImageCol);
    }

    private static void addCellPicture(Map<Integer, List<CellPicture>> cellPictures, ClientAnchor anchor,
                                       PictureData pictureData) {
        int sortKey = anchorSortKey(anchor);
        String ext = normalizeExtension(pictureData.suggestFileExtension());
        byte[] data = pictureData.getData();
        int minRow = anchor.getRow1();
        int minCol = anchor.getCol1();
        int key = cellKey(minRow, minCol);
        cellPictures.computeIfAbsent(key, k -> new ArrayList<>())
                .add(new CellPicture(minRow, minCol, data, ext, sortKey));
    }

    private static Map<Integer, RowEmbeddedImages> buildRowImages(Map<Integer, List<CellPicture>> cellPictures,
                                                                 int headerNum, int mainImageCol, int multiImageCol) {
        Map<Integer, RowEmbeddedImages> result = new HashMap<>();
        int dataStartRow = headerNum + 1;
        Set<Integer> rows = new TreeSet<>();
        for (CellPicture pic : flatten(cellPictures)) {
            if (pic.row >= dataStartRow) {
                rows.add(pic.row);
            }
        }
        for (int row : rows) {
            RowEmbeddedImages rowImages = new RowEmbeddedImages();
            List<CellPicture> mainList = cellPictures.getOrDefault(cellKey(row, mainImageCol), Collections.emptyList());
            mainList.stream()
                    .sorted(Comparator.comparingInt(CellPicture::getSortKey))
                    .findFirst()
                    .ifPresent(pic -> rowImages.setMainImage(toItem(pic)));

            List<RowEmbeddedImages.EmbeddedImageItem> multi = new ArrayList<>();
            for (CellPicture pic : cellPictures.getOrDefault(cellKey(row, multiImageCol), Collections.emptyList())) {
                multi.add(toItem(pic));
            }
            multi.sort(Comparator.comparingInt(RowEmbeddedImages.EmbeddedImageItem::getSortKey));
            rowImages.setMultiImages(multi);
            result.put(row, rowImages);
        }
        return result;
    }

    private static List<CellPicture> flatten(Map<Integer, List<CellPicture>> cellPictures) {
        List<CellPicture> all = new ArrayList<>();
        for (List<CellPicture> list : cellPictures.values()) {
            all.addAll(list);
        }
        return all;
    }

    private static RowEmbeddedImages.EmbeddedImageItem toItem(CellPicture pic) {
        return new RowEmbeddedImages.EmbeddedImageItem(pic.data, pic.extension, pic.sortKey);
    }

    private static int cellKey(int row, int col) {
        return row * 1000 + col;
    }

    private static int anchorSortKey(ClientAnchor anchor) {
        return anchor.getRow1() * 1_000_000 + anchor.getCol1() * 10_000 + anchor.getDx1() + anchor.getDy1();
    }

    private static String normalizeExtension(String ext) {
        if (ext == null || ext.isEmpty()) {
            return ".jpg";
        }
        return ext.startsWith(".") ? ext : "." + ext;
    }

    private static class CellPicture {
        private final int row;
        private final int col;
        private final byte[] data;
        private final String extension;
        private final int sortKey;

        private CellPicture(int row, int col, byte[] data, String extension, int sortKey) {
            this.row = row;
            this.col = col;
            this.data = data;
            this.extension = extension;
            this.sortKey = sortKey;
        }

        private int getSortKey() {
            return sortKey;
        }
    }
}