package org.hzhub.gateway.filter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.regex.Pattern;

/**
 * XSS 防护全局过滤器
 * 对 POST/PUT 请求体进行 XSS 清洗，防止脚本注入
 */
@Slf4j
@Component
public class XssGlobalFilter implements WebFilter, Ordered {

    /**
     * 排除路径（如文件上传等不需要 XSS 清洗的接口）
     */
    private static final List<String> EXCLUDE_PATHS = List.of(
        "/ai/upload"
    );

    /**
     * XSS 清洗规则
     */
    private static final List<Pattern> XSS_PATTERNS = List.of(
        Pattern.compile("<script[^>]*>.*?</script>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL),
        Pattern.compile("javascript\\s*:", Pattern.CASE_INSENSITIVE),
        Pattern.compile("on\\w+\\s*=\\s*['\"]", Pattern.CASE_INSENSITIVE),
        Pattern.compile("<iframe[^>]*>.*?</iframe>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL),
        Pattern.compile("<object[^>]*>.*?</object>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL)
    );

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        HttpMethod method = request.getMethod();
        String path = request.getURI().getPath();

        // GET/DELETE 请求不处理
        if (method == HttpMethod.GET || method == HttpMethod.DELETE || method == null) {
            return chain.filter(exchange);
        }

        // 排除路径放行
        if (EXCLUDE_PATHS.stream().anyMatch(path::startsWith)) {
            return chain.filter(exchange);
        }

        // 装饰请求，对 Body 进行 XSS 清洗
        ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(request) {
            @Override
            public Flux<DataBuffer> getBody() {
                return super.getBody().map(buffer -> {
                    byte[] bytes = new byte[buffer.readableByteCount()];
                    buffer.read(bytes);
                    DataBufferUtils.release(buffer);

                    String body = new String(bytes, StandardCharsets.UTF_8);
                    String cleaned = cleanXss(body);

                    if (!body.equals(cleaned)) {
                        log.debug("XSS 清洗: path={}", path);
                    }

                    byte[] cleanedBytes = cleaned.getBytes(StandardCharsets.UTF_8);
                    return exchange.getResponse().bufferFactory().wrap(cleanedBytes);
                });
            }
        };

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

    /**
     * 执行 XSS 清洗
     */
    private String cleanXss(String value) {
        if (value == null) return null;
        String result = value;
        for (Pattern pattern : XSS_PATTERNS) {
            result = pattern.matcher(result).replaceAll("");
        }
        return result;
    }

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