#!/usr/bin/env python
# vim: tw=0 nofen

import os
import sys
import xmlrpclib

def hrspeed(speed):
    speed = int(speed)
    ranges = [
        (1024**2, '%.2f MiB/s'),
        (1024, '%.2f KiB/s'),
        (1, '%d B/s'),
    ]
    for v in ranges:
        if speed/v[0]:
            return v[1] % (speed/float(v[0]))
    return '0'

def hrsize(size):
    size = int(size)
    ranges = [
        (1024**3, '%.2f GiB'),
        (1024**2, '%.2f MiB'),
        (1024, '%.2f KiB'),
        (1, '%d B')
    ]
    for v in ranges:
        if size/v[0]:
            return v[1] % (size/float(v[0]))
    return '0'


class AriaWrapper(object):
    FMTSTRING = (
        "%(gid)s\t%(status)s\t%(percent_done)s\t%(size_total)s\t"
        "CN:%(connections)s/%(numSeeders)s\t"
        "d:%(downloadSpeed)s\tu:%(uploadSpeed)s\t%(path)s"
    )
    def __init__(self, **kwargs):
        host = kwargs.get('host', '127.0.0.1')
        port = kwargs.get('port', 6800)
        user = kwargs.get('user', '')
        password = kwargs.get('password', '')
        auth = ''
        if all([user, password]):
            auth = '%s:%s' % (user, password)
        URI = 'http://%(auth)s@%(host)s:%(port)d/rpc' % locals()
        self._server = xmlrpclib.ServerProxy(URI)

    def listall(self):
        return self._server.aria2.tellActive()

# {'status': 'active', 'numPieces': '2037', 'downloadSpeed': '55035',
# 'pieceLength': '4194304', 'bitfield':
# 'bbb7b75f7f3ff775ef767ee247f0113d58752cc789f62f0210b44e107238c333e06800200400215018c01e92000b0a2f7640815210e402c88000441924240847a400438810000ea000247262310000000101460203ca18002340c229b4910b3c66108000200150010210e48182040080101051244700201a0a4025125220823c20020004480050200410c128000202c2000005209fc0002010050000810080e0000004024022008001e70414104408000000c4e03a30620001090000081389e00020020a00704000000e00220230b245c410080073a201000000024810600380208480050820870000791a0200008001008a02038012150838448201510208',
# 'connections': '44', 'gid': '1', 'completedLength': '2174400633',
# 'uploadSpeed': '65684', 'uploadLength': '530998862', 'infoHash':
# 'a56975a1517b0a61c4922a0f2effd8884b1e23cf', 'totalLength':
# '8541599865', 'numSeeders': 1}
    def showinfo(self, key):
        data = self._server.aria2.tellStatus(key)
        data['dlsize'] = hrsize(data['completedLength'])
        data['ulsize'] = hrsize(data['uploadLength'])
        FMT = (
            "Status: %(status)s\nDone: %(percent_done)s\n"
            "Total size: %(size_total)s\n"
            "Connected nodes: %(connections)s (%(numSeeders)s seeders)\n"
            "Downloaded: %(dlsize)s (%(downloadSpeed)s)\n"
            "Uploaded: %(ulsize)s (%(uploadSpeed)s)\n"
            "\n%(numfiles)s file(s) in %(path)s"
        )
        return self.prettify(data, tpl=FMT)

    def showfiles(self, key):
        FMT = "%(index)s\t%(selected)s\t%(size)s\t%(path)s"
        files = self._server.aria2.getFiles(key)
        prettylist = []
        for f in files:
            f['size'] = hrsize(f['length'])
            f['selected'] = {'true': '+'}.get(f['selected'], '-')
            prettylist.append(FMT % f)
        return '\n'.join(prettylist)

