package com.infoepoch.pms.agent.observability.ai;

import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.web.filter.OncePerRequestFilter;

import java.io.IOException;

/**
 * 统一注入 AI 观测链路上下文。
 */
public class AiObservationFilter extends OncePerRequestFilter {

    public static final String HEADER_TRACE_ID = "X-Trace-Id";
    public static final String HEADER_REQUEST_ID = "X-Request-Id";
    public static final String REQUEST_ATTRIBUTE_TRACE_ID = "aiObservationTraceId";
    public static final String REQUEST_ATTRIBUTE_REQUEST_ID = "aiObservationRequestId";

    private final AiObservationSupport support;

    public AiObservationFilter(AiObservationSupport support) {
        this.support = support;
    }

    @Override
    protected boolean shouldNotFilter(HttpServletRequest request) {
        String uri = request.getRequestURI();
        return uri == null || !uri.contains("/api/agent/");
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request,
                                    HttpServletResponse response,
                                    FilterChain filterChain) throws ServletException, IOException {
        if (!support.isEnabled()) {
            filterChain.doFilter(request, response);
            return;
        }
        String traceId = headerOrGenerate(request, HEADER_TRACE_ID);
        String requestId = headerOrGenerate(request, HEADER_REQUEST_ID);
        request.setAttribute(REQUEST_ATTRIBUTE_TRACE_ID, traceId);
        request.setAttribute(REQUEST_ATTRIBUTE_REQUEST_ID, requestId);
        AiObservationContext context = AiObservationContext.builder()
                .traceId(traceId)
                .requestId(requestId)
                .flowId(requestId)
                .componentType(AiComponentType.HTTP_INBOUND)
                .componentName(request.getRequestURI())
                .build();
        try {
            support.withObservationContext(context, () -> {
                try {
                    filterChain.doFilter(request, response);
                    return null;
                } catch (IOException exception) {
                    throw new RuntimeIoException(exception);
                } catch (ServletException exception) {
                    throw new RuntimeServletException(exception);
                }
            });
        } catch (RuntimeIoException exception) {
            throw (IOException) exception.getCause();
        } catch (RuntimeServletException exception) {
            throw (ServletException) exception.getCause();
        }
    }

    private String headerOrGenerate(HttpServletRequest request, String headerName) {
        String headerValue = request.getHeader(headerName);
        return headerValue != null && !headerValue.isBlank() ? headerValue.trim() : AiObservationSupport.newObservationId();
    }

    private static final class RuntimeIoException extends RuntimeException {
        private RuntimeIoException(IOException cause) {
            super(cause);
        }
    }

    private static final class RuntimeServletException extends RuntimeException {
        private RuntimeServletException(ServletException cause) {
            super(cause);
        }
    }
}
