package tech.easyflow.agent.runtime.media;

import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.server.ResponseStatusException;
import tech.easyflow.agent.config.AgentMediaProperties;

import javax.imageio.ImageIO;
import javax.imageio.ImageReader;
import javax.imageio.stream.ImageInputStream;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.HexFormat;
import java.util.Iterator;

/**
 * Agent 聊天图片格式识别、尺寸校验与规范化处理器。
 */
@Component
public class AgentImageProcessor {

    private final AgentMediaProperties properties;

    /**
     * 创建图片处理器。
     *
     * @param properties Agent 媒体配置
     */
    public AgentImageProcessor(AgentMediaProperties properties) {
        this.properties = properties;
    }

    /**
     * 校验并规范化上传图片。
     *
     * @param file 上传文件
     * @return 规范化图片
     */
    public AgentImageData process(MultipartFile file) {
        if (file == null || file.isEmpty()) {
            throw badRequest("请选择要上传的图片");
        }
        if (file.getSize() > properties.getMaxImageBytes()) {
            throw badRequest("单张图片不能超过 10 MiB");
        }
        try {
            byte[] source = file.getBytes();
            ImageFormat format = detectFormat(source);
            Dimensions dimensions = format == ImageFormat.WEBP
                    ? webpDimensions(source)
                    : imageIoDimensions(source);
            validateDimensions(dimensions);
            byte[] normalized = source;
            String mimeType = format.mimeType;
            String extension = format.extension;
            // 动图仅保留首帧，BMP 转 PNG，避免模型端格式兼容差异。
            if (format == ImageFormat.GIF || format == ImageFormat.BMP) {
                BufferedImage image = ImageIO.read(new ByteArrayInputStream(source));
                if (image == null) {
                    throw badRequest("图片内容无法解析");
                }
                ByteArrayOutputStream output = new ByteArrayOutputStream();
                if (!ImageIO.write(image, "png", output)) {
                    throw badRequest("图片格式转换失败");
                }
                normalized = output.toByteArray();
                mimeType = "image/png";
                extension = "png";
            }
            if (normalized.length > properties.getMaxImageBytes()) {
                throw badRequest("处理后的图片不能超过 10 MiB");
            }
            return new AgentImageData(normalized, mimeType, extension,
                    dimensions.width, dimensions.height, sha256(normalized));
        } catch (ResponseStatusException error) {
            throw error;
        } catch (Exception error) {
            throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "图片处理失败", error);
        }
    }

    private ImageFormat detectFormat(byte[] bytes) {
        if (bytes.length >= 8
                && bytes[0] == (byte) 0x89 && bytes[1] == 0x50 && bytes[2] == 0x4e && bytes[3] == 0x47) {
            return ImageFormat.PNG;
        }
        if (bytes.length >= 3 && bytes[0] == (byte) 0xff && bytes[1] == (byte) 0xd8 && bytes[2] == (byte) 0xff) {
            return ImageFormat.JPEG;
        }
        if (bytes.length >= 6) {
            String header = new String(bytes, 0, 6, StandardCharsets.US_ASCII);
            if ("GIF87a".equals(header) || "GIF89a".equals(header)) {
                return ImageFormat.GIF;
            }
        }
        if (bytes.length >= 2 && bytes[0] == 'B' && bytes[1] == 'M') {
            return ImageFormat.BMP;
        }
        if (bytes.length >= 12
                && "RIFF".equals(ascii(bytes, 0, 4))
                && "WEBP".equals(ascii(bytes, 8, 4))) {
            return ImageFormat.WEBP;
        }
        throw badRequest("仅支持 PNG、JPG、JPEG、WebP、GIF、BMP 图片");
    }

    private Dimensions imageIoDimensions(byte[] bytes) throws Exception {
        try (ImageInputStream input = ImageIO.createImageInputStream(new ByteArrayInputStream(bytes))) {
            Iterator<ImageReader> readers = ImageIO.getImageReaders(input);
            if (!readers.hasNext()) {
                throw badRequest("图片内容无法解析");
            }
            ImageReader reader = readers.next();
            try {
                reader.setInput(input, true, true);
                return new Dimensions(reader.getWidth(0), reader.getHeight(0));
            } finally {
                reader.dispose();
            }
        }
    }

    private Dimensions webpDimensions(byte[] bytes) {
        int offset = 12;
        while (offset + 8 <= bytes.length) {
            String chunk = ascii(bytes, offset, 4);
            int size = littleEndianInt(bytes, offset + 4);
            int data = offset + 8;
            if (size < 0 || data + size > bytes.length) {
                break;
            }
            if ("VP8X".equals(chunk) && size >= 10) {
                return new Dimensions(1 + littleEndian24(bytes, data + 4),
                        1 + littleEndian24(bytes, data + 7));
            }
            if ("VP8 ".equals(chunk) && size >= 10
                    && bytes[data + 3] == (byte) 0x9d && bytes[data + 4] == 0x01 && bytes[data + 5] == 0x2a) {
                return new Dimensions(littleEndian16(bytes, data + 6) & 0x3fff,
                        littleEndian16(bytes, data + 8) & 0x3fff);
            }
            if ("VP8L".equals(chunk) && size >= 5 && bytes[data] == 0x2f) {
                int b1 = unsigned(bytes[data + 1]);
                int b2 = unsigned(bytes[data + 2]);
                int b3 = unsigned(bytes[data + 3]);
                int b4 = unsigned(bytes[data + 4]);
                int width = 1 + ((b1 | (b2 << 8)) & 0x3fff);
                int height = 1 + (((b2 >> 6) | (b3 << 2) | (b4 << 10)) & 0x3fff);
                return new Dimensions(width, height);
            }
            offset = data + size + (size & 1);
        }
        throw badRequest("WebP 图片内容无法解析");
    }

    private void validateDimensions(Dimensions dimensions) {
        if (dimensions.width <= 0 || dimensions.height <= 0) {
            throw badRequest("图片尺寸无效");
        }
        long pixels = (long) dimensions.width * dimensions.height;
        if (pixels > properties.getMaxImagePixels()) {
            throw badRequest("图片像素过大，请压缩后重试");
        }
    }

    private String sha256(byte[] bytes) throws Exception {
        return HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(bytes));
    }

    private String ascii(byte[] bytes, int offset, int length) {
        return new String(bytes, offset, length, StandardCharsets.US_ASCII);
    }

    private int littleEndian16(byte[] bytes, int offset) {
        return unsigned(bytes[offset]) | (unsigned(bytes[offset + 1]) << 8);
    }

    private int littleEndian24(byte[] bytes, int offset) {
        return unsigned(bytes[offset]) | (unsigned(bytes[offset + 1]) << 8) | (unsigned(bytes[offset + 2]) << 16);
    }

    private int littleEndianInt(byte[] bytes, int offset) {
        return unsigned(bytes[offset]) | (unsigned(bytes[offset + 1]) << 8)
                | (unsigned(bytes[offset + 2]) << 16) | (unsigned(bytes[offset + 3]) << 24);
    }

    private int unsigned(byte value) { return value & 0xff; }

    private ResponseStatusException badRequest(String message) {
        return new ResponseStatusException(HttpStatus.BAD_REQUEST, message);
    }

    private record Dimensions(int width, int height) { }

    private enum ImageFormat {
        PNG("image/png", "png"),
        JPEG("image/jpeg", "jpg"),
        WEBP("image/webp", "webp"),
        GIF("image/gif", "gif"),
        BMP("image/bmp", "bmp");

        private final String mimeType;
        private final String extension;

        ImageFormat(String mimeType, String extension) {
            this.mimeType = mimeType;
            this.extension = extension;
        }
    }
}
