package tech.easyflow.ai.easyagentsflow.llm;

import com.easyagents.core.util.ImageUtil;
import com.easyagents.flow.support.provider.ImageInputResolver;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import tech.easyflow.common.filestorage.FileStorageService;
import tech.easyflow.common.web.exceptions.BusinessException;

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.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.Inet4Address;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Base64;
import java.util.Iterator;
import java.util.Locale;
import java.util.Map;

/**
 * 在模型调用前读取、校验并转换工作流图片输入。
 */
@Component
public class WorkflowImageSourceResolver implements ImageInputResolver {

    static final long MAX_IMAGE_BYTES = 10L * 1024 * 1024;
    static final long MAX_IMAGE_PIXELS = 40_000_000L;
    private static final int MAX_REDIRECTS = 3;
    private static final Duration CONNECT_TIMEOUT = Duration.ofSeconds(8);
    private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(15);
    private static final String DATA_URI_PREFIX = "data:image/";
    private static final String BASE64_MARKER = ";base64,";

    private final FileStorageService fileStorageService;
    private final HttpClient httpClient;

    /**
     * 创建工作流图片解析器。
     *
     * @param fileStorageService 文件存储服务
     */
    @Autowired
    public WorkflowImageSourceResolver(
            @Qualifier("default") FileStorageService fileStorageService) {
        this(fileStorageService, HttpClient.newBuilder()
                .connectTimeout(CONNECT_TIMEOUT)
                .followRedirects(HttpClient.Redirect.NEVER)
                .version(HttpClient.Version.HTTP_1_1)
                .build());
    }

    /**
     * 创建使用指定 HTTP 客户端的图片解析器，供隔离测试使用。
     *
     * @param fileStorageService 文件存储服务
     * @param httpClient HTTP 客户端
     */
    WorkflowImageSourceResolver(FileStorageService fileStorageService, HttpClient httpClient) {
        this.fileStorageService = fileStorageService;
        this.httpClient = httpClient;
    }

    /**
     * 解析图片输入并返回带 MIME 的 Data URI。
     *
     * @param imageInput 图片 URL、文件或结构化图片描述
     * @return 可供模型消费的 Data URI
     * @throws BusinessException 图片不可读取、格式不支持或超过限制时抛出
     */
    @Override
    public String resolve(Object imageInput) {
        if (imageInput instanceof File file) {
            return process(readLocalFile(file));
        }
        if (imageInput instanceof String value) {
            return resolveString(value);
        }
        if (imageInput instanceof Map<?, ?> imageMap) {
            return resolveMap(imageMap);
        }
        throw new BusinessException("图片输入格式不受支持");
    }

    /**
     * 解析字符串形式的旧版图片输入。
     *
     * @param value 图片 URL 或 Data URI
     * @return 规范化 Data URI
     */
    private String resolveString(String value) {
        String normalized = trimToNull(value);
        if (!StringUtils.hasText(normalized)) {
            throw new BusinessException("图片输入不能为空");
        }
        if (normalized.startsWith(DATA_URI_PREFIX)) {
            return process(decodeDataUri(normalized));
        }
        return process(download(normalized));
    }

    /**
     * 解析结构化图片描述。
     *
     * @param imageMap 图片描述
     * @return 规范化 Data URI
     */
    private String resolveMap(Map<?, ?> imageMap) {
        String sourceType = trimObjectToNull(imageMap.get("sourceType"));
        String filePath = trimObjectToNull(imageMap.get("filePath"));
        if (!StringUtils.hasText(sourceType)) {
            sourceType = StringUtils.hasText(filePath) ? "upload" : "url";
        }
        if ("url".equals(sourceType)) {
            return process(download(trimObjectToNull(imageMap.get("url"))));
        }
        if (!"upload".equals(sourceType) && !"resource".equals(sourceType)) {
            throw new BusinessException("图片 sourceType 不受支持");
        }
        if (!StringUtils.hasText(filePath)) {
            throw new BusinessException("图片缺少 filePath");
        }
        try (InputStream input = fileStorageService.readStream(filePath)) {
            return process(readBounded(input));
        } catch (BusinessException exception) {
            throw exception;
        } catch (Exception exception) {
            throw new BusinessException(400, 1, "图片读取失败", exception);
        }
    }

