require 'json'

# takes a JSON object, removes specified subtrees and trims arrays
# returns a pretty printed string
#
#     before = <<-EOF
#     { "foo": {
#         "one": { "a": 1, "b": {"c": 2} },
#         "two": "value",
#         "three": { "b": [1,2,3] }},
#       "bar": [1,2,3,4,5] }
#     EOF
#
#     blacklist = [
#       "foo:!two",   # delete the value before['foo']['two']
#       "foo:*:!b",   # delete all the keys called "b" nested two levels under "foo"
#       "bar:+"       # only keep the first element of the foo-list
#     ]
#
#     after = JSONPrune.cut(before, :prune => blacklist)
#
#      {
#        "foo": {
#          "two": ...,
#          "three": {
#            "b": [ ... ]
#          },
#          "one": {
#            "a": 1,
#            "b": { ... }
#          }
#        },
#        "bar": [
#          1,
#          ...
#        ]
#      }
#
class JSONPrune

  # prune the given JSON-string or ruby Hash according to
  # the rules given in opts[:prune]
  #
  def self.cut(obj, opts = {})
    obj = JSON.parse(obj) if String === obj

    (opts[:prune] || []).each do |ignore_str|
      ignoring = ignore_str.split(":")

      obj = cutr(obj, ignoring)
    end

    js = JSON.pretty_generate(obj)
    js = js.gsub('"IGN_HSH"', "{ ... }")
    js = js.gsub('"IGN_ARY"', "[ ... ]")
    js = js.gsub('"IGN"', "...")
  end
  
  # recursively prune the current ruby object with the given rules
  # replace pruned elements with marker strings
  #
  #   cutr( [1,2,3],         ["+"] )       #=> [ 1, "IGN" ]
  #   cutr( {"a"=>{"b"=>3}}, ["!a"] )      #=> { "a" => "IGN_HSH" }
  #   cutr( {"a"=>{"b"=>3}}, ["a", "!b"])  #=> { "a" => { "b" => "IGN" }}
  #
  def self.cutr(data, rules)
    key, *rest = rules

    if key =~ /^!(.*)/
      name = $1

      data[name] = case data[name]
                     when Hash then "IGN_HSH"
                     when Array then "IGN_ARY"
                     else "IGN"
                   end if data[name]
    elsif key =~ /^\+$/
      if Array === data
        data = [data.first, "IGN"]
      end
    elsif key =~ /^\*$/
      if Array === data
        data = data.map { |elem| cutr(elem, rest) }
      elsif Hash === data
        data = data.inject({}) do |h,(k,v)| h[k] = cutr(v, rest); h end
      end
    else
      name = key

      if data[name] and not rest.empty?
        data[name] = cutr(data[name], rest)
      end
    end

    data
  end
end

if __FILE__ == $0

  require 'shoulda'

  class FooTest < Test::Unit::TestCase
    def assert_equal_ws(expected, actual)
      expected = expected.gsub(/\s/, "")
      actual = actual.gsub(/\s/,"")
      assert_equal expected, actual
    end

    should "pass json through unchanged" do
      assert_equal %Q({\n  "foo": "bar"\n}),
                   JSONPrune.cut(%q({"foo": "bar"}))
    end

    should "ignore" do
      assert_equal({"a" => 3, "b" => "IGN" },
                   JSONPrune.cutr({ "a" => 3, "b" => 4 } , ["!b"]))
      assert_equal({"a" => 3, "b" => "IGN_HSH" },
                   JSONPrune.cutr({ "a" => 3, "b" => { "c" => 4 }} , ["!b"]))
      assert_equal({"a" => 3, "b" => "IGN_ARY" },
                   JSONPrune.cutr({ "a" => 3, "b" => [ 4 ] } , ["!b"]))
    end

    should "distinguish between arrays and hashes" do
      assert_equal %Q({\n  "foo": { ... }\n}),
                   JSONPrune.cut(%q({"foo" : { "bar": "fred"}}),  :prune => ["!foo"])
      assert_equal %Q({\n  "foo": [ ... ]\n}),
                   JSONPrune.cut(%q({"foo" : [ "bar", "fred"]}),  :prune => ["!foo"])
      assert_equal %Q({\n  "foo": ...\n}),
                   JSONPrune.cut(%q({"foo" : 3}),  :prune => ["!foo"])
    end

    should "work for nested attributes" do
      assert_equal %Q({\n  "foo": {\n    "bar": { ... }\n  }\n}),
                   JSONPrune.cut(%q({"foo" : { "bar": {"f":3}}}), :prune => ["foo:!bar"])

      assert_equal %Q({\n  "foo": {\n    "bar": ...\n  }\n}),
                  JSONPrune.cut(%q({"foo" : { "bar": 3}}), :prune => ["foo:!bar"])
    end

    should "not alter json if ignored thing does not exist" do
      assert_equal %Q({\n  "foo": { ... }\n}),
                   JSONPrune.cut(%q({"foo" : { "bar": "fred"}}),
                     :prune => ["!foo", "!doesnotexist", "foo:!doesnotexist"])
    end

    should "allow ignoring of subelements of arrays" do
      assert_equal_ws %Q({"foo":[{"bar": { ... }},{"bar": ...}]}),
                   JSONPrune.cut(%q({"foo" : [{ "bar": {"baz":3}}, {"bar":3}]}), :prune => ["foo:*:!bar"])
    end

    should "allow ignoring of arrays" do
      assert_equal %Q({\n  "foo": [\n    "bar",\n    ...\n  ]\n}),
                   JSONPrune.cut(%q({"foo": ["bar", "baz", "fred"]}),
                     :prune => ["foo:+"])
    end

    should "pass a complex testcase" do
      before = <<-EOF
       { "foo": {
            "one": { "a": 1, "b": {"c": 2} },
            "two": "value",
            "three": { "b": [1,2,3] }},
         "bar": [1,2,3,4,5] }
       EOF

       blacklist = [
         "foo:!two",
         "foo:*:!b",
         "bar:+"
       ]

       expected = <<-EOF
{
  "foo": {
    "two": ...,
    "three": {
      "b": [ ... ]
    },
    "one": {
      "a": 1,
      "b": { ... }
    }
  },
  "bar": [
    1,
    ...
  ]
}
EOF

       actual = JSONPrune.cut(before, :prune => blacklist)
       assert_equal expected.chomp, actual
    end

  end
end