// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package com.intellij.tools.build.bazel.jvmIncBuilder.impl;

import com.intellij.tools.build.bazel.jvmIncBuilder.NodeSourceSnapshot;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.jps.dependency.DataReader;
import org.jetbrains.jps.dependency.GraphDataInput;
import org.jetbrains.jps.dependency.GraphDataOutput;
import org.jetbrains.jps.dependency.NodeSource;

import java.io.IOException;
import java.util.Collections;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.Map;

import static org.jetbrains.jps.util.Iterators.count;

public class SourceSnapshotImpl implements NodeSourceSnapshot {
  private final Map<NodeSource, String> mySources;

  public SourceSnapshotImpl(Map<NodeSource, String> digestSources) {
    // preserve the original map implementation, which may provide a specific iteration order
    mySources = Collections.unmodifiableMap(digestSources);
  }
  
  public SourceSnapshotImpl(GraphDataInput in, DataReader<? extends NodeSource> sourceReader) throws IOException {
    Map<NodeSource, String> sources = new LinkedHashMap<>(); // preserve iteration order
    int count = in.readInt();
    while (count-- > 0) {
      String digest = in.readUTF();
      sources.put(sourceReader.load(in), digest);
    }
    mySources = Collections.unmodifiableMap(sources);
  }

  @Override
  public @NotNull Iterable<@NotNull NodeSource> getElements() {
    return mySources.keySet();
  }

  @Override
  public @NotNull String getDigest(NodeSource src) {
    return mySources.getOrDefault(src, "");
  }

  @Override
  public void write(GraphDataOutput out) throws IOException {
    out.writeInt(mySources.size());
    for (Map.Entry<NodeSource, String> entry : mySources.entrySet()) {
      out.writeUTF(entry.getValue());
      entry.getKey().write(out);
    }
  }

  public static NodeSourceSnapshot composite(Iterable<? extends NodeSourceSnapshot> groups) {

    if (count(groups) <= 1) {
      Iterator<? extends NodeSourceSnapshot> it = groups.iterator();
      return it.hasNext()? it.next() : EMPTY;
    }

    Map<NodeSource, String> data = new LinkedHashMap<>();
    for (NodeSourceSnapshot group : groups) {
      for (NodeSource element : group.getElements()) {
        data.put(element, group.getDigest(element));
      }
    }
    return new SourceSnapshotImpl(data);
  }
}
