#!/usr/bin/env ruby

require 'jgrep'
require 'optparse'

@options = {:flat => false, :start => nil, :field => [], :slice => nil}

def print_json(result)
    unless @options[:flat]
        result = result.first if @options[:stream]
        puts(JSON.pretty_generate(result))
    else
        puts(result.first.to_json)
    end
end

def do_grep(json, expression)
    if @options[:field].empty?
        result = JGrep::jgrep((json), expression, nil, @options[:start])
        result = result.slice(@options[:slice]) if @options[:slice]
        exit 1 if result == []
        unless @options[:quiet] == true
            print_json(result)
        end
    else
        if @options[:field].size > 1
            JGrep::validate_filters(@options[:field])
            result = JGrep::jgrep((json), expression, @options[:field], @options[:start])
            result = result.slice(@options[:slice]) if @options[:slice]
            exit 1 if result == []
            unless @options[:quiet] == true
                print_json(result)
            end

        else
            JGrep::validate_filters(@options[:field][0])
            result = JGrep::jgrep((json), expression, @options[:field][0], @options[:start])
            result = result.slice(@options[:slice]) if @options[:slice]
            exit 1 if result == []
            if result.is_a?(Array) && !(result.first.is_a?(Hash) || result.flatten.first.is_a?(Hash))
                unless @options[:quiet] == true
                    result.map{|x| puts x unless x.nil?}
                end
            else
                unless @options[:quiet] == true
                    print_json(result)
                end
            end
        end
    end
end

begin
    OptionParser.new do |opts|
        opts.banner = "Usage: jgrep [options] \"expression\""
        opts.on("-s", "--simple [FIELDS]", "Display only one or more fields from each of the resulting json documents") do |field|
            unless field.nil?
                @options[:field].concat(field.split(" "))
            else
                raise "-s flag requires a field value"
            end
        end

        opts.on("-c", "--compact", "Display non pretty json") do
            @options[:flat] = true
        end

        opts.on("-n", "--stream", "Display continuous output from continuous input") do
            @options[:stream] = true
        end

        opts.on("-f", "--flatten", "Makes output as flat as possible") do
            JGrep::flatten_on
        end

        opts.on("-i", "--input [FILENAME]", "Specify input file to parse") do |filename|
            @options[:file] = filename
        end

        opts.on("-q", "--quiet", "Quiet; don't write to stdout.  Exit with zero status if match found.") do
            @options[:quiet] = true
        end

        opts.on("-v", "--verbose", "Verbose output") do
            JGrep::verbose_on
        end

        opts.on("--start [FIELD]", "Where in the data to start from") do |field|
            @options[:start] = field
        end

        opts.on("--slice [RANGE]", "A range of the form 'n' or 'n..m', indicating which documents to extract from the final output") do |field|
            range_nums = field.split('..').map{ |x| x.to_i }
            @options[:slice] = range_nums.length == 1 ? range_nums[0] : Range.new(*range_nums)
        end
    end.parse!
rescue OptionParser::InvalidOption => e
    puts e.to_s.capitalize
    exit 1

rescue Exception => e
    puts e
    exit 1
end

begin
    expression = nil

    #Identify the expression from command line arguments
    ARGV.each do |argument|
        if argument =~ /<|>|=|\+|-/
            expression = argument
            ARGV.delete(argument)
        end
    end

    expression = "" if expression.nil?

    #Continuously gets if inputstream in constant
    #Load json from standard input if tty is false
    #else find and load file from command line arugments

    if @options[:stream]
        unless STDIN.tty?
            while json = gets
                do_grep(json, expression)
            end
        else
            raise "No json input specified"
        end
    else
        if @options[:file]
            json = File.read(@options[:file])
            do_grep(json, expression)
        elsif ! STDIN.tty?
            json = STDIN.read
            do_grep(json, expression)
        else
            raise "No json input specified"
        end
    end

rescue Interrupt
    STDERR.puts "Exiting..."
    exit 1
rescue Exception => e
    if e.is_a?(SystemExit)
        exit e.status
    else
        STDERR.puts "Error - #{e}"
        exit 1
    end
end
