#!/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=None, service=None):
    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:
        return False

def prompt(p=None):
    if p:
        stderr.write(str(p))
    return input()

def format_usage(data):
    try:
        quota = data["response"]["usage"]["traffic_types"][0]
        qr = data["response"]["quota_reset"]
        return "{} remaining ({} per day), {} used. Resets on the {}{}".format(
            sizeof_fmt(quota["allocation"] - quota["used"]),
            sizeof_fmt((quota["allocation"] - quota["used"]) /
                qr["days_remaining"]),
            sizeof_fmt(quota["used"]),
            qr["anniversary"],
            "st" if qr["anniversary"] == 1 else "nd" if qr["anniversary"] == 2 else "rd" if qr["anniversary"] == 3 else "th"
        )
    except KeyError:
        return False

def refresh_token():
    print("Tokens need refreshing", file=stderr)
    from getpass import getpass
    try:
        yn = prompt("Want to do that now? [y] ") or "y"
        if yn and yn[0].lower() == "y":
            user = prompt("Username: ")
            westnet = prompt("Westnet account? [n] ") or "n"
            if westnet and westnet[0].lower() == "y":
                user += "@westnet.com.au"
            password = getpass()
            data = get("https://toolbox.iinet.net.au/cgi-bin/api.cgi",
               params={"_USERNAME": user, "_PASSWORD": password}).json()
            if "token" in data:
                print("export IINET_TOKEN=" + data["token"])
            print("\n".join([ "{}\t{}".format(i, sv["pk_v"]) for i, sv in enumerate(data["response"]["service_list"]) if "Usage" in sv["actions"] ]), file=stderr)
            id_num = prompt("Please select a service (using the ID number): ") or ""
            if id_num and id_num.isdigit():
                print("export IINET_SERVICE=" + data["response"]["service_list"][int(id_num)]["s_token"])
            else:
                print("Invalid ID number, please try again", file=stderr)
                return 1
            return 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:
        exit(refresh_token())

if __name__ == "__main__":
    main()