#!/usr/bin/env python3
from requests import get
from os import environ
from sys import exit, stderr, stdout

STDOUT = stdout

# http://stackoverflow.com/a/1094933
def sizeof_fmt(num, suffix='B'):
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)

def get_service_data(token=environ["IINET_TOKEN"], service=environ["IINET_SERVICE"]):
    url = "https://toolbox.iinet.net.au/cgi-bin/api.cgi?Usage&_TOKEN={}&_SERVICE={}"
    try:
        data = get(url.format(environ["IINET_TOKEN"], environ["IINET_SERVICE"])).json()
        return data
    except KeyError:
        print("set environ vars pls", file=stderr)
        return False

def format_usage(data):
    try:
        data = data["response"]["usage"]["traffic_types"][0]
        return "{} remaining ({} used)".format(sizeof_fmt(data["allocation"] - data["used"]), sizeof_fmt(data["used"]))
    except KeyError:
        return False

def refresh_token():
    print("token needs refreshing", file=stderr)
    from getpass import getpass
    try:
        stdout = stderr
        yn = input("Want to do that now? [y] ") or "y"
        if yn and yn[0].lower() == "y":
            user = input("Username: ")
            westnet = input("Westnet account? [n] ") or "n"
            if westnet and westnet[0].lower() == "y":
                user += "@westnet.com.au"
            stdout = STDOUT
            password = getpass()
            print("Token:", file=stderr)
            return get("https://toolbox.iinet.net.au/cgi-bin/api.cgi", params={"_USERNAME": user, "_PASSWORD": password}).json()["token"], 0
        else:
            raise KeyboardInterrupt
    except KeyboardInterrupt:
        return "", 130

def main():
    data = format_usage(get_service_data())
    if data:
        print(data, file=stderr)
        exit(0)
    else:
        pl, retcode = refresh_token()
        print(pl)
        exit(retcode)

if __name__ == "__main__":
    main()