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

import org.springframework.http.HttpRequest;
import org.springframework.http.client.BufferingClientHttpRequestFactory;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestFactory;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestClient;

import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * 带观测能力的 RestClient 工厂。
 */
public class ObservedRestClientFactory {

    public static final String HEADER_TRACE_ID = "X-Trace-Id";
    public static final String HEADER_REQUEST_ID = "X-Request-Id";
    public static final String HEADER_SESSION_ID = "X-Session-Id";

    private final AiObservationSupport support;

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

    public RestClient create(String baseUrl, ClientHttpRequestFactory requestFactory, String componentName) {
        RestClient.Builder builder = RestClient.builder()
                .requestFactory(requestFactory == null ? defaultRequestFactory() : requestFactory)
                .requestInterceptor(new ObservationClientHttpRequestInterceptor(componentName));
        if (StringUtils.hasText(baseUrl)) {
            builder.baseUrl(baseUrl);
        }
        return builder.build();
    }

    public ClientHttpRequestFactory defaultRequestFactory() {
        return new BufferingClientHttpRequestFactory(new SimpleClientHttpRequestFactory());
    }

    private final class ObservationClientHttpRequestInterceptor implements ClientHttpRequestInterceptor {

        private final String componentName;

        private ObservationClientHttpRequestInterceptor(String componentName) {
            this.componentName = componentName;
        }

        @Override
        public ClientHttpResponse intercept(HttpRequest request,
                                            byte[] body,
                                            ClientHttpRequestExecution execution) throws IOException {
            AiObservationContext currentContext = support.currentContext();
            if (support.isHttpHeaderPropagationEnabled() && StringUtils.hasText(currentContext.traceId())) {
                request.getHeaders().set(HEADER_TRACE_ID, currentContext.traceId());
            }
            if (support.isHttpHeaderPropagationEnabled() && StringUtils.hasText(currentContext.requestId())) {
                request.getHeaders().set(HEADER_REQUEST_ID, currentContext.requestId());
            }
            if (support.isHttpHeaderPropagationEnabled() && StringUtils.hasText(currentContext.sessionId())) {
                request.getHeaders().set(HEADER_SESSION_ID, currentContext.sessionId());
            }
            AiObservationContext outboundContext = currentContext.toBuilder()
                    .componentType(AiComponentType.HTTP_OUTBOUND)
                    .componentName(componentName)
                    .build();
            Map<String, Object> requestFields = new LinkedHashMap<>();
            requestFields.put("http.method", request.getMethod());
            requestFields.put("http.url", request.getURI().toString());
            requestFields.put("status", "started");
            AiObservationLogger logger = new AiObservationLogger(support);
            logger.info(ObservedRestClientFactory.class, "ai.http.client.request", outboundContext, requestFields);
            try {
                ClientHttpResponse response = support.withObservationContext(outboundContext, () -> {
                    try {
                        return execution.execute(request, body);
                    } catch (IOException exception) {
                        throw new RuntimeIoException(exception);
                    }
                });
                logger.info(
                        ObservedRestClientFactory.class,
                        "ai.http.client.response",
                        outboundContext,
                        Map.of(
                                "http.status_code", response.getStatusCode().value(),
                                "status", "finished"
                        )
                );
                return response;
            } catch (RuntimeIoException exception) {
                logger.error(
                        ObservedRestClientFactory.class,
                        "ai.request.failed",
                        outboundContext,
                        Map.of(
                                "status", "failed",
                                "error.type", IOException.class.getName(),
                                "error.message", exception.getCause().getMessage()
                        ),
                        exception.getCause()
                );
                throw (IOException) exception.getCause();
            }
        }
    }

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