modeldb-main/backend/server/src/main/java/ai/verta/modeldb/versioning/blob/diff/DiffComputer.java
package ai.verta.modeldb.versioning.blob.diff;

import ai.verta.modeldb.versioning.DiffStatusEnum.DiffStatus;
import ai.verta.modeldb.versioning.autogenerated._public.modeldb.versioning.model.*;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

// TODO: autogenerate?
// TODO: equal comparison should compare just immediate A and B, instead of recursing. This happens
// because some diff types are not isolated
public class DiffComputer {
  public static <B, F, R> R computeDiff(
      B a, B b, Function<B, F> getter, Function3<F, F, R> computer) {
    return computer.apply(Utils.getOrNull(a, getter), Utils.getOrNull(b, getter));
  }

  // This applies an algorithm similar to what I discussed with Ravi for merges: gather sets based
  // on a key
  // 1. if there is a key collision and both sets have size 1, then compute the diff of those
  // elements
  // 2. if there is a key collision and both sets have more than 1 element, consider all A as
  // removal and B as addition, ignoring any modifications (by passing null values)
  // 3. if there is no key collision, then just process the right side
  public static <B, F, R extends ProtoType> List<R> computeListDiff(
      B a,
      B b,
      Function<B, List<F>> getter,
      Function<F, String> hasher,
      Function3<F, F, R> computer) {
    Map<String, HashSet<F>> mapA = a == null ? new HashMap<>() : toMap(getter.apply(a), hasher);
    Map<String, HashSet<F>> mapB = b == null ? new HashMap<>() : toMap(getter.apply(b), hasher);

    HashSet<String> keys = new HashSet<>();
    keys.addAll(mapA.keySet());
    keys.addAll(mapB.keySet());

    List<R> ret =
        keys.stream()
            .flatMap(
                key -> {
                  HashSet<F> elA = mapA.get(key);
                  HashSet<F> elB = mapB.get(key);
                  // Key collision and one element, process it
                  if (elA != null && elB != null && elA.size() == 1 && elB.size() == 1) {
                    final F nextA = elA.iterator().next();
                    final F nextB = elB.iterator().next();
                    if (nextA.equals(nextB)) {
                      return null;
                    }
                    return Stream.of(computer.apply(nextA, nextB));
                  }

                  // Key collision and more elements, consider removal + addition
                  if (elA != null && elB != null) {
                    return Stream.concat(
                        elA.stream().map(el -> computer.apply(el, null)),
                        elB.stream().map(el -> computer.apply(null, el)));
                  } else if (elA != null) {
                    return elA.stream().map(el -> computer.apply(el, null));
                  } else {
                    return elB.stream().map(el -> computer.apply(null, el));
                  }
                })
            .map(Utils::removeEmpty)
            .filter(
                Objects
                    ::nonNull) // Remove elements that became null in the process of computing the
            // diff for some reason
            .collect(Collectors.toList());

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

  public static <F> Map<String, HashSet<F>> toMap(List<F> input, Function<F, String> hasher) {
    if (input == null) return new HashMap<>();
    return input.stream()
        .collect(
            Collectors.toMap(
                hasher,
                el -> new HashSet<>(Collections.singletonList(el)),
                (entry1, entry2) -> {
                  HashSet<F> result2 = new HashSet<>(entry1);
                  result2.addAll(entry2);
                  return result2;
                }));
  }

  public static <T> AutogenDiffStatusEnumDiffStatus getStatus(T a, T b) {
    return AutogenDiffStatusEnumDiffStatus.fromProto(status(a, b));
  }

  private static <T> DiffStatus status(T a, T b) {
    if (a == null && b == null) {
      return DiffStatus.DELETED;
    }
    if (a == null) {
      return DiffStatus.ADDED;
    }
    if (b == null) {
      return DiffStatus.DELETED;
    }
    return DiffStatus.MODIFIED;
  }

  public static AutogenBlobDiff computeBlobDiff(AutogenBlob a, AutogenBlob b) {
    return Utils.removeEmpty(
        new AutogenBlobDiff()
            .setCode(computeDiff(a, b, AutogenBlob::getCode, DiffComputer::computeCodeDiff))
            .setConfig(computeDiff(a, b, AutogenBlob::getConfig, DiffComputer::computeConfigDiff))
            .setDataset(
                computeDiff(a, b, AutogenBlob::getDataset, DiffComputer::computeDatasetDiff))
            .setEnvironment(
                computeDiff(
                    a, b, AutogenBlob::getEnvironment, DiffComputer::computeEnvironmentDiff)));
  }

  public static AutogenCodeDiff computeCodeDiff(AutogenCodeBlob a, AutogenCodeBlob b) {
    return Utils.removeEmpty(
        new AutogenCodeDiff()
            .setGit(computeDiff(a, b, AutogenCodeBlob::getGit, DiffComputer::computeGitCodeDiff))
            .setNotebook(
                computeDiff(
                    a, b, AutogenCodeBlob::getNotebook, DiffComputer::computeNotebookCodeDiff)));
  }

  public static AutogenGitCodeDiff computeGitCodeDiff(AutogenGitCodeBlob a, AutogenGitCodeBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;

    return Utils.removeEmpty(new AutogenGitCodeDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenNotebookCodeDiff computeNotebookCodeDiff(
      AutogenNotebookCodeBlob a, AutogenNotebookCodeBlob b) {
    return Utils.removeEmpty(
        new AutogenNotebookCodeDiff()
            .setGitRepo(
                computeDiff(
                    a, b, AutogenNotebookCodeBlob::getGitRepo, DiffComputer::computeGitCodeDiff))
            .setPath(
                computeDiff(
                    a,
                    b,
                    AutogenNotebookCodeBlob::getPath,
                    DiffComputer::computePathDatasetComponentDiff)));
  }

  public static AutogenConfigDiff computeConfigDiff(AutogenConfigBlob a, AutogenConfigBlob b) {
    return Utils.removeEmpty(
        new AutogenConfigDiff()
            .setHyperparameters(
                computeListDiff(
                    a,
                    b,
                    AutogenConfigBlob::getHyperparameters,
                    AutogenHyperparameterConfigBlob::getName,
                    DiffComputer::computeHyperparameterConfigDiff))
            .setHyperparameterSet(
                computeListDiff(
                    a,
                    b,
                    AutogenConfigBlob::getHyperparameterSet,
                    AutogenHyperparameterSetConfigBlob::getName,
                    DiffComputer::computeHyperparameterSetConfigDiff)));
  }

  public static AutogenHyperparameterConfigDiff computeHyperparameterConfigDiff(
      AutogenHyperparameterConfigBlob a, AutogenHyperparameterConfigBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenHyperparameterConfigDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenHyperparameterSetConfigDiff computeHyperparameterSetConfigDiff(
      AutogenHyperparameterSetConfigBlob a, AutogenHyperparameterSetConfigBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenHyperparameterSetConfigDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenDatasetDiff computeDatasetDiff(AutogenDatasetBlob a, AutogenDatasetBlob b) {
    return Utils.removeEmpty(
        new AutogenDatasetDiff()
            .setPath(
                computeDiff(
                    a, b, AutogenDatasetBlob::getPath, DiffComputer::computePathDatasetDiff))
            .setS3(computeDiff(a, b, AutogenDatasetBlob::getS3, DiffComputer::computeS3DatasetDiff))
            .setQuery(
                computeDiff(
                    a, b, AutogenDatasetBlob::getQuery, DiffComputer::computeQueryDatasetDiff)));
  }

  public static AutogenPathDatasetDiff computePathDatasetDiff(
      AutogenPathDatasetBlob a, AutogenPathDatasetBlob b) {
    return Utils.removeEmpty(
        new AutogenPathDatasetDiff()
            .setComponents(
                computeListDiff(
                    a,
                    b,
                    AutogenPathDatasetBlob::getComponents,
                    AutogenPathDatasetComponentBlob::getPath,
                    DiffComputer::computePathDatasetComponentDiff)));
  }

  public static AutogenPathDatasetComponentDiff computePathDatasetComponentDiff(
      AutogenPathDatasetComponentBlob a, AutogenPathDatasetComponentBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenPathDatasetComponentDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenQueryDatasetDiff computeQueryDatasetDiff(
      AutogenQueryDatasetBlob a, AutogenQueryDatasetBlob b) {
    return Utils.removeEmpty(
        new AutogenQueryDatasetDiff()
            .setComponents(
                computeListDiff(
                    a,
                    b,
                    AutogenQueryDatasetBlob::getComponents,
                    x -> x.toString(), // Revisit
                    DiffComputer::computeQueryDatasetComponentDiff)));
  }

  public static AutogenQueryDatasetComponentDiff computeQueryDatasetComponentDiff(
      AutogenQueryDatasetComponentBlob a, AutogenQueryDatasetComponentBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenQueryDatasetComponentDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenS3DatasetDiff computeS3DatasetDiff(
      AutogenS3DatasetBlob a, AutogenS3DatasetBlob b) {
    return Utils.removeEmpty(
        new AutogenS3DatasetDiff()
            .setComponents(
                computeListDiff(
                    a,
                    b,
                    AutogenS3DatasetBlob::getComponents,
                    x -> x.getPath().getPath(),
                    DiffComputer::computeS3DatasetComponentDiff)));
  }

  public static AutogenS3DatasetComponentDiff computeS3DatasetComponentDiff(
      AutogenS3DatasetComponentBlob a, AutogenS3DatasetComponentBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenS3DatasetComponentDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenEnvironmentDiff computeEnvironmentDiff(
      AutogenEnvironmentBlob a, AutogenEnvironmentBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenEnvironmentDiff()
            .setCommandLine(
                computeDiff(
                    a,
                    b,
                    AutogenEnvironmentBlob::getCommandLine,
                    DiffComputer::computeCommandLineEnvironmentDiff))
            .setDocker(
                computeDiff(
                    a,
                    b,
                    AutogenEnvironmentBlob::getDocker,
                    DiffComputer::computeDockerEnvironmentDiff))
            .setPython(
                computeDiff(
                    a,
                    b,
                    AutogenEnvironmentBlob::getPython,
                    DiffComputer::computePythonEnvironmentDiff))
            .setEnvironmentVariables(
                computeListDiff(
                    a,
                    b,
                    AutogenEnvironmentBlob::getEnvironmentVariables,
                    AutogenEnvironmentVariablesBlob::getName,
                    DiffComputer::computeEnvironmentVariablesDiff)));
  }

  public static AutogenCommandLineEnvironmentDiff computeCommandLineEnvironmentDiff(
      List<String> a, List<String> b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenCommandLineEnvironmentDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenDockerEnvironmentDiff computeDockerEnvironmentDiff(
      AutogenDockerEnvironmentBlob a, AutogenDockerEnvironmentBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenDockerEnvironmentDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenPythonEnvironmentDiff computePythonEnvironmentDiff(
      AutogenPythonEnvironmentBlob a, AutogenPythonEnvironmentBlob b) {
    return Utils.removeEmpty(
        new AutogenPythonEnvironmentDiff()
            .setVersion(
                computeDiff(
                    a,
                    b,
                    AutogenPythonEnvironmentBlob::getVersion,
                    DiffComputer::computeVersionEnvironmentDiff))
            .setConstraints(
                computeListDiff(
                    a,
                    b,
                    AutogenPythonEnvironmentBlob::getConstraints,
                    e -> e.getLibrary() + e.getConstraint(),
                    DiffComputer::computePythonRequirementEnvironmentDiff))
            .setRequirements(
                computeListDiff(
                    a,
                    b,
                    AutogenPythonEnvironmentBlob::getRequirements,
                    e -> e.getLibrary() + e.getConstraint(),
                    DiffComputer::computePythonRequirementEnvironmentDiff)));
  }

  public static AutogenVersionEnvironmentDiff computeVersionEnvironmentDiff(
      AutogenVersionEnvironmentBlob a, AutogenVersionEnvironmentBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenVersionEnvironmentDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenPythonRequirementEnvironmentDiff computePythonRequirementEnvironmentDiff(
      AutogenPythonRequirementEnvironmentBlob a, AutogenPythonRequirementEnvironmentBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;

    return Utils.removeEmpty(
        new AutogenPythonRequirementEnvironmentDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }

  public static AutogenEnvironmentVariablesDiff computeEnvironmentVariablesDiff(
      AutogenEnvironmentVariablesBlob a, AutogenEnvironmentVariablesBlob b) {
    if (a == null && b == null) return null;
    if ((a != null && a.equals(b)) || (b != null && b.equals(a))) return null;
    return Utils.removeEmpty(
        new AutogenEnvironmentVariablesDiff().setA(a).setB(b).setStatus(getStatus(a, b)));
  }
}
