#!/usr/bin/env python

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#

import os
import getopt
import sys
import locale
import socket
import re
import httplib
import json
from qpid.disp import Display, Header, Sorter

_host = "localhost"
_connTimeout = 10
_types = ""
_limit = 200
_increasing = False
_sortcol = None
pattern = re.compile("^\\d+\\.\\d+\\.\\d+\\.\\d+:\\d+$")

def Usage ():
    print "Usage:  qpid-queue-count [OPTIONS] [broker-addr]"
    print
    print "             broker-addr is in the form:   [username/password@] hostname | ip-address [:<port>]"
    print "             ex:  localhost, 10.1.1.7:10000, broker-host:10000, guest/guest@localhost"
    print
    print "General Options:"
    print "    --timeout seconds (10)  Maximum time to wait for broker connection"
#    print "     -n [--numeric]  Don't resolve names"
    print
    print "Display Options:"
    print "     -S [--sort-by] COLNAME  Sort by column name"
    print "     -I [--increasing]       Sort by increasing value (default = decreasing)"
    print "     -L [--limit] NUM        Limit output to NUM rows (default = 200)"
    print
    sys.exit (1)

class RestManager():
    def __init__(self):
        self._host = None
        self._port = 0

    def setPort(self, port):
        self._port = port

    def setHost(self, host):
        self._host = host
    
    def execute(self, service):
        if (self._port is None):
            httpConn = httplib.HTTPConnection(self._host)
        else:
            httpConn = httplib.HTTPConnection(self._host, self._port)

        if (_connTimeout is not None):
            httpConn.timeout = _connTimeout

        httpConn.connect()
        httpConn.request("GET", "/api/latest/" + service)
        response = httpConn.getresponse()
        
        if (response.status != 200):
            print "Unable to post request to server!"
            print response.reason
            sys.exit(1)
    
        return response
##
## Main Program
##
try:
    longOpts = ("top", "numeric", "sort-by=", "limit=", "increasing", "timeout=")
    (optlist, encArgs) = getopt.gnu_getopt(sys.argv[1:], "bceqS:L:I", longOpts)
except:
    Usage()

try:
    encoding = locale.getpreferredencoding()
    cargs = [a.decode(encoding) for a in encArgs]
except:
    cargs = encArgs

for opt in optlist:
    if opt[0] == "--timeout":
        _connTimeout = int(opt[1])
        if _connTimeout == 0:
            _connTimeout = None
    elif opt[0] == "-n" or opt[0] == "--numeric":
        _numeric = True
    elif opt[0] == "-S" or opt[0] == "--sort-by":
        _sortcol = opt[1]
    elif opt[0] == "-I" or opt[0] == "--increasing":
        _increasing = True
    elif opt[0] == "-L" or opt[0] == "--limit":
        _limit = int(opt[1])
    elif len(opt[0]) == 2:
        char = opt[0][1]
        if "bcseq".find(char) != -1:
            _types += char
        else:
            Usage()
    else:
        Usage()

if len(_types) == 0:
    _types='q'

nargs = len(cargs)
rm    = RestManager()
_user = None
_password = None
_port = 8180

if nargs == 1:
    _broker_addr = cargs[0]
    _host = _broker_addr
    
    # check for username
    if _host.find("@") != -1:
        tokens = _host.split("@")
        _user = tokens[0]
        _host = tokens[1]
        # check for password
        if _user.find("/") != -1:
            tokens = _user.split("/")
            _user = tokens[0]
            _password = tokens[1]
    # check for port
    if _host.find(":") != -1:
        tokens = _host.split(":")
        _host = tokens[0]
        _port = int(tokens[1])

try:
    rm.setHost(_host)
    rm.setPort(_port)
    response = rm.execute('queue')
    
    # evaluate the JSON
    jsonStr = response.read()
    jsonObjArray = json.loads(jsonStr)
   
    # table header
    heads = []
    heads.append(Header("queue"))
    heads.append(Header("msg", Header.KMG))
    heads.append(Header("bytes", Header.KMG))
    heads.append(Header("cons", Header.KMG))
 
    # create rows of queues to display
    rows = []
    for staticDict in jsonObjArray:
        row = []
        row.append(staticDict.get("name"))
        statistics = staticDict.get("statistics")
        row.append(statistics.get("queueDepthMessages"))
        row.append(statistics.get("queueDepthBytes"))
        row.append(statistics.get("consumerCount"))
        rows.append(row)

    # sort the queues if required        
    if _sortcol:
        sorter = Sorter(heads, rows, _sortcol, _limit, _increasing)
        dispRows = sorter.getSorted()
    else:
        dispRows = rows
    
    # prepare to build & display the table
    disp = Display(prefix="  ")
    disp.formattedTable("Queues", heads, dispRows)
except KeyboardInterrupt:
    print
except Exception,e:
    print "Failed: %s - %s" % (e.__class__.__name__, e)
    sys.exit(1)
