Untitled

package com.cais.ufg.aspect;

import com.cais.ufg.tracing.UfgRequestContext;
import com.cais.ufg.tracing.UfgRequestTraceService;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.core.Authentication;
import org.springframework.stereotype.Component;
import org.springframework.ws.transport.context.TransportContextHolder;
import org.springframework.ws.transport.http.HttpUrlConnection;

import java.lang.reflect.Field;
import java.util.Map;
import java.util.stream.Collectors;

@Aspect
@Component
@Slf4j
public class UfgRequestLoggingAspect {

    private final UfgRequestTraceService traceService;

    public UfgRequestLoggingAspect(UfgRequestTraceService traceService) {
        this.traceService = traceService;
    }

    @Around("@annotation(com.cais.ufg.aspect.LogUfgRequest)")
    public Object logUfgRequest(ProceedingJoinPoint joinPoint) throws Throwable {
        Object response = null;

        try {
            // Retrieve method arguments (assuming request is the first argument)
            Object[] args = joinPoint.getArgs();
            Object request = args.length > 0 ? args[0] : null;

            // Proceed with the method execution and capture the response
            response = joinPoint.proceed();

            // Retrieve the token and related details from the security context
            String tokenIssuer = getTokenIssuer();
            String tokenUsername = getTokenUsername();

            // Extract headers and other request/response details from the client
            String requestBody = extractRequestBody(request);
            String responseBody = extractResponseBody(response);
            String[] responseHeaders = extractResponseHeaders();

            // Retrieve traceId from Spring Actuator or logging context (e.g., MDC)
            String traceId = getTraceId();

            // Placeholder for callerAddress
            String callerAddress = "someCallerAddress"; // Implement a strategy to extract this if needed

            // Build the UfgRequestContext
            UfgRequestContext context = buildUfgRequestContext(
                joinPoint.getTarget(), 
                requestBody, 
                responseBody, 
                responseHeaders, 
                tokenIssuer, 
                tokenUsername, 
                traceId, 
                callerAddress
            );

            // Log the context using UfgRequestTraceService
            traceService.log(context);

        } catch (Exception e) {
            log.error("Failed to log UfgRequestContext", e);
            throw e; // Ensure the original method behavior is not altered
        }

        return response;
    }

    private String getTokenIssuer() {
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        if (authentication != null && authentication.getCredentials() instanceof SomeTokenType) {
            return ((SomeTokenType) authentication.getCredentials()).getIssuer();
        }
        throw new IllegalStateException("Token issuer is missing.");
    }

    private String getTokenUsername() {
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        return authentication != null ? authentication.getName() : "UnknownUser";
    }

    private String extractRequestBody(Object request) {
        return request != null ? request.toString() : "No Request";
    }

    private String extractResponseBody(Object response) {
        return response != null ? response.toString() : "No Response";
    }

    private String[] extractResponseHeaders() {
        try {
            HttpUrlConnection connection = (HttpUrlConnection) TransportContextHolder.getTransportContext().getConnection();
            return connection.getConnection().getHeaderFields().entrySet().stream()
                    .flatMap(entry -> entry.getValue().stream())
                    .toArray(String[]::new);
        } catch (Exception e) {
            log.error("Unable to extract response headers", e);
            return new String[]{"No Headers Available"};
        }
    }

    private String getTraceId() {
        // Assuming you have an MDC or Spring Actuator trace integration
        return MDC.get("traceId"); // Replace with proper Actuator integration if required
    }

    private UfgRequestContext buildUfgRequestContext(Object clientInstance,
                                                     String requestBody,
                                                     String responseBody,
                                                     String[] responseHeaders,
                                                     String tokenIssuer,
                                                     String tokenUsername,
                                                     String traceId,
                                                     String callerAddress) {
        String endpointUrl = getEndpointUrlFromClient(clientInstance);

        return new UfgRequestContext(
            endpointUrl, 
            "objectId", 
            "organizationCauseSign", 
            "caseSignType", 
            "purpose", 
            "externalLogsCorrelationId", 
            "businessUserId", 
            tokenIssuer, 
            tokenUsername, 
            requestBody, 
            responseBody, 
            responseHeaders, 
            traceId, 
            callerAddress
        );
    }

    private String getEndpointUrlFromClient(Object clientInstance) {
        try {
            Field endpointField = clientInstance.getClass().getDeclaredField("ufgPoliceNotesUrl");
            endpointField.setAccessible(true);
            return (String) endpointField.get(clientInstance);
        } catch (Exception e) {
            log.error("Failed to retrieve endpoint URL from client class", e);
            throw new IllegalStateException("Could not extract endpoint URL", e);
        }
    }
}