/*
 * Copyright 2003-2016 JetBrains s.r.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package jetbrains.mps.smodel;

import jetbrains.mps.components.CoreComponent;
import jetbrains.mps.smodel.adapter.ids.SConceptId;
import jetbrains.mps.smodel.adapter.structure.MetaAdapterFactory;
import jetbrains.mps.smodel.adapter.structure.MetaAdapterFactoryByName;
import jetbrains.mps.smodel.language.ConceptRegistry;
import jetbrains.mps.smodel.language.LanguageRegistry;
import jetbrains.mps.smodel.language.LanguageRegistryListener;
import jetbrains.mps.smodel.language.LanguageRuntime;
import jetbrains.mps.smodel.runtime.ConceptDescriptor;
import jetbrains.mps.smodel.runtime.StructureAspectDescriptor;
import jetbrains.mps.util.IterableUtil;
import jetbrains.mps.util.annotation.ToRemove;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.jetbrains.mps.openapi.language.SAbstractConcept;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * Based on generated structure descriptors
 */
public class ConceptDescendantsCache implements CoreComponent {
  private static ConceptDescendantsCache INSTANCE;

  private final MPSModuleRepository myModuleRepository;
  private final LanguageRegistry myLanguageRegistry;

  private final Map<LanguageRuntime, Set<ConceptDescriptor>> myLoadedLanguageToConceptsMap = new HashMap<LanguageRuntime, Set<ConceptDescriptor>>();
  private final Set<LanguageRuntime> myNotProcessedRuntimes = new LinkedHashSet<LanguageRuntime>();
  private final LanguageRegistryListener myLanguageRegistryListener = new LanguageRegistryListener() {
    @Override
    public void afterLanguagesLoaded(Iterable<LanguageRuntime> languages) {
      synchronized (myNotProcessedRuntimes) {
        myNotProcessedRuntimes.addAll(IterableUtil.asCollection(languages));
      }
    }

    @Override
    public void beforeLanguagesUnloaded(Iterable<LanguageRuntime> languages) {
      HashSet<LanguageRuntime> c = new HashSet<LanguageRuntime>(IterableUtil.asCollection(languages));
      synchronized (myNotProcessedRuntimes) {
        LinkedList<LanguageRuntime> notYetLoaded = new LinkedList<LanguageRuntime>(myNotProcessedRuntimes);
        notYetLoaded.retainAll(c);
        myNotProcessedRuntimes.removeAll(notYetLoaded);
        c.removeAll(notYetLoaded);
      }
      unloadConcepts(c);
    }
  };

  // SConceptId used as key here (instead of SConcept) because of hashCode() method implementation for sub-classes of SConcept.
  // Currently all implementors return 0 from the hashCode() method.
  // Values are unmodifiable linked hash sets to guarantee iteration order.
  // TODO: use SConcept as a key after reimplementing SConcept.hasCode() method properly.
  private final ConcurrentMap<SAbstractConcept, Set<SAbstractConcept>> myDescendantsCache = new ConcurrentHashMap<>();

  private void loadConcepts(Collection<LanguageRuntime> languages) {
    for (LanguageRuntime language : languages) {
      myLoadedLanguageToConceptsMap.put(language, getConcepts(language));
      for (ConceptDescriptor concept : myLoadedLanguageToConceptsMap.get(language)) {
        loadConcept(concept);
      }
    }
  }

  private void unloadConcepts(Collection<LanguageRuntime> languages) {
    for (LanguageRuntime language : languages) {
      Set<ConceptDescriptor> concepts = myLoadedLanguageToConceptsMap.get(language);
      if (concepts == null) {
        throw new IllegalArgumentException("No concepts registered for the language " + language);
      }
      for (ConceptDescriptor concept : concepts) {
        unloadConcept(concept);
      }
      myLoadedLanguageToConceptsMap.remove(language);
    }
  }

  private void loadConcept(ConceptDescriptor concept) {
    List<SConceptId> pids = concept.getParentsIds();
    SAbstractConcept abstractConcept = MetaAdapterFactory.getAbstractConcept(concept);
    for (SConceptId id : pids) {
      SAbstractConcept parentConcept = MetaAdapterFactory.getAbstractConcept(ConceptRegistry.getInstance().getConceptDescriptor(id));
      myDescendantsCache.compute(parentConcept, (key, values) -> add(values, abstractConcept));
    }
  }

