package tech.easyflow.ai.document.support;

import okhttp3.Dns;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Protocol;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okio.Buffer;
import okio.BufferedSource;
import org.junit.Assert;
import org.junit.Test;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetAddress;
import java.net.Proxy;
import java.net.UnknownHostException;
import java.util.List;

/**
 * {@link DocumentInputStreamSupport} 单元测试。
 */
public class DocumentInputStreamSupportTest {

    /**
     * 验证远程文档客户端禁用代理且保留重定向，并让每个目标经过安全 DNS。
     *
     * @throws Exception 地址构造失败时抛出
     */
    @Test
    public void shouldRejectPrivateAddressForRedirectAwareClient()
            throws Exception {
        Dns privateDns = hostname -> List.of(
                InetAddress.getByName("127.0.0.1"));
        OkHttpClient client =
                DocumentInputStreamSupport.createRemoteClient(privateDns);

        Assert.assertEquals(Proxy.NO_PROXY, client.proxy());
        Assert.assertTrue(client.followRedirects());
        try {
            client.dns().lookup("redirect-target.example");
            Assert.fail("expected private address rejection");
        } catch (UnknownHostException error) {
            Assert.assertTrue(error.getMessage().contains("非公网目标"));
        }
    }

    /**
     * 验证正常公网地址仍可通过安全 DNS。
     *
     * @throws Exception 地址构造失败时抛出
     */
    @Test
    public void shouldAllowPublicAddress() throws Exception {
        InetAddress publicAddress =
                InetAddress.getByName("93.184.216.34");
        Dns publicDns = hostname -> List.of(publicAddress);
        OkHttpClient client =
                DocumentInputStreamSupport.createRemoteClient(publicDns);

        Assert.assertEquals(
                List.of(publicAddress),
                client.dns().lookup("public.example"));
    }

    /**
     * 验证常见内网、链路本地和保留地址均被拒绝。
     *
     * @throws Exception 地址构造失败时抛出
     */
    @Test
    public void shouldRejectNonPublicAddressRanges() throws Exception {
        List<String> blockedAddresses = List.of(
                "10.0.0.1",
                "100.64.0.1",
                "127.0.0.1",
                "169.254.169.254",
                "172.16.0.1",
                "192.168.0.1",
                "198.18.0.1",
                "fd00::1",
                "2001:db8::1");

        for (String value : blockedAddresses) {
            Assert.assertFalse(
                    value,
                    DocumentInputStreamSupport.isPublicAddress(
                            InetAddress.getByName(value)));
        }
    }

    /**
     * 验证远端响应声明长度虚高时按实际读取量接受小文件。
     *
     * @throws IOException 响应读取失败时抛出
     */
    @Test
    public void shouldUseActualBytesWhenDeclaredLengthExceedsLimit()
            throws IOException {
        long maxBytes = 20L * 1024L * 1024L;
        byte[] content = new byte[58 * 1024];
        Buffer source = new Buffer().write(content);
        ResponseBody responseBody = new ResponseBody() {
            @Override
            public MediaType contentType() {
                return MediaType.parse(
                        "application/vnd.openxmlformats-officedocument"
                                + ".wordprocessingml.document");
            }

            @Override
            public long contentLength() {
                return maxBytes + 1L;
            }

            @Override
            public BufferedSource source() {
                return source;
            }
        };
        Response response = new Response.Builder()
                .request(new Request.Builder()
                        .url("http://127.0.0.1/document.docx")
                        .build())
                .protocol(Protocol.HTTP_1_1)
                .code(200)
                .message("OK")
                .body(responseBody)
                .build();

        try (InputStream inputStream =
                     DocumentInputStreamSupport.openResponse(
                             response,
                             maxBytes)) {
            Assert.assertArrayEquals(
                    content,
                    DocumentInputStreamSupport.readBytes(
                            inputStream,
                            maxBytes));
        }
    }

    /**
     * 验证真实读取量超过限制时携带实际字节数。
     */
    @Test
    public void shouldRejectActualBytesBeyondLimit() {
        long maxBytes = 1024L;
        byte[] content = new byte[(int) maxBytes + 1];

        try {
            DocumentInputStreamSupport.readBytes(
                    new ByteArrayInputStream(content),
                    maxBytes);
            Assert.fail("expected SizeLimitExceededException");
        } catch (DocumentInputStreamSupport.SizeLimitExceededException error) {
            Assert.assertEquals(maxBytes, error.getMaxBytes());
            Assert.assertEquals(maxBytes + 1L, error.getActualBytes());
        } catch (IOException error) {
            Assert.fail("unexpected IOException: " + error.getMessage());
        }
    }

    /**
     * 验证零长度读取符合 InputStream 约定。
     *
     * @throws IOException 读取失败时抛出
     */
    @Test
    public void shouldAllowZeroLengthRead() throws IOException {
        try (InputStream inputStream = DocumentInputStreamSupport.limit(
                new ByteArrayInputStream(new byte[0]),
                1L)) {
            Assert.assertEquals(0, inputStream.read(new byte[0]));
        }
    }
}
