package cn.reghao.tnb.gateway.log;

import cn.reghao.jutil.jdk.security.RsaCryptor;
import cn.reghao.jutil.jdk.string.IDObfuscation;
import cn.reghao.jutil.tool.jwt.Jwt;
import cn.reghao.tnb.gateway.log.model.UserData;
import cn.reghao.tnb.gateway.log.model.UserLogin;
import cn.reghao.tnb.gateway.config.RedisKeys;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.JwtException;
import io.jsonwebtoken.Jwts;
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.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
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.List;
import java.util.function.Consumer;

/**
 * 检查 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 IDObfuscation userIdObfuscation;

    public GlobalTokenFilter(RedisTemplate<String, String> redisTemplate, IDObfuscation userIdObfuscation) {
        this.redisTemplate = redisTemplate;
        this.userIdObfuscation = userIdObfuscation;
    }

    @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 (JwtException jwtException) {
                // token 过期, 网关返回 401，由前端自行处理
                exchange.getResponse().setStatusCode(HttpStatus.UNAUTHORIZED);

                DataBuffer buffer;
                if (jwtException instanceof ExpiredJwtException) {
                    String msg = "token is expired";
                    byte[] bytes = msg.getBytes(StandardCharsets.UTF_8);
                    buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                } else {
                    String msg = jwtException.getMessage();
                    byte[] bytes = msg.getBytes(StandardCharsets.UTF_8);
                    buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                }
                return exchange.getResponse().writeWith(Flux.just(buffer));
            }
        }

        final UserLogin userLogin1 = userLogin;
        Consumer<HttpHeaders> headers = header -> {
            header.set("x-user-id", userLogin1.getUserId());
            header.set("x-login-id", userLogin1.getLoginId());
        };
        exchange.getRequest().mutate().headers(headers).build();
        return chain.filter(exchange);
    }

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

        MultiValueMap<String, HttpCookie> cookies = request.getCookies();
        for (List<HttpCookie> list : cookies.values()) {
            for (HttpCookie cookie : list) {
                String name = cookie.getName();
                if (name.equalsIgnoreCase(cookieName)) {
                    String userCookie = cookie.getValue();
                    return new UserData("cookie", userCookie);
                }
            }
        }

        return null;
    }

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

    private UserLogin getUserLoginByToken(String jwtToken) {
        RSAPublicKey signKey = getPublicKey();
        Claims claims = Jwts.parser().setSigningKey(signKey).parseClaimsJws(jwtToken).getBody();
        // TODO 帐号登出后, 已颁发 token 的 jti 会放到 blacklist 中, 需要判断当前 token 是否在 blacklist 中
        String jti = (String) claims.get("jti");

        Integer plat = (Integer) claims.get("plat");
        String loginId = (String) claims.get("loginId");
        int loginType = (int) claims.get("loginType");
        String userIdStr = claims.getSubject();
        if (plat == null || loginId == null || userIdStr == null) {
            return new UserLogin();
        }

        long userId = userIdObfuscation.restore(userIdStr);
        return new UserLogin(""+userId, loginId);
    }

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

        return null;
    }
}
