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 myResolveDirectMapVariable;
  private final HashMapVariableImpl myResolveReverseMapVariable;

  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);
    myResolveDirectMapVariable = new HashMapVariableImpl(variablePrefix + "ResolveDirect", GenericType.INT, GenericType.INT, false);
    myResolveReverseMapVariable = new HashMapVariableImpl(variablePrefix + "ResolveReverse", 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 resolveDirect2Array = myResolveDirectMapVariable.convertToArray("resolveDirect", true, true);
    final String resolveReverse2Array = myResolveReverseMapVariable.convertToArray("resolveReverse", true, true);
    final String peekResult =
      "final Object peekResult = " + myPeekTracer.getResultExpression() + ";" + EvaluateExpressionTracerBase.LINE_SEPARATOR;
    return peekPrepare + resolveDirect2Array + resolveReverse2Array + peekResult;
  }

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

  @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 resolveDirectMap = myResolveDirectMapVariable.getName();
    final String resolveReverseMap = myResolveReverseMapVariable.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 +
           resolveReverseMap + ".put(key, time.get());" + newLine +
           "if (value == x && !" + resolveDirectMap + ".containsKey(key)) {" + newLine +
           String.format("%s.put(key, time.get());", resolveDirectMap) + newLine +
           "    }" + newLine +
           "  }" + newLine +
           "}" + newLine;
  }

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

@@ -1968,24 +1968,96 @@
 eResult() %7B%0A
+    final String newLine = EvaluateExpressionTracerBase.LINE_SEPARATOR;%0A
     final St
@@ -2101,24 +2101,25 @@
 reResult();%0A
+%0A
     final St
@@ -2126,30 +2126,31 @@
 ring resolve
-Direct
+Reverse
 2Array = myR
@@ -2147,38 +2147,39 @@
 rray = myResolve
-Direct
+Reverse
 MapVariable.conv
@@ -2197,22 +2197,23 @@
 %22resolve
-Direct
+Reverse
 %22, true,
@@ -2241,104 +2241,944 @@
 ing 
-resolveReverse2Array = myResolveReverseMapVariable.convertToArray(%22resolveReverse%22, true, true);
+storeMapName = myStoreMapVariable.getName();%0A    final String afterMapName = myPeekTracer.getAfterMapName();%0A    final String prepareDirectMap = %22%7B%22 + newLine +%0A                                    %22  for (final int timeAfter : %22 + myResolveReverseMapVariable.getName() + %22.keySet()) %7B%22 + newLine +%0A                                    %22    final Object afterValue = %22 + afterMapName + %22.get(timeAfter);%22 + newLine +%0A                                    %22    final Map%3CInteger, Object%3E valuesBefore = %22 + storeMapName + %22.get(afterValue);%22 + newLine +%0A                                    %22    for (final int timeBefore : valuesBefore.keySet()) %7B%22 + newLine +%0A                                    %22      %22 + myResolveDirectMapVariable.getName() + %22.put(timeBefore, timeAfter);%22 + newLine +%0A                                    %22    %7D%22 + newLine +%0A                                    %22  %7D%22 + newLine +%0A                                    %22%7D%22;%0A
 %0A   
@@ -3337,26 +3337,156 @@
     
-return peekPrepare
+final String resolveDirect2Array = myResolveDirectMapVariable.convertToArray(%22resolveDirect%22, true, true);%0A    return peekPrepare + prepareDirectMap
  + r
@@ -4104,82 +4104,8 @@
 ();%0A
-    final String resolveDirectMap = myResolveDirectMapVariable.getName();%0A
     
@@ -4445,76 +4445,8 @@
 e +%0A
-           resolveReverseMap + %22.put(key, time.get());%22 + newLine +%0A
     
@@ -4483,22 +4483,23 @@
  resolve
-Direct
+Reverse
 Map + %22.
@@ -4568,13 +4568,8 @@
 put(
-key, 
 time
@@ -4574,16 +4574,21 @@
 me.get()
+, key
 );%22, res
@@ -4591,22 +4591,23 @@
  resolve
-Direct
+Reverse
 Map) + n

Fix small issues in generated code for distinct call