package com.intellij.debugger.streams.trace.smart.handler;

import com.intellij.debugger.streams.trace.EvaluateExpressionTracerBase;
import com.intellij.debugger.streams.trace.smart.handler.type.ClassTypeImpl;
import com.intellij.debugger.streams.trace.smart.handler.type.GenericType;
import com.intellij.debugger.streams.wrapper.StreamCall;
import org.jetbrains.annotations.NotNull;

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

/**
 * @author Vitaliy.Bibaev
 */
public class DistinctHandler extends HandlerBase {
  private final PeekTracerHandler myPeekTracer;
  private final HashMapVariableImpl myStoreMapVariable;
  private final HashMapVariableImpl myResolveMapVariable;

  public DistinctHandler(int callNumber, @NotNull String callName) {
    myPeekTracer = new PeekTracerHandler(callNumber, callName);

    final String variablePrefix = callName + callNumber;
    myStoreMapVariable =
      new HashMapVariableImpl(variablePrefix + "Store", GenericType.OBJECT, new ClassTypeImpl("Map<Integer, Object>"), false);
    myResolveMapVariable = new HashMapVariableImpl(variablePrefix + "Resolve", GenericType.INT, GenericType.INT, false);
  }

  @NotNull
  @Override
  public List<StreamCall> additionalCallsBefore() {
    final List<StreamCall> result = new ArrayList<>(myPeekTracer.additionalCallsBefore());

    final PeekCall storeCall = new PeekCall(createStoreLambda());
    result.add(storeCall);
    return result;
  }

  @NotNull
  @Override
  public List<StreamCall> additionalCallsAfter() {
    final List<StreamCall> result = new ArrayList<>(myPeekTracer.additionalCallsAfter());
    result.add(new PeekCall(createResolveLambda()));
    return result;
  }

  @NotNull
  @Override
  public String prepareResult() {
    final String peekPrepare = myPeekTracer.prepareResult();
    final String resolve2Array = myResolveMapVariable.convertToArray("resolveArray", true, true);
    final String peekResult =
      "final Object peekResult = " + myPeekTracer.getResultExpression() + ";" + EvaluateExpressionTracerBase.LINE_SEPARATOR;
    return peekPrepare + resolve2Array + peekResult;
  }

  @NotNull
  @Override
  public String getResultExpression() {
    return "new Object[] { peekResult, resolveArray }";
  }

  @NotNull
  private String createStoreLambda() {
    final String storeMap = myStoreMapVariable.getName();
    return "x -> " + String.format("%s.computeIfAbsent(x, y -> new LinkedHashMap<>()).put(time.get(), x)", storeMap);
  }

  @NotNull
  private String createResolveLambda() {
    final String newLine = EvaluateExpressionTracerBase.LINE_SEPARATOR;
    final String storeMap = myStoreMapVariable.getName();
    final String resolveMap = myResolveMapVariable.getName();

    return "x -> {" + newLine +
           "final Map<Integer, Object> objects = " + String.format("%s.get(x);", storeMap) + newLine +
           "for (final int key: objects.keySet()) {" + newLine +
           "final Object value = objects.get(key);" + newLine +
           "if (value == x) {" + newLine +
           String.format("%s.put(key, time.get());", resolveMap) + newLine +
           "    }" + newLine +
           "  }" + newLine +
           "}" + newLine;
  }

  @NotNull
  @Override
  protected List<Variable> getVariables() {
    final List<Variable> variables = new ArrayList<>(Arrays.asList(myStoreMapVariable, myResolveMapVariable));
    variables.addAll(myPeekTracer.getVariables());
    return variables;
  }
}

@@ -681,32 +681,103 @@
 leImpl myResolve
+DirectMapVariable;%0A  private final HashMapVariableImpl myResolveReverse
 MapVariable;%0A%0A  
@@ -1122,32 +1122,38 @@
 );%0A    myResolve
+Direct
 MapVariable = ne
@@ -1199,16 +1199,157 @@
 %22Resolve
+Direct%22, GenericType.INT, GenericType.INT, false);%0A    myResolveReverseMapVariable = new HashMapVariableImpl(variablePrefix + %22ResolveReverse
 %22, Gener
@@ -2053,24 +2053,30 @@
 ring resolve
+Direct
 2Array = myR
@@ -2073,32 +2073,38 @@
 rray = myResolve
+Direct
 MapVariable.conv
@@ -2122,21 +2122,136 @@
 %22resolve
-Array
+Direct%22, true, true);%0A    final String resolveReverse2Array = myResolveReverseMapVariable.convertToArray(%22resolveReverse
 %22, true,
@@ -2445,16 +2445,45 @@
  resolve
+Direct2Array + resolveReverse
 2Array +
@@ -2609,21 +2609,38 @@
  resolve
-Array
+Direct, resolveReverse
  %7D%22;%0A  %7D
@@ -3078,16 +3078,22 @@
  resolve
+Direct
 Map = my
@@ -3091,32 +3091,114 @@
 tMap = myResolve
+DirectMapVariable.getName();%0A    final String resolveReverseMap = myResolveReverse
 MapVariable.getN
@@ -3484,23 +3484,137 @@
     
-%22if (value == x
+resolveReverseMap + %22.put(key, time.get());%22 + newLine +%0A           %22if (value == x && !%22 + resolveDirectMap + %22.containsKey(key)
 ) %7B%22
@@ -3686,16 +3686,22 @@
  resolve
+Direct
 Map) + n
@@ -3903,16 +3903,22 @@
 iables =
+%0A     
  new Arr
@@ -3969,16 +3969,51 @@
 yResolve
+DirectMapVariable, myResolveReverse
 MapVaria

Add extra information for the trace of distinct call