001/*-
002 * #%L
003 * Smile CDR - CDR
004 * %%
005 * Copyright (C) 2016 - 2024 Smile CDR, Inc.
006 * %%
007 * All rights reserved.
008 * #L%
009 */
010package ca.cdr.api.fhirgw.model;
011
012import ca.uhn.fhir.context.FhirContext;
013import ca.uhn.fhir.model.valueset.BundleEntrySearchModeEnum;
014import ca.uhn.fhir.util.BundleUtil;
015import org.hl7.fhir.instance.model.api.IBaseResource;
016
017import java.util.ArrayList;
018import java.util.Collection;
019import java.util.Collections;
020import java.util.HashMap;
021import java.util.List;
022import java.util.Map;
023import java.util.concurrent.atomic.AtomicBoolean;
024import java.util.concurrent.atomic.AtomicInteger;
025import java.util.stream.Collectors;
026import java.util.stream.Stream;
027
028import static ca.uhn.fhir.model.api.ResourceMetadataKeyEnum.ENTRY_SEARCH_MODE;
029
030/**
031 * This class is used to collect the search results as they are fetched from
032 * individual targets from the gateway.
033 */
034public class BundleResultsAccumulator implements IResultsAccumulator<SearchSingleTargetResponse> {
035
036        private final FhirContext myFhirCtx;
037
038        public BundleResultsAccumulator(FhirContext theFhirCtx) {
039                myFhirCtx = theFhirCtx;
040        }
041
042        private final Map<String, List<IBaseResource>> myTargetIdToMatchResults =
043                        Collections.synchronizedMap(new HashMap<>());
044        private final Map<String, List<IBaseResource>> myTargetIdToIncludeResults =
045                        Collections.synchronizedMap(new HashMap<>());
046
047        private final Map<String, List<IBaseResource>> myTargetIdToOutcomeResults =
048                        Collections.synchronizedMap(new HashMap<>());
049        private final Map<String, List<IBaseResource>> myTargetIdToUncategorizedResults =
050                        Collections.synchronizedMap(new HashMap<>());
051        private final AtomicInteger myTotal = new AtomicInteger();
052        private final AtomicBoolean myTotalOmitted = new AtomicBoolean();
053
054        @Override
055        public void addResults(String theTargetId, List<IBaseResource> theSearchResults, int theTotal) {
056                List<IBaseResource> matchList = myTargetIdToMatchResults.computeIfAbsent(
057                                theTargetId, t -> Collections.synchronizedList(new ArrayList<>()));
058                List<IBaseResource> includeList = myTargetIdToIncludeResults.computeIfAbsent(
059                                theTargetId, t -> Collections.synchronizedList(new ArrayList<>()));
060                List<IBaseResource> outcomeList = myTargetIdToOutcomeResults.computeIfAbsent(
061                                theTargetId, t -> Collections.synchronizedList(new ArrayList<>()));
062                List<IBaseResource> uncategorizedList = myTargetIdToUncategorizedResults.computeIfAbsent(
063                                theTargetId, t -> Collections.synchronizedList(new ArrayList<>()));
064
065                // Split into actual matched results, includes, outcomes, and unknown(missing)
066                for (IBaseResource theSearchResult : theSearchResults) {
067
068                        //                      BundleEntrySearchModeEnum searchMode = ENTRY_SEARCH_MODE.get(theSearchResult);
069                        // TODO GGG fix MetadataKeyEnum to support IBaseResource instead of IAnyResource
070                        Object userData = theSearchResult == null ? null : theSearchResult.getUserData(ENTRY_SEARCH_MODE.name());
071                        if (userData == null) {
072                                uncategorizedList.add(theSearchResult);
073                        } else {
074                                BundleEntrySearchModeEnum searchMode = (BundleEntrySearchModeEnum) userData;
075                                switch (searchMode) {
076                                        case MATCH -> matchList.add(theSearchResult);
077                                        case INCLUDE -> includeList.add(theSearchResult);
078                                        case OUTCOME -> outcomeList.add(theSearchResult);
079                                        default -> uncategorizedList.add(theSearchResult);
080                                }
081                        }
082                }
083                if (!myTotalOmitted.get()) {
084                        myTotal.addAndGet(theTotal);
085                }
086        }
087
088        @Override
089        public List<IBaseResource> getResults(String theTargetId) {
090                return Stream.of(
091                                                getMatchedResults(theTargetId),
092                                                getIncludeResults(theTargetId),
093                                                getOutcomeResults(theTargetId),
094                                                getUncategorizedResults(theTargetId))
095                                .flatMap(Collection::stream)
096                                .collect(Collectors.toList());
097        }
098
099        @Override
100        public List<IBaseResource> getPageableResults(String theTargetId) {
101                return Collections.unmodifiableList(concatenateMatchedAndUncategorized(theTargetId));
102        }
103
104        private List<IBaseResource> concatenateMatchedAndUncategorized(String theTargetId) {
105                return Stream.concat(getMatchedResults(theTargetId).stream(), getUncategorizedResults(theTargetId).stream())
106                                .toList();
107        }
108
109        @Override
110        public List<IBaseResource> getMatchedResults(String theTargetId) {
111                List<IBaseResource> retVal = myTargetIdToMatchResults.getOrDefault(theTargetId, Collections.emptyList());
112                return Collections.unmodifiableList(retVal);
113        }
114
115        @Override
116        public List<IBaseResource> getIncludeResults(String theTargetId) {
117                List<IBaseResource> retVal = myTargetIdToIncludeResults.getOrDefault(theTargetId, Collections.emptyList());
118                return Collections.unmodifiableList(retVal);
119        }
120
121        @Override
122        public List<IBaseResource> getOutcomeResults(String theTargetId) {
123                List<IBaseResource> retVal = myTargetIdToOutcomeResults.getOrDefault(theTargetId, Collections.emptyList());
124                return Collections.unmodifiableList(retVal);
125        }
126
127        @Override
128        public List<IBaseResource> getUncategorizedResults(String theTargetId) {
129                List<IBaseResource> retVal =
130                                myTargetIdToUncategorizedResults.getOrDefault(theTargetId, Collections.emptyList());
131                return Collections.unmodifiableList(retVal);
132        }
133
134        @Override
135        public void removeResults(String theTargetId, List<IBaseResource> theResultsToBeRemoved) {
136                List.of(
137                                                myTargetIdToMatchResults,
138                                                myTargetIdToIncludeResults,
139                                                myTargetIdToUncategorizedResults,
140                                                myTargetIdToOutcomeResults)
141                                .forEach(map -> {
142                                        map.getOrDefault(theTargetId, Collections.emptyList()).removeAll(theResultsToBeRemoved);
143                                });
144        }
145
146        @Override
147        public int getTotal() {
148                return myTotal.get();
149        }
150
151        @Override
152        public boolean getTotalOmitted() {
153                return myTotalOmitted.get();
154        }
155
156        @Override
157        public void setTotalOmitted(boolean theTotalOmitted) {
158                myTotalOmitted.set(theTotalOmitted);
159        }
160
161        @Override
162        public void accumulate(List<IBaseResource> theResults, SearchSingleTargetResponse theResponse, String theId) {
163                Integer total = BundleUtil.getTotal(myFhirCtx, theResponse.getSearchResults());
164
165                if (total != null) {
166                        addResults(theId, theResults, total);
167                } else {
168                        setTotalOmitted(true);
169                        addResults(theId, theResults, 0);
170                }
171        }
172}