package cn.reghao.tnb.gateway.lb;

import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
import org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.ReactiveLoadBalancerClientFilter;
import org.springframework.cloud.gateway.support.DelegatingServiceInstance;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.nio.charset.StandardCharsets;

/**
 * @author reghao
 * @date 2023-02-03 16:23:06
 */
@Slf4j
public class UserReactiveLoadBalancerClientFilter extends ReactiveLoadBalancerClientFilter {
    private static final int LOAD_BALANCER_CLIENT_FILTER_ORDER = 10150;
    private final LoadBalancerProperties properties;
    private final UserLoadBalancer userLoadBalancer;

    public UserReactiveLoadBalancerClientFilter(LoadBalancerProperties properties, UserLoadBalancer userLoadBalancer) {
        super(null,null);
        this.properties = properties;
        this.userLoadBalancer = userLoadBalancer;
    }

    @Override
    public int getOrder() {
        return LOAD_BALANCER_CLIENT_FILTER_ORDER;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI url = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
        String schemePrefix = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR);
        if (url == null || (!"lb".equals(url.getScheme()) && !"lb".equals(schemePrefix))) {
            return chain.filter(exchange);
        }
        // preserve the original url
        ServerWebExchangeUtils.addOriginalRequestUrl(exchange, url);

        if (log.isTraceEnabled()) {
            log.trace(ReactiveLoadBalancerClientFilter.class.getSimpleName() + " url before: " + url);
        }

        try {
            return choose(exchange).doOnNext(response -> {
                if (!response.hasServer()) {
                    //boolean notFound1 = this.properties.isUse404();
                    boolean notFound = true;
                    throw NotFoundException.create(notFound, "Unable to find instance for " + url.getHost());
                }

                URI uri = exchange.getRequest().getURI();

                // if the `lb:<scheme>` mechanism was used, use `<scheme>` as the default,
                // if the loadbalancer doesn't provide one.
                ServiceInstance retrievedInstance = (ServiceInstance)response.getServer();
                String overrideScheme = retrievedInstance.isSecure() ? "https" : "http";
                if (schemePrefix != null) {
                    overrideScheme = url.getScheme();
                }

                DelegatingServiceInstance serviceInstance = new DelegatingServiceInstance(response.getServer(), overrideScheme);
                String host = serviceInstance.getHost();
                int port = serviceInstance.getPort();
                String targetService = String.format("%s:%s", host, port);
                exchange.getAttributes().put("targetService", targetService);
                URI requestUrl = LoadBalancerUriTools.reconstructURI(serviceInstance, uri);
                if (log.isTraceEnabled()) {
                    log.trace("LoadBalancerClientFilter url chosen: " + requestUrl);
                }
                exchange.getAttributes().put(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR, requestUrl);
            }).then(chain.filter(exchange));
        } catch (NotFoundException e) {
            exchange.getAttributes().put("targetService", "SERVICE_UNAVAILABLE");
            exchange.getResponse().setStatusCode(e.getStatusCode());
            String msg = e.getMessage();
            byte[] bytes = msg.getBytes(StandardCharsets.UTF_8);

            DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
            return exchange.getResponse().writeWith(Flux.just(buffer));
        }
    }

    private Mono<Response<ServiceInstance>> choose(ServerWebExchange exchange) {
        URI uri = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
        ServiceInstance serviceInstance = userLoadBalancer.choose(uri.getHost(), exchange.getRequest());
        String path = uri.getPath();
        String host = serviceInstance.getHost();
        int port = serviceInstance.getPort();
        //log.info("{} -> {}:{}", path, host, port);
        return Mono.just(new DefaultResponse(serviceInstance));
    }
}
