# COPYRIGHT:    (c) Awasu Pty. Ltd. 2002-13 (all rights reserved).
#               Unauthorized use of this code is prohibited.
#
# LICENSE:      This software is provided 'as-is', without any express
#               or implied warranty.
#
#               In no event will the author be held liable for any damages
#               arising from the use of this software.
#
#               Permission is granted to anyone to use this software
#               for any purpose and to alter it and redistribute it freely,
#               subject to the following restrictions:
#
#               - The origin of this software must not be misrepresented;
#                 you must not claim that you wrote the original software.
#                 If you use this software, an acknowledgment is requested
#                 but not required.
#
#               - Altered source versions must be plainly marked as such,
#                 and must not be misrepresented as being the original software.
#                 Altered source is encouraged to be submitted back to
#                 the original author so it can be shared with the community.
#                 Please share your changes.
#
#               - This notice may not be removed or altered from any
#                 source distribution.

""" 
This class provides access the Awasu API. 
    http://awasu.com/api
"""
 
import sys
import os
import urllib2
import BaseHTTPServer
from xml.etree import ElementTree
import json
import zlib
import StringIO
import re

DEFAULT_API_URL = "http://localhost:2604"

class AwasuApiException(Exception) : pass

# ---------------------------------------------------------------------

class AwasuApi :

    def __init__( self , url=None , token=None ) :
        self.mApiUrl = url if url is not None else DEFAULT_API_URL
        self.mApiToken = token

    def callApi( self , apiName , apiArgs=None , postData=None , rawMode=False , returnHttpHeaders=False ) :
        """ This is the main entry point for calling the Awasu API.
        Most of the time, you won't need to call this method directly, helper methods are provided for the most common operations.
        
        In raw mode, this function will return the raw response returned by Awasu (typically XML, but it can be something else if you pass in a "format=" parameter).

        Otherwise, it will parse the response and return an ElementTree tree (for XML) or a JSON dictionary. You can then do things like:
            # print the name of every channel
            xml = awasuApi.callApi( "channels/list" )
            for channelNode in xml.findall( "./channel" ) :
                print channelNode.find("name").text
            # dump the contents of the default workpad
            json = awasuApi.callApi( "workpads/get" , {"id":"@","format":"json"} )
            for workpadItem in json["workpad"]["workpadItems"] :
                print "%s => %s" % (workpadItem["title"] , workpadItem["url"])
        """
        # generate the request URL
        url = "%s/%s" % (self.mApiUrl , apiName)
        if url[:7] != "http://" :
            url = "http://" + url
        # initialize the API arguments 
        if apiArgs is None : apiArgs = {}
        if self.mApiToken != None and self.mApiToken != "" :
            apiArgs["token"] = self.mApiToken
        # add the API arguments to the POST data
        # NOTE: We do this to avoid exposing the token in GET requests.
        if len(apiArgs) > 0 :
            if postData is None :
                # no POST data was supplied - create a new data block
                postData = ElementTree.Element( "apiArgs" )
                apiArgsNode = postData
            else :
                # load the POST data
                postData = ElementTree.fromstring( postData ) 
                # NOTE: When parsing the POST data, Awasu stops after it has processed 
                # the <apiArgs> node, so it's advantageous to put it first (to avoid 
                # having to parse the entire XML tree).
                apiArgsNode = ElementTree.Element( "apiArgs" )
                postData.insert( 0 , apiArgsNode )
            # add the API arguments to the POST data (as <apiArgs> attributes)
            for key,val in apiArgs.items() :
                apiArgsNode.set( key , str(val) )
            # convert the POST data back to a string
            postData = ElementTree.tostring( postData )
        # send the request
        req = urllib2.Request( url , postData , {"Accept-Encoding":"deflate"} )
        resp = urllib2.urlopen( req )
        responseHeaders = str( resp.info() )
        responseBody = resp.read()
        resp.close() # nb: try to stop socket exhaustion when stress-testing
        # return the response
        responseHeadersDict = {} # FIXME! how to get the HTTP status code/message?
        for lineBuf in StringIO.StringIO(responseHeaders) :
            mo = re.match( "^(\\s*[^()<>@,;:\\\"/\\[\\]?={} ]+)\\s*:\\s*(.*)$" , lineBuf )
            if mo : responseHeadersDict[ mo.group(1) ] = mo.group(2).strip()
        if responseHeadersDict.get("Content-Encoding") == "deflate" :
            responseBody = zlib.decompressobj(-zlib.MAX_WBITS).decompress( responseBody )
        if not rawMode :
            if getResponseFormat(apiArgs) == "xml" :
                responseBody = ElementTree.fromstring(responseBody) if responseBody.strip() != "" else None
            elif getResponseFormat(apiArgs) == "json" :
                responseBody = json.loads(responseBody) if responseBody.strip() != "" else None
        return (responseHeadersDict,responseBody) if returnHttpHeaders else responseBody
    
    def getAwasuBuildInfo( self ) :
        """ Get the Awasu build info. """
        return self.callApiAndCheck( "buildInfo" , {"format":"json"} )[ "buildInfo" ]
        
    def getAwasuUserInfo( self ) :
        """ Get the Awasu user info. """
        return self.callApiAndCheck( "userInfo" , {"format":"json"} )[ "userInfo" ]
        
    def getAwasuStats( self ) :
        """ Get the Awasu stats. """
        return self.callApiAndCheck( "stats" , {"format":"json"} )[ "stats" ]
        
    def getAwasuActivityLog( self , nLines=None ) :
        """ Get the Awasu Activity log. """
        return self.callApiAndCheck( "logs/activity" , {"lines":nLines} , None , True )
        
    def getAwasuErrorLog( self , nLines=None ) :
        """ Get the Awasu Error log. """
        return self.callApiAndCheck( "logs/error" , {"lines":nLines} , None , True )
        
    def getChannelFolders( self , asTree=True ) :
        """ Get the channel folders. """
        if asTree :
            return self.callApiAndCheck( "channels/folders/tree" , {"format":"json"} )[ "channelFolder" ]
        else :
            return self.callApiAndCheck( "channels/folders/list" , {"format":"json"} )[ "channelFolders" ]

    def createChannelFolder( self , folderName , parentFolder=None  , insertAfterFolder=None ) :
        """ Create a new channel folder. """
        apiArgs = { "name": folderName , "format": "json" }
        if parentFolder : apiArgs["parent"] = parentFolder 
        if insertAfterFolder : apiArgs["after"] = insertAfterFolder
        return self.callApiAndCheck( "channels/folders/create" , apiArgs )["status"]["id"]

    def deleteChannelFolder( self , folderId ) :
        """ Delete a channel folder. """
        json = self.callApiAndCheck( "channels/folders/delete" , {"id":folderId,"format":"json"} )
        
    def getChannelFilters( self ) :
        """ Get the channel filters. """
        json = self.callApiAndCheck( "channels/filters/list" , {"format":"json"} )
        return json["channelFilters"]

    def getChannels( self , channelIds=None , verbose=False ) :
        """ Get the configuration details for the specified channels. """
        apiArgs = { "format": "json" , "verbose": verbose }
        apiArgs = addIdsToApiArgs( apiArgs , channelIds )
        return self.callApiAndCheck( "channels/list" , apiArgs )[ "channels" ]
        
    def getChannelStats( self , channelIds=None ) :
        """ Get the statistics for the specified channels. """
        apiArgs = addIdsToApiArgs( {"format":"json"} , channelIds )
        return self.callApiAndCheck( "channels/stats" , apiArgs )[ "channels" ]
    
    def getChannelErrors( self , channelIds=None ) :
        """ Get the error log for the specified channels. """
        apiArgs = addIdsToApiArgs( {"format":"json"} , channelIds )
        return self.callApiAndCheck( "channels/errors" , apiArgs )[ "channels" ]
    
    def getChannelSummary( self , channelId ) :
        """ Get the summary for the specified channel. """
        if isinstance( channelId , list ) : raise AwasuApiException( "Can't get multiple channels." )        
        return self.callApiAndCheck( "channels/get" , {"id":channelId,"format":"html"} )
        
    def createChannel( self , postData ) :
        """ Create a new channel. """
        json = self.callApiAndCheck( "channels/create" , {"format":"json"} , postData )
        return int( json["status"]["id"] )
    def createChannelByUrl( self , url ) :
        """ Create a new channel (downloaded from the specified URL). """
        return self.createChannel( "<channel type='standard'><feedUrl>%s</feedUrl></channel>" % safeXmlString(url) )
    def createPluginChannel( self , pluginPath , pluginParams ) :
        """ Create a new plugin channel. """
        paramsXml = "".join( [ "<param name='%s'>%s</param>" % (safeXmlString(key),safeXmlString(val)) for key,val in pluginParams.items() ] )
        return self.createChannel( "<channel type='plugin'><pluginChannel path='%s'>%s</pluginChannel></channel>" % (safeXmlString(pluginPath),paramsXml) )
    def createSearchChannel( self , queryString , searchLocations=None , advancedSyntax=False ) :
        """ Create a new search channel. """
        attrs = { "advancedSyntax": boolString(advancedSyntax) }
        if searchLocations :
            attrs["searchInUrls"] = boolString(searchLocations)
            attrs["searchInTitles"] = boolString(searchLocations)
            attrs["searchInDescriptions"] = boolString(searchLocations)
        attrs = " ".join( [ "%s='%s'" % (key,val) for key,val in attrs.items() ] )
        return self.createChannel( "<channel type='search'><searchQuery %s>%s</searchQuery></channel>" % (attrs,safeXmlString(queryString)) )
            
    def deleteChannels( self , channelIds ) :
        """ Delete the specified channels. """
        apiArgs = addIdsToApiArgs( {"format":"json"} , channelIds )
        json = self.callApiAndCheck( "channels/delete" , apiArgs )
        for channelStatus in json["channels"] :
            if channelStatus["status"] != "OK" : 
                raise AwasuApiException( "Can't delete channel \"%s\" (%d): %s" % (channelStatus["name"],channelStatus["id"],channelStatus["status"]) )
              
    def getReports( self , reportIds=None , verbose=False ) :
        """ Get the configuration details for the specified reports. """
        apiArgs = { "format": "json" , "verbose": verbose }
        apiArgs = addIdsToApiArgs( apiArgs , reportIds )
        return self.callApiAndCheck( "reports/list" , apiArgs )[ "channelReports" ]

    def runReports( self , reportIds ) :
        """ Run the specified reports. """
        apiArgs = addIdsToApiArgs( {"format":"json"} , reportIds )
        json = self.callApiAndCheck( "reports/run" , apiArgs )
        for reportStatus in json["channelReports"] :
            if reportStatus["status"] != "OK" : 
                raise AwasuApiException( "Can't run report \"%s\" (%s): %s" % (reportStatus["name"],reportStatus["id"],reportStatus["status"]) )
        
    def getReport( self , reportId ) :
        """ Run the specified report and return the result. """
        if isinstance( reportId , list ) : raise AwasuApiException( "Can't get multiple reports." )        
        return self.callApiAndCheck( "reports/get" , {"id":reportId,"format":"html"} )
        
    def createReport( self , postData ) :
        """ Create a new report. """
        json = self.callApiAndCheck( "reports/create" , {"format":"json"} , postData )
        return json["status"]["id"]
    def createChannelFilterReport( self , reportName , channelFilterName , reportDescription=None ) :
        """ Create a new report, based on the specified channel filter. """
        return self.createReport( "<channelReport><name>%s</name><description>%s</description><dataSource type='channelFilter'><channelFilterName>%s</channelFilterName></dataSource></channelReport>" % (safeXmlString(reportName),safeXmlString(reportDescription),safeXmlString(channelFilterName)) )
    def createChannelFoldersReport( self , reportName , channelFolderIds , includeSubFolders , reportDescription=None ) :
        """ Create a new report, based on the specified channel folders. """
        channelFoldersXml = "".join( [ "<channelFolder id='%s'/>" % safeXmlString(id) for id in channelFolderIds ] ) if channelFolderIds else ""
        return self.createReport( "<channelReport><name>%s</name><description>%s</description><dataSource type='channelFolders' includeSubFolders='%s'>%s</dataSource></channelReport>" % (safeXmlString(reportName),safeXmlString(reportDescription),"true" if includeSubFolders else "false",channelFoldersXml) )
    def createWorkpadReport( self , reportName , workpadId , reportDescription=None ) :
        """ Create a new report, based on the specified workpad. """
        return self.createReport( "<channelReport><name>%s</name><description>%s</description><dataSource type='workpad'><workpad id='%s'/></dataSource></channelReport>" % (safeXmlString(reportName),safeXmlString(reportDescription),safeXmlString(workpadId)) )

    def deleteReports( self , reportIds ) :
        """ Delete the specified reports. """
        apiArgs = addIdsToApiArgs( {"format":"json"} , reportIds )
        json = self.callApiAndCheck( "reports/delete" , apiArgs )
        for reportStatus in json["channelReports"] :
            if reportStatus["status"] != "OK" : 
                raise AwasuApiException( "Can't delete report \"%s\" (%s): %s" % (reportStatus["name"],reportStatus["id"],reportStatus["status"]) )

    def getWorkpads( self , workpadIds=None ) :
        """ Get the configuration details for the specified workpads. """
        apiArgs = addIdsToApiArgs( {"format":"json"} , workpadIds )
        return self.callApiAndCheck( "workpads/list" , apiArgs )[ "workpads" ]
                
    def getWorkpad( self , workpadId ) :
        """ Get the contents of the specified workpad. """
        if isinstance( workpadId , list ) : raise AwasuApiException( "Can't get multiple workpads." )        
        return self.callApiAndCheck( "workpads/get" , {"id":workpadId,"format":"json"} )[ "workpad" ]
                
    def getDefaultWorkpad( self ) :
        """ Get the default workpad. """
        try :
            return self.getWorkpad( "@" )
        except AwasuApiException :
            return None # nb: we assume the error happened because there are no workpads
                
    def getWorkpadFeed( self , workpadId ) :
        """ Get the feed XML for the specified workpad. """
        xml = self.callApiAndCheck( "workpads/feed" , {"id":workpadId} )
        errorMsgNodes = xml.findall( "./errorMsg" )
        if len(errorMsgNodes) > 0 : raise AwasuApiException(errorMsgNodes[0].text)
        return ElementTree.tostring( xml )
        
    def addWorkpadItem( self , workpadIds , url , title=None , customFields=None ) :
        """ Add a new item to the specified workpads. """
        apiArgs = { "url":url , "title":title , "format":"json" }
        if customFields is not None : apiArgs.update( customFields )
        apiArgs = addIdsToApiArgs( apiArgs , workpadIds )
        json = self.callApiAndCheck( "workpads/addItem" , apiArgs ) 
        for workpadStatus in json["workpads"] :
            if workpadStatus["status"] != "OK" :
                raise AwasuApiException( "Can't add item to workpad \"%s\" (%s) : %s" % (workpadStatus["name"],workpadStatus["id"],workpadStatus["status"]) )
                
    def createWorkpad( self , workpadName , workpadDescription=None ) :
        """ Create a new workpad. """
        postData = "<workpad><name>%s</name><description>%s</description></workpad>" % (safeXmlString(workpadName),safeXmlString(workpadDescription))
        json = self.callApiAndCheck( "workpads/create" , {"format":"json"} , postData )
        return json["status"]["id"]

    def deleteWorkpads( self , workpadIds ) :
        """ Delete the specified workpads. """
        apiArgs = addIdsToApiArgs( {"format":"json"} , workpadIds )
        json = self.callApiAndCheck( "workpads/delete" , apiArgs )
        for workpadStatus in json["workpads"] :
            if workpadStatus["status"] != "OK" : 
                raise AwasuApiException( "Can't delete workpad \"%s\" (%s): %s" % (workpadStatus["name"],workpadStatus["id"],workpadStatus["status"]) )
        
    def runSearchQuery( self , queryString , searchLocations=None , resultsFormat="excerpt" , advancedSyntax=False , pageNo=1 , pageSize=10 ) :
        """ Run the specified search query. """
        apiArgs = { "query": queryString , "fidf": resultsFormat , "advsyn": advancedSyntax , "page": pageNo , "pageSize": pageSize , "format": "json" }
        if searchLocations :
            apiArgs["locations"] = ",".join(searchLocations) if isinstance(searchLocations,list) else searchLocations
        return self.callApiAndCheck( "search/query" , apiArgs )[ "searchResults" ]
        
    def callApiAndCheck( self , apiName , apiArgs=None , postData=None , rawMode=False ) :
        """ Call the Awasu API and check for errors. """
        if apiArgs is None : apiArgs={}
        apiArgs["quiet"] = False
        response = self.callApi( apiName , apiArgs , postData , rawMode , True )
        # FIXME! Since we can't get the HTTP status code, we can't check it :-/
        if not rawMode :
            if getResponseFormat(apiArgs) == "json" :
                if response[1] is not None and "status" in response[1] and "errorMsg" in response[1]["status"] :
                    raise AwasuApiException( response[1]["status"]["errorMsg"] )
            elif getResponseFormat(apiArgs) == "xml" :
                if response[1] is not None and response[1].find("./errorMsg") is not None :
                    raise AwasuApiException( response[1].find("./errorMsg").text )
            elif getResponseFormat(apiArgs) == "html" :
                mo = re.search( "<td class=\"error-msg value\">(.+?)</td>" , response[1] )
                if mo : raise AwasuApiException( mo.group(1).strip() )
        return response[1]

    def __str__( self ) : return "AwasuApi @ %s" % self.mApiUrl
    
