annotation.api/src/main/java/ch/benlegiang/annotation/api/services/PredictionService.java
<filename>annotation.api/src/main/java/ch/benlegiang/annotation/api/services/PredictionService.javapackage ch.benlegiang.annotation.api.services;

import ch.benlegiang.annotation.api.entities.AnnotationEntity;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectReader;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;

import org.springframework.web.client.HttpClientErrorException;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class PredictionService {

    private final WebClient webClient;


    public PredictionService(WebClient.Builder webClientBuilder) {
        this.webClient = webClientBuilder.baseUrl("http://syntax-highlighting-service-load-balancer:7777").build();
    }

    public AnnotationEntity setPrediction(AnnotationEntity annotationEntity) throws IOException {
        try {
            Map<String, Object> requestBody = new HashMap();

            requestBody.put("codeLanguage", annotationEntity.getCodeLanguage());
            requestBody.put("tokenIds", annotationEntity.getLexed());

            String response = this.webClient.post()
                    .uri("/api/v1/predict")
                    .header("Content-Type", "application/json")
                    .contentType(MediaType.APPLICATION_JSON)
                    .accept(MediaType.APPLICATION_JSON)
                    .body(BodyInserters.fromValue(requestBody))
                    .retrieve()
                    .bodyToMono(String.class)
                    .block();

            ObjectMapper mapper = new ObjectMapper();
            JsonNode rootNode = mapper.readTree(response).get("prediction");

            ObjectReader reader = mapper.readerFor(new TypeReference<List<Integer>>() {
            });

            List<Integer> prediction = reader.readValue(rootNode);

            annotationEntity.setPrediction(prediction);

            return annotationEntity;
        } catch (HttpClientErrorException e) {
            e.printStackTrace();
            return null;
        }
    }
}
