package cc.mrbird.febs.auth.configure;

import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.ArrayUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.*;

@Slf4j
public class HeaderMapRequestWrapper extends HttpServletRequestWrapper {

    private byte[] bodyContent;
    private String requestBody;
    private Map<String, String[]> parameterMap = new HashMap<>();
    /**
     * construct a wrapper for this request
     *
     * @param request
     */
    public HeaderMapRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        // 读取并缓存请求体
        this.bodyContent = IOUtils.toByteArray(request.getInputStream());
        this.requestBody = new String(bodyContent, StandardCharsets.UTF_8);
    }

    private Map<String, String> headerMap = new HashMap<>();

    /**
     * add a header with given name and value
     *
     * @param name
     * @param value
     */
    public void addHeader(String name, String value) {
        headerMap.put(name, value);
    }

    @Override
    public String getHeader(String name) {
        String headerValue = super.getHeader(name);
        if (headerMap.containsKey(name)) {
            headerValue = headerMap.get(name);
        }
        return headerValue;
    }

    /**
     * get the Header names
     */
    @Override
    public Enumeration<String> getHeaderNames() {
        List<String> names = Collections.list(super.getHeaderNames());
        for (String name : headerMap.keySet()) {
            names.add(name);
        }
        return Collections.enumeration(names);
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        List<String> values = Collections.list(super.getHeaders(name));
        if (headerMap.containsKey(name)) {
            values = Arrays.asList(headerMap.get(name));
        }
        return Collections.enumeration(values);
    }


    // 新增setBody方法
    public void setBody(String content) {
        this.bodyContent = content.getBytes(StandardCharsets.UTF_8);
        this.requestBody = content; // 更新requestBody字段
        this.addHeader("Content-Length", String.valueOf(this.bodyContent.length));
    }


    // 添加参数
    public void setParameter(String name, String value) {
        parameterMap.put(name, new String[]{value});
    }

    @Override
    public String getParameter(String name) {
        String[] values = getParameterMap().get(name);
        return values != null && values.length > 0 ? values[0] : null;
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        // 合并原始参数、手动设置参数和解析的body参数
        Map<String, String[]> combined = new HashMap<>(super.getParameterMap());
        combined.putAll(parseParameters()); // 解析请求体中的参数
        combined.putAll(parameterMap); // 添加手动设置的参数
        return Collections.unmodifiableMap(combined);
    }

    @Override
    public Enumeration<String> getParameterNames() {
        return Collections.enumeration(getParameterMap().keySet());
    }

    @Override
    public String[] getParameterValues(String name) {
        return getParameterMap().get(name);
    }

    private Map<String, String[]> parseParameters() {
        Map<String, String[]> params = new HashMap<>();
        String contentType = getContentType();
        String body = getRequestBody();

        if (contentType == null || body == null) {
            return params;
        }
        if (contentType != null && contentType.contains("application/json")) {
            // 解析 JSON 格式
            try {
                JSONObject jsonBody = JSONObject.parseObject(body);
                for (String key : jsonBody.keySet()) {
                    String value = jsonBody.getString(key);
                    params.put(key, new String[]{value});
                }
            } catch (Exception e) {
                log.info("Failed to parse JSON body", e);
            }
        }
        // 仅处理 application/x-www-form-urlencoded
        else if (contentType.contains("application/x-www-form-urlencoded")) {
            String[] pairs = body.split("&");
            for (String pair : pairs) {
                int idx = pair.indexOf('=');
                if (idx > 0) {
                    String key = pair.substring(0, idx);
                    String value = pair.substring(idx + 1);
                    params.put(key, new String[]{value});
                }
            }
        }
        // 不再处理其他格式（如JSON）
        return params;
    }

    @Override
    public ServletInputStream getInputStream() {
        return new ServletInputStreamWrapper(this.bodyContent);
    }

    // 新增获取缓存请求体的方法
    public String getRequestBody() {
        return this.requestBody;
    }

    // 确保返回正确的Content-Length
    @Override
    public int getContentLength() {
        return this.bodyContent.length;
    }


    @Override
    public BufferedReader getReader() {
        return new BufferedReader(new InputStreamReader(this.getInputStream()));
    }

    // 内部输入流包装类
    private static class ServletInputStreamWrapper extends ServletInputStream {
        private final ByteArrayInputStream inputStream;

        public ServletInputStreamWrapper(byte[] content) {
            this.inputStream = new ByteArrayInputStream(content);
        }

        @Override
        public int read() {
            return inputStream.read();
        }

        @Override
        public boolean isFinished() {
            return inputStream.available() == 0;
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setReadListener(ReadListener listener) {
            throw new UnsupportedOperationException();
        }

    }
}