# ---------------------------------------------------------------------

def addIdsToApiArgs( apiArgs , idList ) :
    """ Add the specified ID's to the argument list. """
    if idList is not None :
        if isinstance( idList , list ) :
            apiArgs["id"] = ",".join( [ str(x) for x in idList ] )
        else :
            apiArgs["id"] = idList 
    return apiArgs 

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

def convertApiArgs( apiArgs ) :
    """ Convert a list of "key=val" arguments to a dictionary. """
    apiArgsDict = {}
    for arg in apiArgs :
        pos = arg.find( "=" )
        if pos > 0 :
            apiArgsDict[ arg[:pos] ] = arg[pos+1:]
    return apiArgsDict
        
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

def getResponseFormat( apiArgs ) :
    """ Determine what the response format will be. """
    if "format" in apiArgs : return apiArgs["format"]
    if "f" in apiArgs : return apiArgs["f"]
    return "xml"

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

def safeXmlString( val ) : 
    """ Convert the string into something that's safe for inclusion in XML. """
    return str(val).replace( "&" , "&amp;" ) \
                   .replace( "<" , "&lt;" ) \
                   .replace( ">" , "&gt;" ) \
                   .replace( "\"" , "&quot;" ) if val else ""

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

def boolString( val ) : return "true" if val else "false"