    /**
     * 安全下载外部图片，并在每次重定向后重新校验目标地址。
     *
     * @param value 外部图片 URL
     * @return 图片字节
     */
    private byte[] download(String value) {
        URI current = parseRemoteUri(value);
        for (int redirectCount = 0; redirectCount <= MAX_REDIRECTS; redirectCount++) {
            validateRemoteUri(current);
            HttpRequest request = HttpRequest.newBuilder(current)
                    .timeout(REQUEST_TIMEOUT)
                    .header("Accept", "image/png,image/jpeg,image/webp,image/gif,image/bmp")
                    .GET()
                    .build();
            try {
                HttpResponse<InputStream> response =
                        httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream());
                int status = response.statusCode();
                if (status >= 300 && status < 400) {
                    closeQuietly(response.body());
                    if (redirectCount == MAX_REDIRECTS) {
                        throw new BusinessException("图片 URL 重定向次数超过限制");
                    }
                    String location = response.headers().firstValue("location")
                            .orElseThrow(() -> new BusinessException("图片 URL 重定向缺少目标地址"));
                    current = current.resolve(location);
                    continue;
                }
                if (status < 200 || status >= 300) {
                    closeQuietly(response.body());
                    throw new BusinessException("图片 URL 请求失败，HTTP 状态码: " + status);
                }
                try (InputStream input = response.body()) {
                    return readBounded(input);
                }
            } catch (BusinessException exception) {
                throw exception;
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
                throw new BusinessException(400, 1, "图片 URL 请求被中断", exception);
            } catch (Exception exception) {
                throw new BusinessException(400, 1, "图片 URL 请求失败", exception);
            }
        }
        throw new BusinessException("图片 URL 请求失败");
    }

    /**
     * 校验远程图片 URI，阻止访问本机、内网和云元数据地址。
     *
     * @param uri 待访问 URI
     * @throws BusinessException URI 不安全时抛出
     */
    void validateRemoteUri(URI uri) {
        if (uri == null
                || (!"http".equalsIgnoreCase(uri.getScheme())
                && !"https".equalsIgnoreCase(uri.getScheme()))) {
            throw new BusinessException("图片 URL 仅支持 HTTP/HTTPS");
        }
        if (uri.getUserInfo() != null) {
            throw new BusinessException("图片 URL 不能包含用户信息");
        }
        String host = trimToNull(uri.getHost());
        if (!StringUtils.hasText(host)) {
            throw new BusinessException("图片 URL 缺少有效主机");
        }
        String normalizedHost = host.toLowerCase(Locale.ROOT);
        if ("localhost".equals(normalizedHost)
                || normalizedHost.endsWith(".localhost")
                || "metadata.google.internal".equals(normalizedHost)) {
            throw new BusinessException("图片 URL 不能访问本机或云元数据地址");
        }
        try {
            InetAddress[] addresses = InetAddress.getAllByName(host);
            if (addresses.length == 0) {
                throw new BusinessException("图片 URL 主机无法解析");
            }
            for (InetAddress address : addresses) {
                if (!isPublicAddress(address)) {
                    throw new BusinessException("图片 URL 不能访问内网或保留地址");
                }
            }
        } catch (BusinessException exception) {
            throw exception;
        } catch (Exception exception) {
            throw new BusinessException(400, 1, "图片 URL 主机解析失败", exception);
        }
    }

    /**
     * 判断解析后的地址是否为允许访问的公网地址。
     *
     * @param address IP 地址
     * @return 是否为公网地址
     */
    private boolean isPublicAddress(InetAddress address) {
        if (address.isAnyLocalAddress()
                || address.isLoopbackAddress()
                || address.isLinkLocalAddress()
                || address.isSiteLocalAddress()
                || address.isMulticastAddress()) {
            return false;
        }
        byte[] bytes = address.getAddress();
        if (address instanceof Inet4Address && bytes.length == 4) {
            int first = bytes[0] & 0xff;
            int second = bytes[1] & 0xff;
            int third = bytes[2] & 0xff;
            int fourth = bytes[3] & 0xff;
            return first != 0
                    && first != 10
                    && first != 127
                    && !(first == 168 && second == 63 && third == 129 && fourth == 16)
                    && !(first == 169 && second == 254)
                    && !(first == 100 && second >= 64 && second <= 127)
                    && !(first == 172 && second >= 16 && second <= 31)
                    && !(first == 192 && (second == 0 || second == 168))
                    && !(first == 198 && (second == 18 || second == 19))
                    && first < 224;
        }
        if (address instanceof Inet6Address && bytes.length == 16) {
            int first = bytes[0] & 0xff;
            return (first & 0xfe) != 0xfc;
        }
        return false;
    }

    /**
     * 校验图片格式、尺寸并规范化不兼容格式。
     *
     * @param source 原始图片字节
     * @return 带 MIME 的 Data URI
     */
    private String process(byte[] source) {
        if (source.length == 0) {
            throw new BusinessException("图片内容为空");
        }
        ImageFormat format = detectFormat(source);
        try {
            Dimensions dimensions = format == ImageFormat.WEBP
                    ? webpDimensions(source)
                    : imageIoDimensions(source);
            validateDimensions(dimensions);

            byte[] normalized = source;
            String mimeType = format.mimeType;
            if (format == ImageFormat.GIF || format == ImageFormat.BMP) {
                BufferedImage image = ImageIO.read(new ByteArrayInputStream(source));
                if (image == null) {
                    throw new BusinessException("图片内容无法解析");
                }
                ByteArrayOutputStream output = new ByteArrayOutputStream();
                if (!ImageIO.write(image, "png", output)) {
                    throw new BusinessException("图片格式转换失败");
                }
                normalized = output.toByteArray();
                mimeType = "image/png";
            }
            if (normalized.length > MAX_IMAGE_BYTES) {
                throw new BusinessException("处理后的图片不能超过 10 MiB");
            }
            return ImageUtil.imageBytesToDataUri(normalized, mimeType);
        } catch (BusinessException exception) {
            throw exception;
        } catch (Exception exception) {
            throw new BusinessException(400, 1, "图片处理失败", exception);
        }
    }

    /**
     * 根据文件签名识别真实图片格式。
     *
     * @param bytes 图片字节
     * @return 图片格式
     */
    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 = ascii(bytes, 0, 6);
            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 new BusinessException("仅支持 PNG、JPG、JPEG、WebP、GIF、BMP 图片");
    }

    /**
     * 使用 ImageIO 读取图片尺寸。
     *
     * @param bytes 图片字节
     * @return 图片尺寸
     * @throws IOException 无法读取图片时抛出
     */
    private Dimensions imageIoDimensions(byte[] bytes) throws IOException {
        try (ImageInputStream input =
                     ImageIO.createImageInputStream(new ByteArrayInputStream(bytes))) {
            Iterator<ImageReader> readers = ImageIO.getImageReaders(input);
            if (!readers.hasNext()) {
                throw new BusinessException("图片内容无法解析");
            }
            ImageReader reader = readers.next();
            try {
                reader.setInput(input, true, true);
                return new Dimensions(reader.getWidth(0), reader.getHeight(0));
            } finally {
                reader.dispose();
            }
        }
    }

    /**
     * 读取 WebP 图片尺寸。
     *
     * @param bytes 图片字节
     * @return 图片尺寸
     */
    private Dimensions webpDimensions(byte[] bytes) {
        int offset = 12;
        while (offset + 8 <= bytes.length) {
            String chunk = ascii(bytes, offset, 4);
            long size = Integer.toUnsignedLong(littleEndianInt(bytes, offset + 4));
            int data = offset + 8;
            if (size > bytes.length - data) {
                break;
            }
            int chunkSize = (int) size;
            if ("VP8X".equals(chunk) && chunkSize >= 10) {
                return new Dimensions(
                        1 + littleEndian24(bytes, data + 4),
                        1 + littleEndian24(bytes, data + 7));
            }
            if ("VP8 ".equals(chunk) && chunkSize >= 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) && chunkSize >= 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]);
                return new Dimensions(
                        1 + ((b1 | b2 << 8) & 0x3fff),
                        1 + (((b2 >> 6) | b3 << 2 | b4 << 10) & 0x3fff));
            }
            long nextOffset = (long) data + chunkSize + (chunkSize & 1);
            if (nextOffset > Integer.MAX_VALUE) {
                break;
            }
            offset = (int) nextOffset;
        }
        throw new BusinessException("WebP 图片内容无法解析");
    }

    /**
     * 校验图片像素数量。
     *
     * @param dimensions 图片尺寸
     */
    private void validateDimensions(Dimensions dimensions) {
        if (dimensions.width <= 0 || dimensions.height <= 0) {
            throw new BusinessException("图片尺寸无效");
        }
        long pixels = (long) dimensions.width * dimensions.height;
        if (pixels > MAX_IMAGE_PIXELS) {
            throw new BusinessException("图片不能超过 4000 万像素");
        }
    }

    /**
     * 读取文件并应用大小上限。
     *
     * @param file 本地文件
     * @return 文件字节
     */
    private byte[] readLocalFile(File file) {
        if (file == null || !file.isFile()) {
            throw new BusinessException("图片文件不存在");
        }
        try (InputStream input = java.nio.file.Files.newInputStream(file.toPath())) {
            return readBounded(input);
        } catch (BusinessException exception) {
            throw exception;
        } catch (Exception exception) {
            throw new BusinessException(400, 1, "图片文件读取失败", exception);
        }
    }

    /**
     * 流式读取图片，并拒绝超过 10 MiB 的内容。
     *
     * @param input 图片输入流
     * @return 图片字节
     * @throws IOException 读取失败时抛出
     */
    private byte[] readBounded(InputStream input) throws IOException {
        if (input == null) {
            throw new BusinessException("图片读取失败");
        }
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        byte[] buffer = new byte[8192];
        long total = 0L;
        int read;
        while ((read = input.read(buffer)) != -1) {
            total += read;
            if (total > MAX_IMAGE_BYTES) {
                throw new BusinessException("单张图片不能超过 10 MiB");
            }
            output.write(buffer, 0, read);
        }
        return output.toByteArray();
    }

    /**
     * 解码并校验图片 Data URI。
     *
     * @param dataUri 图片 Data URI
     * @return 图片字节
     */
    private byte[] decodeDataUri(String dataUri) {
        int markerIndex = dataUri.indexOf(BASE64_MARKER);
        if (markerIndex <= DATA_URI_PREFIX.length()) {
            throw new BusinessException("图片 Data URI 格式不正确");
        }
        String encoded = dataUri.substring(markerIndex + BASE64_MARKER.length());
        long maxEncodedLength = (MAX_IMAGE_BYTES + 2L) / 3L * 4L;
        if (encoded.length() > maxEncodedLength + 2L) {
            throw new BusinessException("单张图片不能超过 10 MiB");
        }
        try {
            byte[] bytes = Base64.getDecoder().decode(encoded);
            if (bytes.length > MAX_IMAGE_BYTES) {
                throw new BusinessException("单张图片不能超过 10 MiB");
            }
            return bytes;
        } catch (IllegalArgumentException exception) {
            throw new BusinessException(400, 1, "图片 Data URI 编码无效", exception);
        }
    }

    /**
     * 解析远程 URI。
     *
     * @param value 原始 URL
     * @return URI
     */
    private URI parseRemoteUri(String value) {
        if (!StringUtils.hasText(value)) {
            throw new BusinessException("图片 URL 不能为空");
        }
        try {
            return new URI(value.trim());
        } catch (URISyntaxException exception) {
            throw new BusinessException(400, 1, "图片 URL 格式不正确", exception);
        }
    }

    private String trimObjectToNull(Object value) {
        return trimToNull(value == null ? null : String.valueOf(value));
    }

    private String trimToNull(String value) {
        return StringUtils.hasText(value) ? value.trim() : null;
    }

    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 void closeQuietly(InputStream input) {
        if (input == null) {
            return;
        }
        try {
            input.close();
        } catch (IOException ignored) {
            // 响应已失败，关闭异常不覆盖原始业务错误。
        }
    }

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

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

        private final String mimeType;

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