# {'downloadSpeed': '0', 'amChoking': 'true', 'ip': 'XX.XXX.XXX.XXX',
# 'bitfield':
# 'bbf7b75f7f37fff5eff6f800000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000',
# 'seeder': 'false', 'peerId':
# '%2DUT160B%2D%D3%81%3F%E3J%231%D1%FA%ECX%95', 'peerChoking': 'true',
# 'uploadSpeed': '0', 'port': '14688'}
    def showpeers(self, key):
        FMT = "%(ip)s\t%(seeder)s\t%(amChoking)s %(peerChoking)s\t%(percent_done)3d %%\t%(dlspeed)12s dl\t%(upspeed)12s up"
        peers = self._server.aria2.getPeers(key)
        prettylist = []
        for p in peers:
            p['dlspeed'] = hrspeed(p['downloadSpeed'])
            p['upspeed'] = hrspeed(p['uploadSpeed'])
            p['seeder'] = 'S' if p['seeder'] == 'true' else 'L'
            # I'm not sure if this is the right way to detect percentage, but
            # it works more or less accurate
            # We ignore last bit since it may be "incomplete"
            p['percent_done'] = 100*float(sum(int(i,16) for i in p['bitfield'][:-1]))/sum(int(i,16) for i in 'f'*len(p['bitfield'][:-1]))
            p['amChoking'] = 'x' if p['amChoking'] == 'true' else ' '
            p['peerChoking'] = 'c' if p['peerChoking'] == 'true' else ' '
            prettylist.append(FMT % p)
        return '\n'.join(prettylist)

    def add_torrent(self, torrent, path=None):
        options = {}
        if path:
            assert os.path.isdir(path)
            options['dir'] = path
        assert os.path.isfile(torrent)
        assert os.stat(torrent).st_size <= 1024**3, "Torrent file size should be under 1MiB"
        data = xmlrpclib.Binary(data=open(torrent).read())
        return self._server.aria2.addTorrent(data, [], options)

    def get_path(self, key):
        files = [f['path'] for f in self._server.aria2.getFiles(key)]
        if len(files) == 1:
            path = files[0]
        else:
            path = os.path.dirname(os.path.commonprefix(files))
        if path.startswith(os.environ['HOME']):
            return (path.replace(os.environ['HOME'], '~', 1), len(files))
        return (path, len(files))


    def prettify(self, r, tpl=None):
        if not tpl: tpl = self.FMTSTRING
        if r['completedLength'] == r['totalLength']:
            r['percent_done'] = 'done'
        else:
            percent_done = float(r['completedLength'])/float(r['totalLength'])*100
            r['percent_done'] = '%.2f%%' % percent_done
        path, numfiles = self.get_path(r['gid'])
        r['path'] = path
        r['numfiles'] = numfiles
        r['uploadSpeed'] = hrspeed(r['uploadSpeed'])
        r['downloadSpeed'] = hrspeed(r['downloadSpeed'])
        r['size_total'] = hrsize(r['totalLength'])
        return tpl % r

    def prettylist(self):
        data = self.listall()
        res = []
        for r in data:
            res.append(self.prettify(r))
        return '\n'.join(res)


if __name__ == "__main__":
    from optparse import OptionParser, OptionGroup, OptionValueError

    def check_depends(option, opt_str, value, parser):
        if not parser.values.tid:
            raise OptionValueError("Please select download with -t")
        setattr(parser.values, option.dest, True)

    parser = OptionParser()
    parser.add_option("-l", "--list", action="store_true", dest="showlist", default=False, help="list all downloads")
    parser.add_option("-a", "--add", dest="torrent", help="add torrent file")
    parser.add_option("-d", "--dir", dest="directory", help="directory to download added torrent to")
    torrentactiongroup = OptionGroup(parser, "Per-download actions")
    torrentactiongroup.add_option("-t", "--torrent", dest="tid", type="int", help="select torrent by id")
    torrentactiongroup.add_option("-i", "--info", action="callback", callback=check_depends, dest="showinfo", default=False, help="show torrent info")
    torrentactiongroup.add_option("-f", "--files", action="callback", callback=check_depends, dest="showfiles", default=False, help="show torrent files")
    torrentactiongroup.add_option("-p", "--peers", action="callback", callback=check_depends, dest="showpeers", default=False, help="show peers")
    parser.add_option_group(torrentactiongroup)
    conngroup = OptionGroup(parser, "XML-RPC connection options")
    conngroup.add_option("--host", dest="host", default="127.0.0.1", help="XML-RPC service host")
    conngroup.add_option("--port", dest="port", default=6800, help="XML-RPC service port", type="int")
    conngroup.add_option("--user", dest="user", help="XML-RPC service user")
    conngroup.add_option("--pass", dest="password", help="XML-RPC service password")
    parser.add_option_group(conngroup)

    if not sys.argv[1:]:
        parser.print_help()
        raise SystemExit()

    (options, args) = parser.parse_args()
    #print options, args # DEBUG

    # try to get user/password from aria2 config
    aria2conf = os.path.expanduser('~/.aria2/aria2.conf')
    if not all([options.user, options.password]) and os.path.exists(aria2conf):
        aria2settings = open(aria2conf).read()
        aria2settings = [l for l in aria2settings.splitlines() if any(['xml-rpc-passwd' in l, 'xml-rpc-user' in l])]
        aria2settings = dict(map(lambda x: x.split('=',1), aria2settings))
        if 'xml-rpc-user' in aria2settings:
            options.user = aria2settings['xml-rpc-user']
        if 'xml-rpc-passwd' in aria2settings:
            options.password = aria2settings['xml-rpc-passwd']


    aria = AriaWrapper(**options.__dict__)

    if options.showlist:
        print aria.prettylist()
        raise SystemExit()

    if options.torrent:
        if not os.path.isfile(options.torrent):
            raise SystemExit("No such file: %s" % options.torrent)
        added_key = aria.add_torrent(options.torrent, options.directory)
        print "Added torrent id:", added_key
        raise SystemExit()

    if options.tid:
        options.tid = str(options.tid)
        if options.showinfo:
            print aria.showinfo(options.tid)
            raise SystemExit()
        if options.showfiles:
            print aria.showfiles(options.tid)
            raise SystemExit()
        if options.showpeers:
            print aria.showpeers(options.tid)
            raise SystemExit()

