package tech.easyflow.ai.document.support;

import com.easyagents.flow.core.util.OkHttpClientUtil;
import okhttp3.Dns;
import okhttp3.HttpUrl;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

/**
 * 文档输入流的共享连接与实际字节数保护工具。
 */
public final class DocumentInputStreamSupport {

    private static final Logger LOG =
            LoggerFactory.getLogger(DocumentInputStreamSupport.class);
    private static final String BLOCKED_REMOTE_ADDRESS_MESSAGE =
            "远端文档地址不允许访问非公网目标";

    private DocumentInputStreamSupport() {
    }

    /**
     * 使用共享 HTTP 客户端打开远端文档流。
     *
     * @param url 远端 URL
     * @param maxBytes 最大允许读取字节数；小于等于 0 时不限制
     * @return 关闭时会同步释放 HTTP 响应的输入流
     * @throws IOException 请求失败、非成功状态或响应超过限制时抛出
     */
    public static InputStream openRemote(String url, long maxBytes) throws IOException {
        HttpUrl remoteUrl = parseRemoteUrl(url);
        Response response = RemoteClientHolder.CLIENT
                .newCall(new Request.Builder().url(remoteUrl).get().build())
                .execute();
        return openResponse(response, maxBytes);
    }

    /**
     * 解析并校验远端文档 URL 的协议。
     *
     * @param url 原始 URL
     * @return 可用于请求的 HTTP URL
     * @throws IOException URL 为空、格式错误或协议不受支持时抛出
     */
    private static HttpUrl parseRemoteUrl(String url) throws IOException {
        if (url == null || url.isBlank()) {
            throw new IOException("远端文档 URL 不能为空");
        }
        HttpUrl remoteUrl = HttpUrl.parse(url);
        if (remoteUrl == null
                || (!"http".equals(remoteUrl.scheme())
                && !"https".equals(remoteUrl.scheme()))) {
            throw new IOException("远端文档仅支持 HTTP 或 HTTPS URL");
        }
        return remoteUrl;
    }

    /**
     * 创建只允许公网目标的远程文档客户端。
     *
     * <p>显式禁用代理，避免代理服务器重新解析目标域名后绕过本机 DNS 校验。
     * OkHttp 的重定向请求会继续使用同一 DNS 实现，因此每个新目标都会重新校验。</p>
     *
     * @param delegateDns 实际执行域名解析的 DNS
     * @return 带公网地址约束的 HTTP 客户端
     */
    static OkHttpClient createRemoteClient(Dns delegateDns) {
        if (delegateDns == null) {
            throw new IllegalArgumentException("DNS 解析器不能为空");
        }
        return OkHttpClientUtil.buildDefaultClient()
                .newBuilder()
                .proxy(Proxy.NO_PROXY)
                .dns(new PublicAddressDns(delegateDns))
                .build();
    }

    /**
     * 判断地址是否属于允许访问的公网范围。
     *
     * @param address 已解析的目标地址
     * @return 公网单播地址返回 true
     */
    static boolean isPublicAddress(InetAddress address) {
        if (address == null
                || address.isAnyLocalAddress()
                || address.isLoopbackAddress()
                || address.isLinkLocalAddress()
                || address.isSiteLocalAddress()
                || address.isMulticastAddress()) {
            return false;
        }
        byte[] bytes = address.getAddress();
        if (bytes.length == 4) {
            return isPublicIpv4(bytes);
        }
        if (bytes.length == 16) {
            return isPublicIpv6(bytes);
        }
        return false;
    }

