package com.tofly.getway.balance;

import com.tofly.common.core.util.StringUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancerClient;
import org.springframework.web.server.ServerWebExchange;

import java.net.URI;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;

/**
 * @author tuonbed
 * @date 2020/7/21  16:19
 */
@Slf4j
public class ToflyLoadBalancerClientFilter extends LoadBalancerClientFilter {

    public ToflyLoadBalancerClientFilter(LoadBalancerClient loadBalancer, LoadBalancerProperties properties) {
        super(loadBalancer, properties);
    }

    @Override
    protected ServiceInstance choose(ServerWebExchange exchange) {
//        return loadBalancer.choose(
//                ((URI) exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR)).getHost());

        if (this.loadBalancer instanceof RibbonLoadBalancerClient) {
            RibbonLoadBalancerClient client = (RibbonLoadBalancerClient) this.loadBalancer;
            String serviceId = ((URI) exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR)).getHost();
            //这里使用userId做为选择服务实例的key
            if(serviceId.equalsIgnoreCase("tofly-base")){
                //这里可以拿到web请求的上下文，可以从header中取出来自己定义的数据。
                String token = exchange.getRequest().getHeaders().getFirst("Authorization");
                log.info("从Authorization获取到的token为:{}",token);
                String access_token = exchange.getRequest().getQueryParams().getFirst("access_token");
                token=StringUtil.isEmpty(token)?access_token:token;
                if (!StringUtil.isEmpty(token)) {
                    token=token.replace("bearer ","");
                    return client.choose(serviceId, token);
                }
            }
        }
        return super.choose(exchange);

    }
}
