import "ansiLib" as ansiLib;

def colorEscapes:
  def parseJQColours:
    (. | split(":")) |
    {
      "null": .[0],
      "false": .[1],
      "true": .[2],
      "number": .[3],
      "string": .[4],
      "array": .[5],
      "object": .[6],
      "objectKey": .[7],
    };

  def parseGOJQColours:
    (. | split(":")) |
    {
      "null": .[0],
      "false": .[1],
      "true": .[2],
      "number": .[3],
      "string": .[4],
      "objectKey": .[5],
      "array": .[6],
      "object": .[7],
    };
  def defaultColours:
    if ($ENV["JQ_FLAVOUR"] | tostring | startswith("gojq")) then
      "90:33:33:36:32:34;1" | parseGOJQColours
    else 
      "0;90:0;39:0;39:0;39:0;32:1;39:1;39:1;34" | parseJQColours
    end;

  try (
    if $ENV["JQ_COLORS"] != null then 
      $ENV["JQ_COLORS"] | parseJQColours
    elif $ENV["GOJQ_COLORS"] != null then
      $ENV["GOJQ_COLORS"] | parseGOJQColours
    else
      defaultColours
    end
  ) catch (
    defaultColours
  );

def encodeJSONValue($withColour; $escapes; $indent; $currentDepth):
  def colourText($kind; $escapes):
    . as $text |
    ansiLib::CSR + $escapes[$kind] + "m" + $text + ansiLib::CSR + "0m";

  def maybeColourText($withColour; $kind; $escapes): if $withColour then colourText($kind; $escapes) end;

  def maybeIndent($indent; $currentDepth):
    if $indent > 0 then [range($indent * $currentDepth) | " "] | join("") else "" end;
  def maybeNewline($indent):
    if $indent > 0 then "\n" else "" end;
  def maybeSpace($indent):
    if $indent > 0 then " " else "" end;

  . as $value |
  (. | type) as $valueType |
  if $valueType == "null" then
    "null" | maybeColourText($withColour; "null"; $escapes)
  elif $valueType == "boolean" then
    if $value then
      "true" | maybeColourText($withColour; "true"; $escapes)
    else 
      "false" | maybeColourText($withColour; "false"; $escapes)
    end
  elif $valueType == "number" then
    $value | tostring | maybeColourText($withColour; "number"; $escapes)
  elif $valueType == "string" then
    $value | tojson | maybeColourText($withColour; "string"; $escapes)
  elif $valueType == "array" then
    ($value | length) as $numValues |
    ("[" | maybeColourText($withColour; "array"; $escapes)) +
    if $numValues > 0 then maybeNewline($indent) else "" end +
    ([
      $value[] | 
      maybeIndent($indent; $currentDepth) + 
      encodeJSONValue($withColour; $escapes; $indent; $currentDepth + 1)
    ] | join("," + maybeNewline($indent))) +
    if $numValues > 0 then maybeNewline($indent) else "" end +
    if $numValues > 0 then maybeIndent($indent; $currentDepth-1) else "" end +
    ("]" | maybeColourText($withColour; "array"; $escapes))
  elif $valueType == "object" then
    ($value | keys | length) as $numKeys |
    ("{" | maybeColourText($withColour; "object"; $escapes)) + 
    if $numKeys > 0 then maybeNewline($indent) else "" end +
    ([
      $value | keys | sort[] |
      . as $objectKey |
      $value[$objectKey] as $objectValue |
      maybeIndent($indent; $currentDepth) + 
      ($objectKey | tojson | maybeColourText($withColour; "objectKey"; $escapes)) +
      ":" + maybeSpace($indent) +
      ($objectValue | encodeJSONValue($withColour; $escapes; $indent; $currentDepth + 1))
    ] | join("," + maybeNewline($indent))) +
    if $numKeys > 0 then maybeNewline($indent) else "" end +
    if $numKeys > 0 then maybeIndent($indent; $currentDepth-1) else "" end +
    ("}" | maybeColourText($withColour; "object"; $escapes))
  end;

def encodeJSON($withColour; $indent):
  . | encodeJSONValue($withColour; colorEscapes; $indent // 0; 1);