package com.yaoai.ai.providers.config;

import org.apache.hc.client5.http.DnsResolver;
import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
import org.apache.hc.client5.http.impl.classic.HttpClients;
import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManagerBuilder;
import org.springframework.http.client.HttpComponentsClientHttpRequestFactory;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestClient;

import java.net.InetAddress;
import java.net.URI;
import java.net.UnknownHostException;
import java.time.Duration;

@Component
public class ProviderRestClientFactory {

    private static final String AGNES_API_HOST = "apihub.agnes-ai.com";
    private static final String AGNES_API_RESOLVE_IP = "104.18.18.62";

    public RestClient create(ResolvedModelConfig config) {
        return create(config, Duration.ofSeconds(10), Duration.ofMinutes(2));
    }

    public RestClient create(ResolvedModelConfig config, Duration connectTimeout, Duration readTimeout) {
        RestClient.Builder builder = RestClient.builder()
                .baseUrl(config.baseUrl())
                .defaultHeader("Content-Type", "application/json");

        String resolveIp = resolveIp(config);
        if (StringUtils.hasText(resolveIp)) {
            builder.requestFactory(requestFactory(config, resolveIp, connectTimeout, readTimeout));
        }

        return builder.build();
    }

    private HttpComponentsClientHttpRequestFactory requestFactory(ResolvedModelConfig config,
                                                                  String resolveIp,
                                                                  Duration connectTimeout,
                                                                  Duration readTimeout) {
        String host = URI.create(config.baseUrl()).getHost();
        DnsResolver resolver = new FixedHostDnsResolver(host, resolveIp.trim());
        CloseableHttpClient httpClient = HttpClients.custom()
                .setConnectionManager(PoolingHttpClientConnectionManagerBuilder.create()
                        .setDnsResolver(resolver)
                        .build())
                .build();
        HttpComponentsClientHttpRequestFactory factory = new HttpComponentsClientHttpRequestFactory(httpClient);
        factory.setConnectTimeout(connectTimeout);
        factory.setConnectionRequestTimeout(connectTimeout);
        return factory;
    }

    private String resolveIp(ResolvedModelConfig config) {
        if (StringUtils.hasText(config.resolveIp())) {
            return config.resolveIp();
        }
        String host = URI.create(config.baseUrl()).getHost();
        if (AGNES_API_HOST.equalsIgnoreCase(host)) {
            return AGNES_API_RESOLVE_IP;
        }
        return null;
    }

    private record FixedHostDnsResolver(String host, String resolveIp) implements DnsResolver {
        @Override
        public InetAddress[] resolve(String requestedHost) throws UnknownHostException {
            if (host != null && host.equalsIgnoreCase(requestedHost)) {
                return new InetAddress[]{InetAddress.getByName(resolveIp)};
            }
            return InetAddress.getAllByName(requestedHost);
        }

        @Override
        public String resolveCanonicalHostname(String requestedHost) throws UnknownHostException {
            return requestedHost;
        }
    }
}
