/*
 * YamlTree.java
 *
 * Copyright (C) 2022 by Posit Software, PBC
 *
 * Unless you have received this program directly from Posit Software pursuant
 * to the terms of a commercial license agreement with Posit Software, then
 * this program is licensed to you under the terms of version 3 of the
 * GNU Affero General Public License. This program is distributed WITHOUT
 * ANY EXPRESS OR IMPLIED WARRANTY, INCLUDING THOSE OF NON-INFRINGEMENT,
 * MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. Please refer to the
 * AGPL (http://www.gnu.org/licenses/agpl-3.0.txt) for more details.
 *
 */
package org.rstudio.studio.client.rmarkdown.model;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.rstudio.core.client.Debug;
import org.rstudio.core.client.StringUtil;

import com.google.gwt.regexp.shared.MatchResult;
import com.google.gwt.regexp.shared.RegExp;

// NOTE: This class does not fully parse YAML fields; it just provides some facilities
// for inspecting and extracting structured pieces of a YAML tree.
public class YamlTree
{
   private class YamlTreeNode
   {
      public YamlTreeNode(String line)
      {
         // detect comment and blank lines before extracting key, so that
         // comments with colons (e.g. "# Note: important") don't produce
         // bogus keys
         String trimmed = line.trim();
         isBlank = trimmed.isEmpty();
         isComment = !isBlank && trimmed.startsWith("#");

         // the YAML package outputs tildes to represent null values; these
         // are valid YAML but "null" is easier to read for R users
         if (!isComment && !isBlank && line.endsWith(": ~"))
         {
            line = StringUtil.substring(line, 0, line.length() - 1) + "null";
         }
         yamlLine = line;

         key = (isComment || isBlank) ? "" : getKey(line);
         indentLevel = getIndentLevel();
      }
      
      // add a child node; returns true iff the node was added (we may choose
      // to treat the child as a continuation)
      public boolean addChild(YamlTreeNode child)
      {
         // if this node has a key, is a comment, or is a blank line, add
         // it as a tree node; otherwise, add its line data to this node
         // as a multi-line continuation
         if (child.key.length() > 0 || child.isComment || child.isBlank)
         {
            children.add(child);
            child.parent = this;
            return true;
         }

         yamlLine += ("\n" + child.yamlLine);
         return false;
      }
      
      public void adopt(YamlTree tree)
      {
         children.clear();
         adoptNode(tree.root_, true);
      }
      
      public void adoptNode(YamlTreeNode node, boolean isRoot)
      {
         for (YamlTreeNode child: node.children)
         {
            // increase the indentation level of the children to match our
            // own; blank lines stay empty (no spurious whitespace)
            if (!child.isBlank)
            {
               child.indentLevel += indentLevel + 2;
               child.yamlLine = getIndent() + "  " + child.yamlLine;
            }
            
            // if this is the root node, adopt is children as our own
            if (isRoot)
            {
               children.add(child);
            }
            adoptNode(child, false);
         }
      }
      
      public String getIndent()
      {
         // consider the list element indicator (-) to be part of the node's
         // indentation, to prevent list continuations from being treated as
         // sibling nodes
         RegExp whitespace = RegExp.compile("^\\s*-?\\s*");
         MatchResult result = whitespace.exec(yamlLine);
         if (result == null)
            return "";
         else
            return result.getGroup(0);
      }
      
      public String getValue()
      {
         int idx = yamlLine.indexOf(":");
         if (idx < 0)
            return "";
         return StringUtil.substring(yamlLine, idx + 2).trim();
      }
      
      public void setValue(String value)
      {
         int idx = yamlLine.indexOf(":");
         
         // can't set a value if we don't currently have one
         if (idx < 0)
            return;

         // if the ":" is at the end, ensure there's a space preceding the value
         if (idx == yamlLine.length() - 1)
            value = " " + value;

         yamlLine = StringUtil.substring(yamlLine, 0, idx + 2) + value;
      }
      
      public String yamlLine;
      public String key;
      public int indentLevel = 0;
      public boolean isComment = false;
      public boolean isBlank = false;
      public YamlTreeNode parent = null;
      public List<YamlTreeNode> children = new ArrayList<>();
      
      private String getKey(String line)
      {
         RegExp keyReg = RegExp.compile("^\\s*([^:]+(?:::[^:]+)?):");
         MatchResult result = keyReg.exec(line);
         if (result == null)
            return "";
          else
            return result.getGroup(1).trim();
      }

      private int getIndentLevel()
      {
         return getIndent().length();
      }
   }
   
