#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#   Copyright 2012 Bruce Yang
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#

"""
parse a file with a headline and json content like this:

GET /index/type/_search?pretty
{
    "query": {
        "match": {
            "title": "hello world"
        }
    }
}

to make a elasticsearch restful query.
multiple queries can in one query file, splited by blank line.
"""

import sys
import json
import requests
import click

def print_json(content, pretty=False):
    if not pretty:
        print content
    else:
        data = json.loads(content)
        content_pretty = json.dumps(data, indent=2)
        print content_pretty

def parse_and_query(query_file, pretty=False):
    while True:
        headline = query_file.readline()
        if not headline:
            return False
        if headline.startswith('#'):
            continue
        break
    try:
        method, url = headline.split()
        method = method.upper()
    except:
        print 'Bad headline, exit!'
        return False
    if method not in ('GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS'):
        print 'Wrong method: %s' % method
        return False
    if url.startswith('/'):
        url = '%s%s' % (globals().get('_host', 'http://localhost:9200'), url)
    print '-' * 80
    print method, url
    lines = []
    while True:
        line = query_file.readline()
        if not line or not line.strip():
            break
        if headline.lstrip().startswith('#'):
            continue
        lines.append(line)
    data = ''.join(lines)
    resp = requests.request(method, url, data=data)
    if resp.ok:
        print 'OK', resp.status_code
        print_json(resp.content, pretty)
    else:
        print 'FAIL', resp.status_code
        print_json(resp.content, pretty)
    return True

def parse_dsl_file(file, pretty):
    while True:
        ret = parse_and_query(file, pretty=pretty)
        if not ret:
            break

@click.command()
@click.option('-p', '--pretty', is_flag=True, help='Pretty output json')
@click.option('-h', '--host', default='http://localhost:9200', help='Default host and port of elasticsearch server')
@click.argument('filenames', nargs=-1, required=False)
def main(pretty, host, filenames):
    global _host
    if host.find('://') == -1:
        host = 'http://%s' % host
    _host = host
    for filename in filenames:
        with open(filename) as f:
            parse_dsl_file(f, pretty)
    if not filenames:
        parse_dsl_file(sys.stdin, pretty)

if __name__ == '__main__':
    main()
