package org.opencds.cqf.providers;

import ca.uhn.fhir.jpa.provider.dstu3.JpaResourceProviderDstu3;
import ca.uhn.fhir.rest.server.exceptions.ResourceNotFoundException;
import org.hl7.fhir.dstu3.model.CodeSystem;
import org.hl7.fhir.dstu3.model.IdType;
import org.hl7.fhir.dstu3.model.ValueSet;
import org.opencds.cqf.cql.runtime.Code;
import org.opencds.cqf.cql.terminology.CodeSystemInfo;
import org.opencds.cqf.cql.terminology.TerminologyProvider;
import org.opencds.cqf.cql.terminology.ValueSetInfo;

import java.util.ArrayList;
import java.util.List;

/**
 * Created by Christopher Schuler on 7/17/2017.
 */
public class JpaTerminologyProvider implements TerminologyProvider {
    private JpaResourceProviderDstu3<ValueSet> valueSetProvider;
    private JpaResourceProviderDstu3<CodeSystem> codeSystemProvider;

    public JpaTerminologyProvider(JpaResourceProviderDstu3<ValueSet> valueSetProvider, JpaResourceProviderDstu3<CodeSystem> codeSystemProvider) {
        this.valueSetProvider = valueSetProvider;
        this.codeSystemProvider = codeSystemProvider;
    }

    @Override
    public boolean in(Code code, ValueSetInfo valueSet) throws ResourceNotFoundException {
        for (Code c : expand(valueSet)) {
            if (c == null) continue;
            if (c.getCode().equals(code.getCode()) && c.getSystem().equals(code.getSystem())) {
                return true;
            }
        }
        return false;
    }

    @Override
    public Iterable<Code> expand(ValueSetInfo valueSet) throws ResourceNotFoundException {
        ValueSet vs = valueSetProvider.getDao().read(new IdType(valueSet.getId()));
        List<Code> codes = new ArrayList<>();
        for (ValueSet.ValueSetExpansionContainsComponent expansion : vs.getExpansion().getContains()) {
            codes.add(new Code().withCode(expansion.getCode()).withSystem(expansion.getSystem()));
        }

        if (codes.isEmpty()) {
            for (ValueSet.ConceptReferenceComponent component : vs.getCompose().getInclude().get(0).getConcept()) {
                codes.add(new Code().withCode(component.getCode()).withSystem(vs.getCompose().getInclude().get(0).getSystem()));
            }
        }

        return codes;
    }

    @Override
    public Code lookup(Code code, CodeSystemInfo codeSystem) throws ResourceNotFoundException {
        CodeSystem cs = codeSystemProvider.getDao().read(new IdType(codeSystem.getId()));
        for (CodeSystem.ConceptDefinitionComponent concept : cs.getConcept()) {
            if (concept.getCode().equals(code.getCode()))
                return code.withSystem(codeSystem.getId()).withDisplay(concept.getDisplay());
        }
        return code;
    }
}

@@ -1946,17 +1946,11 @@
 cept
-Reference
+Set
 Comp
@@ -1955,25 +1955,23 @@
 mponent 
-component
+include
  : vs.ge
@@ -1989,31 +1989,156 @@
 getInclude()
-.get(0)
+) %7B%0A                String system = include.getSystem();%0A                for (ValueSet.ConceptReferenceComponent component : include
 .getConcept(
@@ -2134,32 +2134,36 @@
 getConcept()) %7B%0A
+    
                 
@@ -2228,58 +2228,35 @@
 tem(
-vs.getCompose().getInclude().get(0).getS
+s
 ystem
-(
 ))
-)
 ;
+%0A                %7D
 %0A   

Fix for issue https://github.com/DBCG/cqf-ruler/issues/31 ... incomplete expansion logic in JpaTerminologyProvider