import "ansiLib" as ansiLib;

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

  def parseGOJQColours:
    split(":") |
    {
      "null": (.[0] // "0"),
      "false": (.[1] // "0"),
      "true": (.[2] // "0"),
      "number": (.[3] // "0"),
      "string": (.[4] // "0"),
      "objectKey": (.[5] // "0"),
      "array": (.[6] // "0"),
      "object": (.[7] // "0"),
    };
  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 _encodeJSONValuePlainPretty($indent; $currentDepth):
  ([range($indent * $currentDepth) | " "] | join("")) as $currentIndentDepthIndent |
  ([range($indent * ($currentDepth + 1)) | " "] | join("")) as $nextDepthIndent |
  "\n" as $currentIndentDepthNewline |
  " " as $currentIndentDepthSpace |

  . as $value |
  (. | type) as $valueType |
  if $valueType == "array" then
    ($value | length) as $numValues |

    if $numValues > 0 then
      "[" +
      $currentIndentDepthNewline +
      ([
        $value[] | 
          $nextDepthIndent + 
          _encodeJSONValuePlainPretty($indent; $currentDepth + 1)
      ] | join("," + $currentIndentDepthNewline)) +
      $currentIndentDepthNewline +
      $currentIndentDepthIndent +
      "]"
    else
      "[]" 
    end
  elif $valueType == "object" then
    ($value | keys | length) as $numKeys |

    if $numKeys > 0 then
      "{" + 
      $currentIndentDepthNewline +
      ([
        $value | keys | sort[] |
          . as $objectKey |
          $value[$objectKey] as $objectValue |

          $nextDepthIndent + 
          ($objectKey | tojson) +
          ":" + 
          $currentIndentDepthSpace +
          ($objectValue | _encodeJSONValuePlainPretty($indent; $currentDepth + 1))
      ] | join("," + $currentIndentDepthNewline)) +
      $currentIndentDepthNewline +
      $currentIndentDepthIndent +
      "}"
    else 
      "{}"
    end
  else 
    tojson
  end;

def _encodeJSONValueColourCompact:
  colourEscapes as $escapes |

  def colourText($text; $kind):
    ansiLib::CSR + $escapes[$kind] + "m" + $text + ansiLib::CSR + "0m";

  . as $value |
  (. | type) as $valueType |
  if $valueType == "null" then
    colourText("null"; "null")
  elif $valueType == "boolean" then
    ($value | tostring) as $boolStr |
    colourText($boolStr; $boolStr)
  elif $valueType == "number" then
    colourText($value | tostring; "number")
  elif $valueType == "string" then
     colourText($value | tojson; "string")
  elif $valueType == "array" then
    ($value | length) as $numValues |

    if $numValues > 0 then
      colourText("["; "array") +
      ([
        $value[] | 
          _encodeJSONValueColourCompact
      ] | join(colourText(","; "array"))) +
      colourText("]"; "array")
    else
      colourText("[]"; "array")
    end
  elif $valueType == "object" then
    ($value | keys | length) as $numKeys |

    if $numKeys > 0 then
      colourText("{"; "object") + 
      ([
        $value | keys | sort[] |
          . as $objectKey |
          $value[$objectKey] as $objectValue |

          colourText($objectKey | tojson; "objectKey") +
          colourText(":"; "object") +
          ($objectValue | _encodeJSONValueColourCompact)
      ] | join(colourText(","; "object"))) +
      colourText("}"; "object")
    else 
      colourText("{}"; "object")
    end
  end;

def _encodeJSONValueColourPretty($indent; $currentDepth):
  colourEscapes as $escapes |

  def colourText($text; $kind):
    ansiLib::CSR + $escapes[$kind] + "m" + $text + ansiLib::CSR + "0m";

  ([range($indent * $currentDepth) | " "] | join("")) as $currentIndentDepthIndent |
  ([range($indent * ($currentDepth + 1)) | " "] | join(""))  as $nextDepthIndent |
  "\n" as $currentIndentDepthNewline |
  " " as $currentIndentDepthSpace |

  . as $value |
  (. | type) as $valueType |
  if $valueType == "null" then
    colourText("null"; "null")
  elif $valueType == "boolean" then
    ($value | tostring) as $boolStr |
    colourText($boolStr; $boolStr)
  elif $valueType == "number" then
    colourText($value | tostring; "number")
  elif $valueType == "string" then
     colourText($value | tojson; "string")
  elif $valueType == "array" then
    ($value | length) as $numValues |

    if $numValues > 0 then
      colourText("["; "array") +
      $currentIndentDepthNewline +
      ([
        $value[] | 
          $nextDepthIndent + 
          _encodeJSONValueColourPretty($indent; $currentDepth + 1)
      ] | join(colourText(","; "array") + $currentIndentDepthNewline)) +
      $currentIndentDepthNewline +
      $currentIndentDepthIndent +
      colourText("]"; "array")
    else
      colourText("[]"; "array")
    end
  elif $valueType == "object" then
    ($value | keys | length) as $numKeys |

    if $numKeys > 0 then
      colourText("{"; "object") + 
      $currentIndentDepthNewline +
      ([
        $value | keys | sort[] |
          . as $objectKey |
          $value[$objectKey] as $objectValue |

          $nextDepthIndent + 
          colourText($objectKey | tojson; "objectKey") +
          colourText(":"; "object") + 
          $currentIndentDepthSpace +
          ($objectValue | _encodeJSONValueColourPretty($indent; $currentDepth + 1))
      ] | join(colourText(","; "object") + $currentIndentDepthNewline)) +
      $currentIndentDepthNewline +
      $currentIndentDepthIndent +
      colourText("}"; "object")
    else 
      colourText("{}"; "object")
    end
  end;


def encodeJSON($withColour; $indent):
  if ($withColour == false) and ($indent == 0) then
    tojson
  else
    if $withColour then
      if $indent > 0 then 
        _encodeJSONValueColourPretty($indent; 0)
      else 
        _encodeJSONValueColourCompact
      end
    else 
      _encodeJSONValuePlainPretty($indent; 0)
    end
  end;