# ---------------------------------------------------------------------

def dumpXmlTree( xmlTree , prefix="" ) : __doDumpXmlTree( xmlTree , prefix , 0 )

def __doDumpXmlTree( xmlNode , prefix , tagFieldWidth ) :
    if xmlNode is None : return 
    # dump the XML node
    print prefix + (str(xmlNode.tag)+":").ljust(1+tagFieldWidth) , \
          xmlNode.text.strip().encode("utf-8") if xmlNode.text else ""
    # dump any attributes 
    for attr in xmlNode.items() :
        print prefix , " @"+attr[0] , "=" , attr[1].encode("utf-8")
    # dump any childNodes
    childNodes = list( xmlNode )
    if len(childNodes) > 0 :
        maxTagLen = max( [ len(cn.tag) for cn in childNodes ] )
        for cn in childNodes :
            __doDumpXmlTree( cn , prefix+"    " , maxTagLen )
            
# ---------------------------------------------------------------------

def printHelp() :
    scriptName = os.path.split(sys.argv[0])[ 1 ]
    print "%s {options} {api-name} {arg1} {arg2} ..." % scriptName
    print "  Calls the Awasu API."
    print
    print "Options:"
    print "  -u --url       Invocation URL (default=%s)" % DEFAULT_API_URL
    print "  -t --token     Awasu API token."
    print "  -h --headers   Output the HTTP response headers."
    print "  -r --raw       Output the raw response."
    print 
    print "The arguments following {api-name} are passed on to Awasu via the API call and are specified as they would normally be in a URL (i.e. \"key=val\" pairs)."
    print "For API calls that expect POST data, pipe the data into stdin."
    print
    print "Examples:"
    print "  Get a full list of channels and their configuration:"
    print "    %s channels/list verbose=1" % scriptName
    print 
    print "  Get the summary page for a channel:"
    print "    %s --raw channels/get name=... sfim=all" % scriptName
    print 
    print "  Update the configuration for a report:"
    print "    %s reports/update id=... <newReportConfig.xml" % scriptName
    print 
    print "  Add an item to the default workpad:"
    print "    %s workpads/addItem id=@ url=http://awasu.com title=Awasu" % scriptName
    print
    
