package com.seer.teach.config;

import com.alibaba.cloud.nacos.balancer.NacosBalancer;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients;
import org.springframework.cloud.loadbalancer.core.ReactorLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import reactor.core.publisher.Mono;

import java.net.InetAddress;
import java.net.NetworkInterface;
import java.net.SocketException;
import java.net.URI;
import java.util.*;

@Slf4j
@LoadBalancerClients(defaultConfiguration = LoadBalancerConfig.class)
@Configuration
public class LoadBalancerConfig {

    @Bean
    public ServiceInstanceListSupplier discoveryClientServiceInstanceListSupplier(
            ConfigurableApplicationContext context,
            DiscoveryClient discoveryClient) {
        return ServiceInstanceListSupplier.builder()
                .withDiscoveryClient()
                .withCaching()
                .build(context);
    }

    @Bean
    public GlobalFilter ipForwardFilter() {
        return (exchange, chain) -> {
            String ip = Optional.ofNullable(exchange.getRequest().getRemoteAddress())
                    .map(addr -> addr.getAddress().getHostAddress())
                    .orElse("unknown");
            log.info("ipForwardFilter 请求来自客户端: {}", ip);
            ServerHttpRequest mutatedRequest = exchange.getRequest().mutate()
                    .header("X-Forwarded-For", ip)
                    .build();

            return chain.filter(exchange.mutate().request(mutatedRequest).build());
        };
    }

    @Bean
    public ReactorLoadBalancer<ServiceInstance> ipAwareLoadBalancer(
            LoadBalancerClientFactory clientFactory,
            ServiceInstanceListSupplier supplier, Environment environment) {
        String serviceId = LoadBalancerClientFactory.getName(environment);
        return new IpAwareLoadBalancer(supplier, serviceId);
    }

    static class IpAwareLoadBalancer implements ReactorServiceInstanceLoadBalancer {

        private final String serviceId;
        private final ServiceInstanceListSupplier supplier;
        private final List<String> localIps;

        public IpAwareLoadBalancer(ServiceInstanceListSupplier supplier, String serviceId) {
            this.supplier = supplier;
            this.serviceId = serviceId;
            localIps = getAllLocalIps();
        }

        @Override
        public Mono<Response<ServiceInstance>> choose(Request request) {
            return supplier.get()
                    .next()
                    .map(instances -> processInstances(instances, request));
        }

        private Response<ServiceInstance> processInstances(List<ServiceInstance> instances, Request request) {
            if (instances.isEmpty()) {
                log.warn("服务 {} 没有可用实例", serviceId);
                return new EmptyResponse();
            }

            // 1. 获取请求的客户端IP（从X-Forwarded-For或直接RemoteAddr）
            String clientIp = extractClientIp(request);
            log.info("请求来自客户端: {}", clientIp);

            // 2. 优先匹配策略（本地IP > 客户端IP匹配 > 随机）
            ServiceInstance chosenInstance = selectBestInstance(instances, clientIp);

            // 3. 记录日志（请求URL + 选中实例）
            logSelectedInstance(request, chosenInstance);

            return new DefaultResponse(chosenInstance);
        }

        /**
         * 选择最优实例（新优先级：客户端IP匹配 > 本地实例 > 随机）
         */
        private ServiceInstance selectBestInstance(List<ServiceInstance> instances, String clientIp) {
            // 1. 优先匹配客户端IP相同的实例
            if (clientIp != null) {
                Optional<ServiceInstance> clientIpMatched = instances.stream()
                        .filter(instance -> clientIp.equals(instance.getHost()))
                        .findFirst();
                if (clientIpMatched.isPresent()) {
                    log.debug("根据客户端IP匹配选择实例: {}", clientIp);
                    return clientIpMatched.get();
                }
            }

            // 2. 次选本地实例
            Optional<ServiceInstance> localInstance = findLocalInstance(instances);
            if (localInstance.isPresent()) {
                log.debug("选择本地实例: {}", localInstance.get().getHost());
                return localInstance.get();
            }

            // 3. 默认NacosBalancer
            return NacosBalancer.getHostByRandomWeight3(instances);
        }

        /**
         * 从Request中提取客户端真实IP（支持X-Forwarded-For）
         */
        private String extractClientIp(Request request) {
            if (request.getContext() instanceof RequestDataContext) {
                RequestDataContext context = (RequestDataContext) request.getContext();
                HttpHeaders headers = context.getClientRequest().getHeaders();

                // 尝试从多种可能的请求头中获取客户端IP
                String[] headersArr = {"X-Forwarded-For", "X-Real-IP", "Proxy-Client-IP", "WL-Proxy-Client-IP"};
                for (String header : headersArr) {
                    String ip = headers.getFirst(header);
                    if (ip != null && !ip.isEmpty() && !"unknown".equalsIgnoreCase(ip)) {
                        // 如果是X-Forwarded-For，可能包含多个IP，取第一个
                        if ("X-Forwarded-For".equalsIgnoreCase(header)) {
                            int index = ip.indexOf(",");
                            if (index != -1) {
                                String firstIp = ip.substring(0, index);
                                log.debug("Extracted client IP from X-Forwarded-For header: {}", firstIp);
                                return firstIp;
                            }
                        }
                        log.debug("Extracted client IP from {} header: {}", header, ip);
                        return ip;
                    }
                }
            }
            return null;
        }

        /**
         * 记录选中的实例和请求信息
         */
        private void logSelectedInstance(Request request, ServiceInstance instance) {
            if (request.getContext() instanceof RequestDataContext) {
                RequestDataContext context = (RequestDataContext) request.getContext();
                URI url = context.getClientRequest().getUrl();
                String fullUrl = url.getScheme() + "://" + instance.getHost() + ":" + instance.getPort() + url.getPath();
                log.info("[LoadBalancer][{}] 选中实例: {} | 请求URL: {}, 请求方法: {}",
                        serviceId,
                        instance.getHost() + ":" + instance.getPort(),
                        fullUrl,
                        context.getClientRequest().getHttpMethod().name());
            }
        }

        // 其他原有方法保持不变...
        private Optional<ServiceInstance> findLocalInstance(List<ServiceInstance> instances) {
            return instances.stream()
                    .filter(instance -> localIps.contains(instance.getHost()))
                    .findAny();
        }

        private List<String> getAllLocalIps() {
            List<String> ipList = new ArrayList<>();
            try {
                Enumeration<NetworkInterface> interfaces = NetworkInterface.getNetworkInterfaces();
                while (interfaces.hasMoreElements()) {
                    NetworkInterface iface = interfaces.nextElement();
                    if (iface.isLoopback() || !iface.isUp()) {
                        continue;
                    }

                    Enumeration<InetAddress> addresses = iface.getInetAddresses();
                    while (addresses.hasMoreElements()) {
                        InetAddress addr = addresses.nextElement();
                        if (addr.getHostAddress().contains(":")) {
                            continue;
                        }
                        ipList.add(addr.getHostAddress());
                    }
                }
            } catch (SocketException e) {
                log.warn("Could not determine local IP addresses", e);
            }
            return Collections.unmodifiableList(ipList);
        }
    }
}