  private void unloadConcept(ConceptDescriptor concept) {
    List<SConceptId> pids = concept.getParentsIds();

    SAbstractConcept abstractConcept = MetaAdapterFactory.getAbstractConcept(concept);
    for (SConceptId id : pids) {
      SAbstractConcept parentConcept = MetaAdapterFactory.getAbstractConcept(ConceptRegistry.getInstance().getConceptDescriptor(id));
      myDescendantsCache.compute(parentConcept, (key, values) -> remove(values, abstractConcept));
    }
  }

  @NotNull
  private static Set<SAbstractConcept> add(@Nullable Set<SAbstractConcept> values, @NotNull SAbstractConcept abstractConcept) {
    // linked guarantees iteration order
    Set<SAbstractConcept> descendants = new LinkedHashSet<>(values == null ? Collections.emptySet() : values);
    descendants.add(abstractConcept);
    return Collections.unmodifiableSet(descendants);
  }

  @Nullable
  private static Set<SAbstractConcept> remove(@Nullable Set<SAbstractConcept> values, @NotNull SAbstractConcept abstractConcept) {
    if (values == null) {
      return null;
    }

    Set<SAbstractConcept> descendants = new LinkedHashSet<>(values);
    if (!descendants.remove(abstractConcept)) {
      return values;
    }

    if (descendants.isEmpty()) {
      return null;
    }

    return Collections.unmodifiableSet(descendants);
  }

  public ConceptDescendantsCache(MPSModuleRepository moduleRepository, LanguageRegistry languageRegistry) {
    myModuleRepository = moduleRepository;
    myLanguageRegistry = languageRegistry;
  }

  @Override
  public void init() {
    if (INSTANCE != null) {
      throw new IllegalStateException("double initialization");
    }
    INSTANCE = this;

    myLanguageRegistry.addRegistryListener(myLanguageRegistryListener);
  }

  @Override
  public void dispose() {
    myLanguageRegistry.removeRegistryListener(myLanguageRegistryListener);

    INSTANCE = null;
  }

  public static ConceptDescendantsCache getInstance() {
    return INSTANCE;
  }

  /**
   * Collect all descendant concepts
   * Iteration order is guaranteed to be stable
   *
   * @param concept concept to start from
   * @return non-empty set of descendant concepts including the one supplied.
   */
  public Set<SAbstractConcept> getDescendants(SAbstractConcept concept) {
    loadConceptsFromNotProcessedRuntimes();
    //note that linked here guarantees the iteration order
    Set<SAbstractConcept> result = new LinkedHashSet<SAbstractConcept>();
    collectDescendants(concept, result);
    return result;
  }

  public Set<SAbstractConcept> getDirectDescendants(SAbstractConcept concept) {
    loadConceptsFromNotProcessedRuntimes();
    return getDirectDescendantsInternal(concept);
  }

  private void loadConceptsFromNotProcessedRuntimes() {
    synchronized (myNotProcessedRuntimes) {
      if (!myNotProcessedRuntimes.isEmpty()) {
        loadConcepts(myNotProcessedRuntimes);
        myNotProcessedRuntimes.clear();
      }
    }
  }

  private void collectDescendants(SAbstractConcept concept, Set<SAbstractConcept> result) {
    if (result.contains(concept)) return;
    result.add(concept);

    for (SAbstractConcept descendant : getDirectDescendantsInternal(concept)) {
      collectDescendants(descendant, result);
    }
  }

  @NotNull
  private Set<SAbstractConcept> getDirectDescendantsInternal(SAbstractConcept concept) {
    myModuleRepository.getModelAccess().checkReadAccess();
    Set<SAbstractConcept> result = myDescendantsCache.get(concept);
    return result != null ? result : Collections.<SAbstractConcept>emptySet();
  }

  private Set<ConceptDescriptor> getConcepts(LanguageRuntime languageRuntime) {
    StructureAspectDescriptor structureDescriptor = languageRuntime.getAspect(StructureAspectDescriptor.class);
    if (structureDescriptor == null) {
      return Collections.emptySet();
    }
    return new LinkedHashSet<ConceptDescriptor>(structureDescriptor.getDescriptors());
  }
}
