import org.drools.core.ClassObjectFilter;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.kie.api.KieBaseConfiguration;
import org.kie.api.KieServices;
import org.kie.api.builder.KieBuilder;
import org.kie.api.builder.KieFileSystem;
import org.kie.api.builder.ReleaseId;
import org.kie.api.conf.EqualityBehaviorOption;
import org.kie.api.runtime.KieSession;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;

public class QueryingCorruptsClassStoreMapTest {

    KieSession dataLoadDestination;

    @Test
    public void worksIfWeInsertBeforeQuery() throws Exception {
        dataLoadDestination.insert(new BaseClass());
        final DerivedClass element = new DerivedClass();
        dataLoadDestination.insert(element);
        assertContains(queryFor(BaseClass.class), element);
    }

    @Test
    public void doesntWorkIfWeQueryBeforeInsert() throws Exception {
        dataLoadDestination.insert(new BaseClass());

        queryFor(DerivedClass.class); // Creates non-concrete class store

        final DerivedClass element = new DerivedClass();
        dataLoadDestination.insert(element); // Makes concrete but doesn't resolve dependent stores

        assertContains(queryFor(BaseClass.class), element);

    }

/////////////////// Everything below here is just setup

    @SuppressWarnings("unchecked")
    private <T> Collection<T> queryFor(Class<T> clazz) {
        return new ArrayList(dataLoadDestination.getObjects(classFilter(clazz)));
    }

    private <T> void assertContains(Iterable<? extends T> iterable, T item) {
        for (T t : iterable) {
            if (t == item) {
                return;
            }
        }
        Assert.fail("Item not found");
    }

    private ClassObjectFilter classFilter(Class<?> clazz) {
        return new ClassObjectFilter(clazz);
    }

    public  static class BaseClass {

    }

    public static class DerivedClass extends BaseClass {

    }

    @Before
    public void setUp() throws Exception {
        dataLoadDestination = knowledgeSession();
    }

    @After
    public void tearDown() throws Exception {
        dataLoadDestination.dispose();
    }

    /** This can't be the easiest way to set up a KieSession... */
    public KieSession knowledgeSession() throws IOException {
        KieServices ks = KieServices.Factory.get();
        KieFileSystem kieFileSystem = ks.newKieFileSystem();

        ReleaseId rid = ks.newReleaseId("com.acme.company", "rules", "1.0");
        kieFileSystem.generateAndWritePomXML(rid);

        KieBuilder kieBuilder = ks.newKieBuilder(kieFileSystem);
        kieBuilder.buildAll();

        KieBaseConfiguration kieBaseConfiguration = ks.newKieBaseConfiguration();
        kieBaseConfiguration.setOption(EqualityBehaviorOption.EQUALITY);
        return ks.newKieContainer(rid).newKieBase(kieBaseConfiguration).newKieSession();
    }
}

    @Test
    public void worksIfWeAvoidTheOptimization() throws Exception {
        dataLoadDestination.insert(new BaseClass());

        queryWithoutOptimizationFor(DerivedClass.class); // Creates non-concrete class store

        final DerivedClass element = new DerivedClass();
        dataLoadDestination.insert(element); // Makes concrete but doesn't resolve dependent stores

        assertContains(queryFor(BaseClass.class), element);

    }

    @SuppressWarnings("unchecked")
    private <T> Collection<T> queryWithoutOptimizationFor(Class<T> clazz) {
        return new ArrayList(dataLoadDestination.getObjects(clazz::isInstance));
    }