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.Stream;
026
027import static ca.uhn.fhir.model.api.ResourceMetadataKeyEnum.ENTRY_SEARCH_MODE;
028
029/**
030 * This class is used to collect the search results as they are fetched from
031 * individual targets from the gateway.
032 * An instance is created for each set of request sent from the gateway to target servers
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                        if (theSearchResult == null) {
068                                continue;
069                        }
070
071                        //                      BundleEntrySearchModeEnum searchMode = ENTRY_SEARCH_MODE.get(theSearchResult);
072                        // TODO GGG fix MetadataKeyEnum to support IBaseResource instead of IAnyResource
073                        BundleEntrySearchModeEnum searchMode =
074                                        (BundleEntrySearchModeEnum) theSearchResult.getUserData(ENTRY_SEARCH_MODE.name());
075                        if (searchMode == null) {
076                                uncategorizedList.add(theSearchResult);
077                        } else {
078                                switch (searchMode) {
079                                        case MATCH -> matchList.add(theSearchResult);
080                                        case INCLUDE -> includeList.add(theSearchResult);
081                                        case OUTCOME -> outcomeList.add(theSearchResult);
082                                        default -> uncategorizedList.add(theSearchResult);
083                                }
084                        }
085                }
086                if (!myTotalOmitted.get()) {
087                        myTotal.addAndGet(theTotal);
088                }
089        }
090
091        @Override
092        public List<IBaseResource> getResults(String theTargetId) {
093                return Stream.of(
094                                                getMatchedResults(theTargetId),
095                                                getIncludeResults(theTargetId),
096                                                getOutcomeResults(theTargetId),
097                                                getUncategorizedResults(theTargetId))
098                                .flatMap(Collection::stream)
099                                .toList();
100        }
101
102        @Override
103        public List<IBaseResource> getPageableResults(String theTargetId) {
104                return Collections.unmodifiableList(concatenateMatchedAndUncategorized(theTargetId));
105        }
106
107        private List<IBaseResource> concatenateMatchedAndUncategorized(String theTargetId) {
108                return Stream.concat(getMatchedResults(theTargetId).stream(), getUncategorizedResults(theTargetId).stream())
109                                .toList();
110        }
111
112        @Override
113        public List<IBaseResource> getMatchedResults(String theTargetId) {
114                List<IBaseResource> retVal = myTargetIdToMatchResults.getOrDefault(theTargetId, Collections.emptyList());
115                return Collections.unmodifiableList(retVal);
116        }
117
118        @Override
119        public List<IBaseResource> getIncludeResults(String theTargetId) {
120                List<IBaseResource> retVal = myTargetIdToIncludeResults.getOrDefault(theTargetId, Collections.emptyList());
121                return Collections.unmodifiableList(retVal);
122        }
123
124        @Override
125        public List<IBaseResource> getOutcomeResults(String theTargetId) {
126                List<IBaseResource> retVal = myTargetIdToOutcomeResults.getOrDefault(theTargetId, Collections.emptyList());
127                return Collections.unmodifiableList(retVal);
128        }
129
130        @Override
131        public List<IBaseResource> getUncategorizedResults(String theTargetId) {
132                List<IBaseResource> retVal =
133                                myTargetIdToUncategorizedResults.getOrDefault(theTargetId, Collections.emptyList());
134                return Collections.unmodifiableList(retVal);
135        }
136
137        @Override
138        public void removeResults(String theTargetId, List<IBaseResource> theResultsToBeRemoved) {
139                List.of(
140                                                myTargetIdToMatchResults,
141                                                myTargetIdToIncludeResults,
142                                                myTargetIdToUncategorizedResults,
143                                                myTargetIdToOutcomeResults)
144                                .forEach(map ->
145                                                map.getOrDefault(theTargetId, Collections.emptyList()).removeAll(theResultsToBeRemoved));
146        }
147
148        @Override
149        public int getTotal() {
150                return myTotal.get();
151        }
152
153        @Override
154        public boolean getTotalOmitted() {
155                return myTotalOmitted.get();
156        }
157
158        @Override
159        public void setTotalOmitted(boolean theTotalOmitted) {
160                myTotalOmitted.set(theTotalOmitted);
161        }
162
163        @Override
164        public void accumulate(List<IBaseResource> theResults, SearchSingleTargetResponse theResponse, String theId) {
165                Integer total = BundleUtil.getTotal(myFhirCtx, theResponse.getSearchResults());
166
167                if (total != null) {
168                        addResults(theId, theResults, total);
169                } else {
170                        setTotalOmitted(true);
171                        addResults(theId, theResults, 0);
172                }
173        }
174}