package cn.myapps.gateway;

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.RedirectToGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.RewritePathGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.SetPathGatewayFilterFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriTemplate;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.HashMap;
import java.util.Map;

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

@Component
public class RewriteResourcePathGatewayFilterFactory extends SetPathGatewayFilterFactory {
    @Override
    public GatewayFilter apply(Config config) {
        UriTemplate uriTemplate = new UriTemplate(config.getTemplate());
        final String QUERY_PATH = "path";

        return new GatewayFilter() {
            @Override
            public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
                ServerHttpRequest req = exchange.getRequest();
                if(!CustomRedirectToGatewayFilterFactory.isAllowedAccess(exchange)) {
                    HttpStatus httpStatus = HttpStatus.UNAUTHORIZED;
                    setResponseStatus(exchange, httpStatus);
                    ServerHttpResponse httpResponse = exchange.getResponse();
                    return httpResponse.setComplete();
                }

                addOriginalRequestUrl(exchange, req.getURI());

                Map<String, String> uriVariables = new HashMap<>();
                Map<String,String> tmp = getUriTemplateVariables(exchange);
                if(CollectionUtils.isEmpty(tmp)) {
                    for (String key : tmp.keySet()) {
                        uriVariables.put(key, tmp.get(key));
                    }
                }

                if (exchange.getRequest().getQueryParams().containsKey(QUERY_PATH)) {
                    String path = (String) exchange.getRequest().getQueryParams().get(QUERY_PATH).get(0);
                    uriVariables.put("query_path", path);
                }

                URI uri = uriTemplate.expand(uriVariables);
                String newPath = uri.getRawPath();
                exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, uri);


                ServerHttpRequest request = req.mutate().path(newPath).build();

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