    /**
     * 判断 IPv4 地址是否为公网单播地址。
     *
     * @param bytes IPv4 地址字节
     * @return 公网单播地址返回 true
     */
    private static boolean isPublicIpv4(byte[] bytes) {
        int first = bytes[0] & 0xff;
        int second = bytes[1] & 0xff;
        int third = bytes[2] & 0xff;
        if (first == 0 || first == 10 || first == 127 || first >= 224) {
            return false;
        }
        if (first == 100 && second >= 64 && second <= 127) {
            return false;
        }
        if (first == 169 && second == 254) {
            return false;
        }
        if (first == 172 && second >= 16 && second <= 31) {
            return false;
        }
        if (first == 192 && second == 168) {
            return false;
        }
        if (first == 192 && second == 0 && (third == 0 || third == 2)) {
            return false;
        }
        if (first == 192 && second == 88 && third == 99) {
            return false;
        }
        if (first == 198 && (second == 18 || second == 19)) {
            return false;
        }
        if (first == 198 && second == 51 && third == 100) {
            return false;
        }
        return !(first == 203 && second == 0 && third == 113);
    }

    /**
     * 判断 IPv6 地址是否为公网单播地址。
     *
     * @param bytes IPv6 地址字节
     * @return 公网单播地址返回 true
     */
    private static boolean isPublicIpv6(byte[] bytes) {
        int first = bytes[0] & 0xff;
        int second = bytes[1] & 0xff;
        if ((first & 0xfe) == 0xfc) {
            return false;
        }
        if (first == 0x20 && second == 0x01) {
            int third = bytes[2] & 0xff;
            int fourth = bytes[3] & 0xff;
            if ((third == 0x0d && fourth == 0xb8)
                    || (third == 0x00 && fourth == 0x00)) {
                return false;
            }
        }
        if (first == 0x20 && second == 0x02) {
            return false;
        }
        return !(first == 0x00
                && second == 0x64
                && (bytes[2] & 0xff) == 0xff
                && (bytes[3] & 0xff) == 0x9b);
    }

    /**
     * 将 HTTP 响应转换为按实际读取字节数受限的输入流。
     *
     * <p>响应声明长度仅用于诊断。最终限制以流中实际读到的字节数为准，
     * 避免代理或存储服务返回错误 Content-Length 时误判正常文件。</p>
     *
     * @param response HTTP 响应
     * @param maxBytes 最大允许读取字节数；小于等于 0 时不限制
     * @return 关闭时会同步释放 HTTP 响应的输入流
     * @throws IOException 非成功状态或响应体为空时抛出
     */
    static InputStream openResponse(Response response, long maxBytes)
            throws IOException {
        if (!response.isSuccessful()) {
            int status = response.code();
            response.close();
            throw new IOException("Document download failed with HTTP status " + status);
        }
        ResponseBody body = response.body();
        if (body == null) {
            response.close();
            throw new IOException("Document download response body is empty");
        }
        long contentLength = body.contentLength();
        if (maxBytes > 0L && contentLength > maxBytes) {
            LOG.warn(
                    "远端文档响应声明长度超过限制，将按实际读取量确认: "
                            + "declaredBytes={}, maxBytes={}",
                    contentLength,
                    maxBytes);
        }
        return limit(new ResponseInputStream(body.byteStream(), response), maxBytes);
    }

    /**
     * 为已有输入流增加实际读取字节数限制。
     *
     * @param inputStream 原始输入流
     * @param maxBytes 最大允许读取字节数；小于等于 0 时直接返回原始流
     * @return 受限输入流
     */
    public static InputStream limit(InputStream inputStream, long maxBytes) {
        if (inputStream == null || maxBytes <= 0L) {
            return inputStream;
        }
        return new LimitedInputStream(inputStream, maxBytes);
    }

    /**
     * 在实际字节数限制下读取完整内容。
     *
     * @param inputStream 文档输入流
     * @param maxBytes 最大允许读取字节数；小于等于 0 时不限制
     * @return 文档字节
     * @throws IOException 读取失败或内容超过限制时抛出
     */
    public static byte[] readBytes(InputStream inputStream, long maxBytes) throws IOException {
        try (InputStream limited = limit(inputStream, maxBytes)) {
            return limited.readAllBytes();
        }
    }

    /**
     * 文档实际读取字节数超过保护值异常。
     */
    public static final class SizeLimitExceededException extends IOException {

        private final long maxBytes;
        private final long actualBytes;

