package com.ruoyi.framework.web.filter;

import java.io.IOException;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import com.ruoyi.framework.shiro.web.filter.embed.EmbedRequestUtils;

/**
 * 仅改写会话类 Cookie（JSESSIONID / rememberMe），不影响业务 Cookie。
 * <p>
 * 目标：在无法改 Nginx 时，尽量让浏览器收下会话 Cookie，供 HIS 同主机:8899 iframe 使用。
 * 非嵌入主流访问：只清理「Domain=IP」这类明显错误，不强制改其它合法 Domain。
 */
public class SessionCookieRewriteFilter implements Filter
{
    private static final Pattern COOKIE_NAME = Pattern.compile("^\\s*([^=\\s;]+)\\s*=");
    private static final Pattern DOMAIN_ATTR = Pattern.compile("(?i);\\s*Domain=([^;]*)");
    private static final Pattern IP_DOMAIN = Pattern.compile(
            "^\\[?[0-9a-fA-F:.]+\\]?$|^\\d{1,3}(?:\\.\\d{1,3}){3}$");

    @Override
    public void init(FilterConfig filterConfig)
    {
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException
    {
        HttpServletRequest req = (HttpServletRequest) request;
        HttpServletResponse res = (HttpServletResponse) response;
        boolean embed = EmbedRequestUtils.isEmbed(req);
        boolean secure = req.isSecure()
                || "https".equalsIgnoreCase(req.getHeader("X-Forwarded-Proto"));
        chain.doFilter(request, new CookieRewriteResponse(res, embed, secure));
    }

    @Override
    public void destroy()
    {
    }

    static boolean isSessionCookieName(String name)
    {
        if (name == null)
        {
            return false;
        }
        String n = name.trim();
        return "JSESSIONID".equalsIgnoreCase(n)
                || "rememberMe".equalsIgnoreCase(n)
                || "SHIROJSESSIONID".equalsIgnoreCase(n);
    }

    static boolean isIpLiteral(String domain)
    {
        if (domain == null)
        {
            return false;
        }
        String d = domain.trim();
        if (d.startsWith("."))
        {
            d = d.substring(1);
        }
        return IP_DOMAIN.matcher(d).matches();
    }

    /**
     * @param embedMode HIS iframe 嵌入：会话 Cookie 强制 host-only + SameSite=Lax
     * @param secureRequest 当前是否 HTTPS（含 X-Forwarded-Proto）
     */
    static String rewriteSetCookie(String raw, boolean embedMode, boolean secureRequest)
    {
        if (raw == null || raw.isEmpty())
        {
            return raw;
        }
        Matcher nameMatcher = COOKIE_NAME.matcher(raw);
        if (!nameMatcher.find())
        {
            return raw;
        }
        String cookieName = nameMatcher.group(1);
        if (!isSessionCookieName(cookieName))
        {
            // 业务 Cookie 原样返回，不影响主流任务
            return raw;
        }

        String value = raw;
        Matcher domainMatcher = DOMAIN_ATTR.matcher(value);
        if (domainMatcher.find())
        {
            String domainVal = domainMatcher.group(1) == null ? "" : domainMatcher.group(1).trim();
            // embed：一律 host-only；主流：仅去掉 Domain=IP（反代常见误配）
            if (embedMode || isIpLiteral(domainVal))
            {
                value = domainMatcher.replaceFirst("");
            }
        }

        if (!secureRequest)
        {
            value = value.replaceAll("(?i);\\s*Secure", "");
            value = value.replaceAll("(?i)SameSite=None", "SameSite=Lax");
        }

        // 会话 Cookie 补 SameSite=Lax（同站不同端口 iframe 可带上；不使用 None 以免无 Secure 被丢）
        if (!value.matches("(?i).*\\bSameSite\\s*=.*"))
        {
            value = value + "; SameSite=Lax";
        }
        return value;
    }

    private static final class CookieRewriteResponse extends HttpServletResponseWrapper
    {
        private final boolean embedMode;
        private final boolean secureRequest;

        CookieRewriteResponse(HttpServletResponse response, boolean embedMode, boolean secureRequest)
        {
            super(response);
            this.embedMode = embedMode;
            this.secureRequest = secureRequest;
        }

        @Override
        public void addCookie(Cookie cookie)
        {
            if (cookie != null && isSessionCookieName(cookie.getName()))
            {
                String domain = cookie.getDomain();
                if (embedMode || isIpLiteral(domain))
                {
                    cookie.setDomain(null);
                }
                if (cookie.getPath() == null || cookie.getPath().isEmpty())
                {
                    cookie.setPath("/");
                }
                if (!secureRequest)
                {
                    cookie.setSecure(false);
                }
            }
            super.addCookie(cookie);
        }

        @Override
        public void setHeader(String name, String value)
        {
            if ("Set-Cookie".equalsIgnoreCase(name))
            {
                super.setHeader(name, rewriteSetCookie(value, embedMode, secureRequest));
                return;
            }
            super.setHeader(name, value);
        }

        @Override
        public void addHeader(String name, String value)
        {
            if ("Set-Cookie".equalsIgnoreCase(name))
            {
                super.addHeader(name, rewriteSetCookie(value, embedMode, secureRequest));
                return;
            }
            super.addHeader(name, value);
        }
    }
}