   public YamlTree(String yaml)
   {
      root_ = createYamlTree(yaml);
      keyMap_ = new HashMap<>();
      createKeyMap(root_, keyMap_);
   }
   
   public String toString()
   {
      return yamlFromTree(root_);
   }
   
   public void reorder(List<String> orderedKeys)
   {
      if (!keyMap_.containsKey(orderedKeys.get(0)))
         return;
      YamlTreeNode parent = keyMap_.get(orderedKeys.get(0)).parent;

      // Group children into blocks: each block is a run of leading
      // comment/blank nodes followed by a data node. Comments "stick"
      // to the data node they precede so they move together on reorder.
      List<List<YamlTreeNode>> groups = new ArrayList<>();
      List<YamlTreeNode> current = new ArrayList<>();
      for (YamlTreeNode child : parent.children)
      {
         current.add(child);
         if (!child.isComment && !child.isBlank)
         {
            groups.add(current);
            current = new ArrayList<>();
         }
      }
      // trailing comments/blanks not followed by a data node
      List<YamlTreeNode> trailing = current;

      // Swap groups into the positions specified by orderedKeys
      for (int i = 0; i < orderedKeys.size() && i < groups.size(); i++)
      {
         for (int j = i; j < groups.size(); j++)
         {
            List<YamlTreeNode> group = groups.get(j);
            YamlTreeNode dataNode = group.get(group.size() - 1);
            if (orderedKeys.get(i) == dataNode.key)
            {
               List<YamlTreeNode> temp = groups.get(i);
               groups.set(i, groups.get(j));
               groups.set(j, temp);
               break;
            }
         }
      }

      // Flatten groups back into the children list
      parent.children.clear();
      for (List<YamlTreeNode> group : groups)
      {
         parent.children.addAll(group);
      }
      parent.children.addAll(trailing);
   }

   public List<String> getChildKeys(String parentKey)
   {
      try
      {
         if (!keyMap_.containsKey(parentKey))
            return null;
         YamlTreeNode parent = keyMap_.get(parentKey);
         ArrayList<String> result = new ArrayList<>();
         for (YamlTreeNode child: parent.children)
         {
            if (!child.isComment && !child.isBlank)
               result.add(child.key);
         }
         return result;
      }
      catch (Exception e)
      {
         // case 3826: it's not clear why but we've seen an exception thrown
         // from the above in non-reproducible circumstances; handle gracefully
         Debug.log("YamlTree: Couldn't get child keys for " + parentKey + 
                   " (" + e.getMessage() + ")");
      }
      return null;
   }
   
   // add a simple YAML value to the tree under the given parent, or under
   // the root if parentKey is null.
   public void addYamlValue(String parentKey, String key, String value)
   {
      // if a key was specified but doesn't exist in the flattened map, abort
      if (parentKey != null && !keyMap_.containsKey(parentKey))
         return;
      
      String line = key + ": " + value;
      YamlTreeNode parent = parentKey == null ? root_ : keyMap_.get(parentKey);
      
      // check to see if we should be replacing a value rather than adding
      YamlTreeNode child = null;
      for (YamlTreeNode childNode: parent.children)
      {
         if (childNode.key == key)
         {
            // found an existing child node
            child = childNode;
            
            // set its value to the one supplied by the caller and return
            child.setValue(value);
            return;
         }
      }
      
      // no existing value, create a wholly new node
      child = parentKey == null
         ? new YamlTreeNode(line)
         : new YamlTreeNode(parent.getIndent() + "  " + line);

      keyMap_.put(key, child);
      parent.addChild(child);
   }
   
   
   public String getKeyValue(String key)
   {
      if (keyMap_.containsKey(key))
         return keyMap_.get(key).getValue();
      return "";
   }
   
   public String getChildValue(String parentKey, String childKey)
   {
      for (YamlTreeNode parent: root_.children)
      {
         if (parent.key == parentKey)
         {
            for (YamlTreeNode child: parent.children)
            {
               if (child.key == childKey)
               {
                  return child.getValue();
               }
            }
         }
      }
      return "";
   }
   
   public String getGrandchildValue(String grandparentKey, String parentKey, String childKey)
   {
      for (YamlTreeNode grandparent: root_.children)
      {
         if (grandparent.key == grandparentKey)
         {
            for (YamlTreeNode parent: grandparent.children)
            {
               if (parent.key == parentKey)
               {
                  for (YamlTreeNode child: parent.children)
                  {
                     if (child.key == childKey)
                     {
                        return child.getValue();
                     }
                  }
               }
            }
         }
      }
      return "";
   }
   
