package com.hubspot.jinjava.lib.tag.eager;

import com.hubspot.jinjava.interpret.DeferredValue;
import com.hubspot.jinjava.interpret.DeferredValueException;
import com.hubspot.jinjava.interpret.InterpretException;
import com.hubspot.jinjava.interpret.JinjavaInterpreter;
import com.hubspot.jinjava.lib.tag.ForTag;
import com.hubspot.jinjava.objects.serialization.PyishObjectMapper;
import com.hubspot.jinjava.tree.TagNode;
import com.hubspot.jinjava.tree.parse.TagToken;
import com.hubspot.jinjava.util.EagerExpressionResolver;
import com.hubspot.jinjava.util.EagerExpressionResolver.EagerExpressionResult;
import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
import com.hubspot.jinjava.util.LengthLimitingStringJoiner;
import java.util.HashSet;
import java.util.List;
import java.util.Map.Entry;
import java.util.stream.Collectors;
import org.apache.commons.lang3.tuple.Pair;

public class EagerForTag extends EagerTagDecorator<ForTag> {

  public EagerForTag() {
    super(new ForTag());
  }

  public EagerForTag(ForTag forTag) {
    super(forTag);
  }

  @Override
  public String eagerInterpret(
    TagNode tagNode,
    JinjavaInterpreter interpreter,
    InterpretException e
  ) {
    LengthLimitingStringBuilder result = new LengthLimitingStringBuilder(
      interpreter.getConfig().getMaxOutputSize()
    );
    String prefix = "";

    // separate getEagerImage from renderChildren because the token gets evaluated once
    // while the children are evaluated 0...n times.
    result.append(
      executeInChildContext(
          eagerInterpreter ->
            EagerExpressionResult.fromString(
              getEagerImage(
                buildToken(tagNode, e, interpreter.getLineNumber()),
                eagerInterpreter
              )
            ),
          interpreter,
          true,
          false,
          false
        )
        .asTemplateString()
    );

    EagerExecutionResult eagerExecutionResult = executeInChildContext(
      eagerInterpreter -> {
        eagerInterpreter.getContext().put("loop", DeferredValue.instance());
        return EagerExpressionResult.fromString(
          renderChildren(tagNode, eagerInterpreter)
        );
      },
      interpreter,
      false,
      false,
      true
    );
    if (!eagerExecutionResult.getSpeculativeBindings().isEmpty()) {
      // Values cannot be modified within a for loop because we don't know many times, if any it will run
      prefix =
        buildSetTagForDeferredInChildContext(
          eagerExecutionResult
            .getSpeculativeBindings()
            .entrySet()
            .stream()
            .collect(
              Collectors.toMap(
                Entry::getKey,
                entry -> PyishObjectMapper.getAsPyishString(entry.getValue())
              )
            ),
          interpreter,
          true
        );
      eagerExecutionResult =
        executeInChildContext(
          eagerInterpreter -> {
            eagerInterpreter.getContext().put("loop", DeferredValue.instance());
            return EagerExpressionResult.fromString(
              renderChildren(tagNode, eagerInterpreter)
            );
          },
          interpreter,
          false,
          false,
          true
        );
      if (!eagerExecutionResult.getSpeculativeBindings().isEmpty()) {
        throw new DeferredValueException(
          "Modified values in deferred for loop: " +
          String.join(", ", eagerExecutionResult.getSpeculativeBindings().keySet())
        );
      }
    }
    result.append(eagerExecutionResult.asTemplateString());
    result.append(reconstructEnd(tagNode));
    return prefix + result;
  }

  @Override
  public String getEagerTagImage(TagToken tagToken, JinjavaInterpreter interpreter) {
    Pair<List<String>, String> loopVarsAndExpression = getTag()
      .getLoopVarsAndExpression(tagToken);
    List<String> loopVars = loopVarsAndExpression.getLeft();
    String loopExpression = loopVarsAndExpression.getRight();

    EagerExpressionResult eagerExpressionResult = EagerExpressionResolver.resolveExpression(
      loopExpression,
      interpreter
    );

    LengthLimitingStringJoiner joiner = new LengthLimitingStringJoiner(
      interpreter.getConfig().getMaxOutputSize(),
      " "
    );

    joiner
      .add(tagToken.getSymbols().getExpressionStartWithTag())
      .add(tagToken.getTagName())
      .add(String.join(", ", loopVars))
      .add("in")
      .add(eagerExpressionResult.toString())
      .add(tagToken.getSymbols().getExpressionEndWithTag());
    String newlyDeferredFunctionImages = reconstructFromContextBeforeDeferring(
      eagerExpressionResult.getDeferredWords(),
      interpreter
    );

    interpreter
      .getContext()
      .handleEagerToken(
        new EagerToken(
          new TagToken(
            joiner.toString(),
            tagToken.getLineNumber(),
            tagToken.getStartPosition(),
            tagToken.getSymbols()
          ),
          eagerExpressionResult
            .getDeferredWords()
            .stream()
            .filter(
              word -> !(interpreter.getContext().get(word) instanceof DeferredValue)
            )
            .collect(Collectors.toSet()),
          new HashSet<>(loopVars)
        )
      );
    return (newlyDeferredFunctionImages + joiner.toString());
  }
}

@@ -1942,493 +1942,179 @@
 t = 
-executeInChildContext(%0A      eagerInterpreter -%3E %7B%0A        eagerInterpreter.getContext().put(%22loop%22, DeferredValue.instance());%0A        return EagerExpressionResult.fromString(%0A          renderChildren(tagNode, eagerInterpreter)%0A        );%0A      %7D,%0A      interpreter,%0A      false,%0A      false,%0A      true%0A    );%0A    if (!eagerExecutionResult.getSpeculativeBindings().isEmpty()) %7B%0A      // Values cannot be modified within a for loop because we don't know many times, if any it will run
+runLoopOnce(tagNode, interpreter);%0A    if (!eagerExecutionResult.getSpeculativeBindings().isEmpty()) %7B%0A      // Defer any variables that we tried to modify during the loop
 %0A   
@@ -2561,370 +2561,41 @@
 lt =
-%0A        executeInChildContext(%0A          eagerInterpreter -%3E %7B%0A            eagerInterpreter.getContext().put(%22loop%22, DeferredValue.instance());%0A            return EagerExpressionResult.fromString(%0A              renderChildren(tagNode, eagerInterpreter)%0A            );%0A          %7D,%0A          interpreter,%0A          false,%0A          false,%0A          true%0A        
+ runLoopOnce(tagNode, interpreter
 );%0A 
@@ -3000,24 +3000,458 @@
 esult;%0A  %7D%0A%0A
+  private EagerExecutionResult runLoopOnce(%0A    TagNode tagNode,%0A    JinjavaInterpreter interpreter%0A  ) %7B%0A    return executeInChildContext(%0A      eagerInterpreter -%3E %7B%0A        eagerInterpreter.getContext().put(%22loop%22, DeferredValue.instance());%0A        return EagerExpressionResult.fromString(%0A          renderChildren(tagNode, eagerInterpreter)%0A        );%0A      %7D,%0A      interpreter,%0A      false,%0A      false,%0A      true%0A    );%0A  %7D%0A%0A
   @Override%0A

Extract common method