# ---------------------------------------------------------------------

if __name__=="__main__" :
    # initialize
    sys.tracebacklimit = 0
    # parse the command-line arguments
    url = None
    token = None
    dumpHeaders = False
    rawMode = False
    import getopt 
    try :
        opts , args = getopt.getopt( sys.argv[1:] , "u:t:hr?" , ["url=","token=","headers","raw","help"] )
    except getopt.GetoptError , err :
        raise AwasuApiException , "Can't parse arguments: %s" % err
    for opt,val in opts :
        if opt in ["-u","--url"] :
            url = val
        elif opt in ["-t","--token"] :
            token = val
        elif opt in ["-h","--headers"] :
            dumpHeaders = True
        elif opt in ["-r","--raw"] :
            rawMode = True
        elif opt in ["-?","--help"] :
            printHelp()
            sys.exit()
        else :
            raise AwasuApiException , "Invalid command line option: %s" % opt
    if len(args) == 0 :
        printHelp()
        sys.exit()
    # call the Awasu API
    awasuApi = AwasuApi( url , token )
    postData = None if os.isatty(0) else sys.stdin.read()
    try :
        apiArgs = convertApiArgs( args[1:] )
        responseHeaders,responseBody = awasuApi.callApi( args[0] , apiArgs , postData , rawMode , True )
    except urllib2.HTTPError , xcptn :
        print "HTTP %d: %s" % (xcptn.code,BaseHTTPServer.BaseHTTPRequestHandler.responses[xcptn.code][0])
        print xcptn.read()
        responseHeaders = {}
        responseBody = None
    # output the results 
    if dumpHeaders :
        print "Response headers:"
        if len(responseHeaders) > 0 :
            maxKeyLen = max( [ len(k) for k in responseHeaders.keys() ] )
            for key,val in responseHeaders.items() :
                print "  %s %s" % ((str(key)+":").ljust(1+maxKeyLen) , val )
        print 
    if isinstance( responseBody , ElementTree.Element ) :
        print ElementTree.tostring( responseBody )
    elif isinstance( responseBody , dict ) :
        print json.dumps( responseBody )
    elif responseBody is not None :
        print responseBody 
