package tech.easyflow.ai.easyagentsflow.llm;

import org.junit.Assert;
import org.junit.Test;
import org.springframework.web.multipart.MultipartFile;
import tech.easyflow.common.filestorage.FileStorageService;
import tech.easyflow.common.web.exceptions.BusinessException;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.InputStream;
import java.net.URI;
import java.net.http.HttpClient;
import java.util.Map;

/**
 * 工作流图片源解析器测试。
 */
public class WorkflowImageSourceResolverTest {

    /**
     * 验证存储中的 PNG 图片转换为完整 Data URI。
     *
     * @throws Exception 图片构造失败
     */
    @Test
    public void shouldResolveStoredPngToDataUri() throws Exception {
        byte[] png = imageBytes("png");
        WorkflowImageSourceResolver resolver = resolver(Map.of("/images/a.png", png));

        String dataUri = resolver.resolve(Map.of(
                "sourceType", "upload",
                "fileName", "a.png",
                "filePath", "/images/a.png"));

        Assert.assertTrue(dataUri.startsWith("data:image/png;base64,"));
        Assert.assertArrayEquals(
                png,
                java.util.Base64.getDecoder().decode(dataUri.substring(dataUri.indexOf(',') + 1)));
    }

    /**
     * 验证 BMP 图片会在模型调用前规范化为 PNG。
     *
     * @throws Exception 图片构造失败
     */
    @Test
    public void shouldNormalizeBmpToPng() throws Exception {
        WorkflowImageSourceResolver resolver =
                resolver(Map.of("/images/a.bmp", imageBytes("bmp")));

        String dataUri = resolver.resolve(Map.of(
                "sourceType", "resource",
                "fileName", "a.bmp",
                "filePath", "/images/a.bmp"));

        Assert.assertTrue(dataUri.startsWith("data:image/png;base64,"));
    }

    /**
     * 验证旧版 Data URI 会经过真实图片校验后继续使用。
     *
     * @throws Exception 图片构造失败
     */
    @Test
    public void shouldValidateLegacyDataUri() throws Exception {
        byte[] png = imageBytes("png");
        String input = "data:image/png;base64,"
                + java.util.Base64.getEncoder().encodeToString(png);

        String dataUri = resolver(Map.of()).resolve(input);

        Assert.assertTrue(dataUri.startsWith("data:image/png;base64,"));
    }

    /**
     * 验证本机、内网和云元数据地址会被拦截。
     */
    @Test
    public void shouldRejectUnsafeRemoteAddresses() {
        WorkflowImageSourceResolver resolver = resolver(Map.of());

        assertUnsafe(resolver, "http://127.0.0.1/image.png");
        assertUnsafe(resolver, "http://192.168.1.2/image.png");
        assertUnsafe(resolver, "http://168.63.129.16/metadata/instance");
        assertUnsafe(resolver, "http://169.254.169.254/latest/meta-data");
        assertUnsafe(resolver, "http://metadata.google.internal/image.png");
    }

    private static void assertUnsafe(WorkflowImageSourceResolver resolver, String value) {
        try {
            resolver.validateRemoteUri(URI.create(value));
            Assert.fail("expected BusinessException for " + value);
        } catch (BusinessException expected) {
            Assert.assertTrue(expected.getMessage().contains("不能访问"));
        }
    }

    private static WorkflowImageSourceResolver resolver(Map<String, byte[]> files) {
        return new WorkflowImageSourceResolver(
                new InMemoryStorage(files),
                HttpClient.newHttpClient());
    }

    private static byte[] imageBytes(String format) throws Exception {
        BufferedImage image = new BufferedImage(2, 2, BufferedImage.TYPE_INT_RGB);
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        Assert.assertTrue(ImageIO.write(image, format, output));
        return output.toByteArray();
    }

    /**
     * 测试用内存文件存储。
     */
    private static final class InMemoryStorage implements FileStorageService {

        private final Map<String, byte[]> files;

        private InMemoryStorage(Map<String, byte[]> files) {
            this.files = files;
        }

        @Override
        public String save(MultipartFile file) {
            throw new UnsupportedOperationException();
        }

        @Override
        public void delete(String path) {
            throw new UnsupportedOperationException();
        }

        @Override
        public String save(File file, String prePath) {
            throw new UnsupportedOperationException();
        }

        @Override
        public InputStream readStream(String path) {
            byte[] bytes = files.get(path);
            if (bytes == null) {
                throw new IllegalArgumentException("missing file: " + path);
            }
            return new ByteArrayInputStream(bytes);
        }

        @Override
        public long getFileSize(String path) {
            byte[] bytes = files.get(path);
            return bytes == null ? 0 : bytes.length;
        }
    }
}
