package cn.reghao.tnb.gateway.token;

import cn.reghao.jutil.jdk.web.result.WebResult;
import cn.reghao.jutil.jdk.security.RsaCryptor;
import cn.reghao.tnb.auth.JwtUtil;
import cn.reghao.tnb.auth.model.JwtPayload;
import cn.reghao.tnb.gateway.config.AppProperties;
import cn.reghao.tnb.gateway.token.model.UserData;
import cn.reghao.tnb.gateway.token.model.UserLogin;
import cn.reghao.tnb.gateway.config.RedisKeys;
import io.jsonwebtoken.ExpiredJwtException;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.http.HttpCookie;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.MultiValueMap;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.security.interfaces.RSAPublicKey;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 检查 JWT token(若存在) 是否过期
 *
 * @author reghao
 * @date 2021-07-28 16:50:50
 */
@Component
public class GlobalTokenFilter implements GlobalFilter, Ordered {
    private final String cookieName = "USERDATA";
    private final RedisTemplate<String, String> redisTemplate;
    private final Map<String, String> resourceMap = new HashMap<>();

    public GlobalTokenFilter(RedisTemplate<String, String> redisTemplate, AppProperties appProperties) {
        this.redisTemplate = redisTemplate;
        appProperties.getResources().forEach((key, value) -> {
            this.resourceMap.put(key.replace("_", "/"), value);
        });
    }

    @Override
    public int getOrder() {
        return -1;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();

        UserLogin userLogin = new UserLogin();
        UserData userData = getUserData(request);
        if (userData != null) {
            try {
                if (userData.getType().equals("token")) {
                    userLogin = getUserLoginByToken(userData.getData());
                } else {
                    userLogin = getUserLoginByCookie(userData.getData());
                }
            } catch (Exception exception) {
                // token 过期, 网关返回 401，由前端自行处理
                exchange.getResponse().setStatusCode(HttpStatus.UNAUTHORIZED);

                DataBuffer buffer;
                if (exception instanceof ExpiredJwtException) {
                    String msg = "token is expired";
                    byte[] bytes = msg.getBytes(StandardCharsets.UTF_8);
                    buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                } else {
                    String msg = exception.getMessage();
                    byte[] bytes = msg.getBytes(StandardCharsets.UTF_8);
                    buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                }
                return exchange.getResponse().writeWith(Flux.just(buffer));
            }
        }
        Set<String> roleSet = Arrays.stream(userLogin.getRoles().split(","))
                .map(role -> role.toLowerCase(Locale.ROOT))
                .collect(Collectors.toSet());
        String url = request.getURI().getPath();
        for (String urlPrefix : resourceMap.keySet()) {
            if (url.startsWith(urlPrefix)) {
                String role = resourceMap.get(urlPrefix);
                if (!roleSet.contains(role)) {
                    String errMsg = String.format("access %s require role %s", url, role);
                    return fastFinish(exchange, errMsg);
                }
            }
        }

        final UserLogin userLogin1 = userLogin;
        // 1. 构建包含新 Header 的请求
        ServerHttpRequest newRequest = exchange.getRequest().mutate()
                .header("x-user-id", ""+userLogin1.getUserId())
                .header("x-login-id", userLogin1.getLoginId())
                .header("x-user-roles", userLogin1.getRoles())
                .build();
        // 2. 先生成包含新请求的 Exchange
        ServerWebExchange newExchange = exchange.mutate()
                .request(newRequest)
                .build();
        return chain.filter(newExchange);
    }

    private UserData getUserData(ServerHttpRequest request) {
        List<String> headerValues = request.getHeaders().get(JwtUtil.AUTH_HEADER);
        if (headerValues != null && headerValues.size() == 1) {
            String headerValue = headerValues.get(0);
            if (headerValue != null && headerValue.startsWith(JwtUtil.JWT_PREFIX)) {
                String jwtToken = headerValue.replace(JwtUtil.JWT_PREFIX, "");
                return new UserData("token", jwtToken);
            }
        }

        MultiValueMap<String, HttpCookie> cookies = request.getCookies();
        Map<String, String> cookieMap = cookies.values().stream()
                .map(list -> list.get(0))
                .collect(Collectors.toMap(HttpCookie::getName, HttpCookie::getValue));
        String value = cookieMap.get(cookieName);
        if (value != null) {
            return new UserData("cookie", value);
        }

        return null;
    }

    private UserLogin getUserLoginByCookie(String cookie) {
        return new UserLogin();
    }

    private UserLogin getUserLoginByToken(String jwtToken) {
        RSAPublicKey signKeyPublic = getPublicKey();
        JwtPayload jwtPayload = JwtUtil.getJwtPayload(jwtToken, signKeyPublic);
        if (jwtPayload != null) {
            // TODO 帐号登出后, 已颁发 token 的 jti 会放到 blacklist 中, 需要判断当前 token 是否在 blacklist 中
            String jti = jwtPayload.getJti();
            long userId = jwtPayload.getUserId();
            String loginId = jwtPayload.getLoginId();
            String roles = jwtPayload.getAuthorities();
            String rolesStr = roles.replace("[", "").replace("]", "").replace(" ", "");
            Boolean exist = redisTemplate.hasKey(RedisKeys.getAuthTokenKey(userId, loginId));
            if (exist != null && exist) {
                return new UserLogin(userId, loginId, rolesStr);
            }
        }

        return new UserLogin();
    }

    private RSAPublicKey getPublicKey() {
        String key = RedisKeys.getJwtSignKey("pubkey");
        ValueOperations<String, String> operations = redisTemplate.opsForValue();
        String pubkey = operations.get(key);
        return pubkey != null ? RsaCryptor.getRSAPublicKey(pubkey) : null;
    }

    /**
     * 快速返回(不调用业务系统)
     *
     * @param
     * @return
     * @date 2025-10-17 21:10:660
     */
    private Mono<Void> fastFinish(ServerWebExchange exchange, String msg){
        ServerHttpResponse response = exchange.getResponse();
        response.setStatusCode(HttpStatus.OK);
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
        String fastResult = WebResult.failWithMsg(msg);
        DataBuffer dataBuffer = response.bufferFactory().allocateBuffer().write(fastResult.getBytes(StandardCharsets.UTF_8));
        return response.writeWith(Mono.just(dataBuffer));
    }
}