        /**
         * 创建文档大小限制异常。
         *
         * @param maxBytes 最大允许字节数
         */
        public SizeLimitExceededException(long maxBytes) {
            this(maxBytes, -1L);
        }

        /**
         * 创建包含实际读取量的文档大小限制异常。
         *
         * @param maxBytes 最大允许字节数
         * @param actualBytes 已确认的实际字节数；未知时传入负数
         */
        public SizeLimitExceededException(long maxBytes, long actualBytes) {
            super("Document exceeds size limit: max=" + maxBytes
                    + " bytes, actual=" + actualBytes + " bytes");
            this.maxBytes = maxBytes;
            this.actualBytes = actualBytes;
        }

        /**
         * 获取最大允许字节数。
         *
         * @return 最大允许字节数
         */
        public long getMaxBytes() {
            return maxBytes;
        }

        /**
         * 获取已确认的实际读取量。
         *
         * @return 实际字节数；未知时为负数
         */
        public long getActualBytes() {
            return actualBytes;
        }
    }

    /**
     * 按实际读取量执行限制的输入流。
     */
    private static final class LimitedInputStream extends FilterInputStream {

        private final long maxBytes;
        private long consumed;

        private LimitedInputStream(InputStream inputStream, long maxBytes) {
            super(inputStream);
            this.maxBytes = maxBytes;
        }

        @Override
        public int read() throws IOException {
            int value = super.read();
            if (value >= 0) {
                recordRead(1L);
            }
            return value;
        }

        @Override
        public int read(byte[] buffer, int offset, int length) throws IOException {
            if (length == 0) {
                return 0;
            }
            int allowed = (int) Math.min(
                    Math.max(0L, maxBytes - consumed + 1L),
                    (long) length);
            if (allowed <= 0) {
                throw new SizeLimitExceededException(maxBytes, consumed + 1L);
            }
            int count = super.read(buffer, offset, allowed);
            if (count > 0) {
                recordRead(count);
            }
            return count;
        }

        private void recordRead(long count) throws SizeLimitExceededException {
            consumed += count;
            if (consumed > maxBytes) {
                throw new SizeLimitExceededException(maxBytes, consumed);
            }
        }
    }

    /**
     * 关闭输入流时一并关闭 OkHttp 响应。
     */
    private static final class ResponseInputStream extends FilterInputStream {

        private final Response response;

        private ResponseInputStream(InputStream inputStream, Response response) {
            super(inputStream);
            this.response = response;
        }

        @Override
        public void close() {
            response.close();
        }
    }

    /**
     * 延迟创建远程文档客户端，避免本地流处理路径初始化网络资源。
     */
    private static final class RemoteClientHolder {

        private static final OkHttpClient CLIENT = createRemoteClient(Dns.SYSTEM);

        private RemoteClientHolder() {
        }
    }

    /**
     * 对 DNS 解析结果执行公网地址约束。
     */
    private static final class PublicAddressDns implements Dns {

        private final Dns delegate;

        /**
         * 创建安全 DNS 包装器。
         *
         * @param delegate 实际 DNS 解析器
         */
        private PublicAddressDns(Dns delegate) {
            this.delegate = delegate;
        }

        /**
         * 解析主机并拒绝任一非公网地址，避免连接回退到私有地址。
         *
         * @param hostname 目标主机名
         * @return 全部通过校验的解析地址
         * @throws UnknownHostException 解析失败或包含非公网地址时抛出
         */
        @Override
        public List<InetAddress> lookup(String hostname) throws UnknownHostException {
            List<InetAddress> addresses = delegate.lookup(hostname);
            if (addresses == null || addresses.isEmpty()) {
                throw new UnknownHostException("远端文档域名没有可用地址");
            }
            for (InetAddress address : addresses) {
                if (!isPublicAddress(address)) {
                    throw new UnknownHostException(BLOCKED_REMOTE_ADDRESS_MESSAGE);
                }
            }
            return addresses;
        }
    }
}