   public void setGrandchildValue(String grandparentKey, String parentKey, String childKey, String value)
   {
      for (YamlTreeNode grandparent: root_.children)
      {
         if (grandparent.key == grandparentKey)
         {
            for (YamlTreeNode parent: grandparent.children)
            {
               if (parent.key == parentKey)
               {
                  for (YamlTreeNode child: parent.children)
                  {
                     if (child.key == childKey)
                     {
                        child.setValue(value);
                        return;
                     }
                  }
                  // if we didn't find an existing value then add a new one
                  YamlTreeNode newChild = new YamlTreeNode(parent.getIndent() + "  " + childKey + ": " + value);
                  keyMap_.put(childKey, newChild);
                  parent.addChild(newChild);
                  return;
               }
            }
         }
      }
   }
   
   public boolean containsKey(String key)
   {
      return keyMap_.containsKey(key);
   }

   public void setKeyValue(String key, String value)
   {
      if (keyMap_.containsKey(key))
         keyMap_.get(key).setValue(value);
   }
   
   public void setKeyValue(String key, YamlTree subtree)
   {
      if (keyMap_.containsKey(key))
      {
         YamlTreeNode node = keyMap_.get(key);
         node.adopt(subtree);
         node.setValue("");
         keyMap_.clear();
         createKeyMap(root_, keyMap_);
      }
   }
   
   public void clearChildren(String key)
   {
      if (keyMap_.containsKey(key))
      {
         keyMap_.get(key).children.clear();
         keyMap_.clear();
         createKeyMap(root_, keyMap_);
      }
   }
   
   public static boolean isTrue(String value)
   {
      return value == "y"    || value == "yes" || 
             value == "true" || value == "on";
   }
   
   private static void flushPending(List<YamlTreeNode> pending,
                                    YamlTreeNode parent)
   {
      for (YamlTreeNode node : pending)
         parent.addChild(node);
      pending.clear();
   }

   private static YamlTreeNode findParentForIndent(
         int childIndent,
         int currentIndent,
         YamlTreeNode currentParent,
         YamlTreeNode lastNode,
         YamlTreeNode root)
   {
      if (childIndent > currentIndent)
      {
         return lastNode;
      }
      else if (childIndent < currentIndent)
      {
         YamlTreeNode parent = currentParent;
         do
         {
            parent = parent.parent;
         }
         while (parent != null && parent.indentLevel >= childIndent);
         return parent != null ? parent : root;
      }
      return currentParent;
   }

   private YamlTreeNode createYamlTree(String yaml)
   {
      String[] lines = yaml.split("\n");
      YamlTreeNode root = new YamlTreeNode("");
      root.indentLevel = -2;
      YamlTreeNode currentParent = root;
      YamlTreeNode lastNode = root;
      int currentIndentLevel = 0;
      List<YamlTreeNode> pending = new ArrayList<>();
      for (String line: lines)
      {
         YamlTreeNode child = new YamlTreeNode(line);

         // buffer comments and blank lines so they attach to the next
         // data node's parent; neither affects tree-building state
         if (child.isComment || child.isBlank)
         {
            pending.add(child);
            continue;
         }

         currentParent = findParentForIndent(
               child.indentLevel,
               currentIndentLevel,
               currentParent,
               lastNode,
               root);

         flushPending(pending, currentParent);
         currentIndentLevel = child.indentLevel;
         if (currentParent.addChild(child))
            lastNode = child;
      }

      // flush any trailing comments/blanks
      flushPending(pending, currentParent);
      return root;
   }
   
   private static String yamlFromTree(YamlTreeNode root)
   {
      String yaml = "";
      for (YamlTreeNode child: root.children)
      {
         yaml += (child.yamlLine + '\n');
         yaml += yamlFromTree(child);
      }
      return yaml;
   }
   
   private static void createKeyMap(YamlTreeNode root,
                                    Map<String, YamlTreeNode> output)
   {
      for (YamlTreeNode child: root.children)
      {
         String key = child.key;
         // add this key if it doesn't already exist, or if it does exist and
         // this key is closer to the root than the existing key
         // (comments and blanks are skipped implicitly: their key is "")
         if (key.length() > 0 &&
             (!output.containsKey(key) ||
              child.indentLevel < output.get(key).indentLevel))
         {
            output.put(child.key, child);
         }
         createKeyMap(child, output);
      }
   }
   
   // A flattened version of the tree as a list of key-value pairs. In the 
   // case where multiple keys exist, the key closest to the root is stored.
   private final Map<String, YamlTreeNode> keyMap_;
   private final YamlTreeNode root_;
}
