#!/usr/bin/python3

import logging
import logging.handlers

import os
from pathlib import Path

import csv
from datetime import datetime, timedelta
import time
import dateutil.parser as dp

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.cbook as cbook
import matplotlib.dates as mdates
from matplotlib.ticker import FuncFormatter
from matplotlib import rcParams

rcParams["font.family"] = "sans-serif"
rcParams["font.sans-serif"] = ["DejaVu Sans"]
rcParams["font.weight"] = "medium"
rcParams["font.size"] = 10.0
rcParams["axes.linewidth"] = 0.4
rcParams["legend.fontsize"] = 8.0

from scipy.interpolate import splev, splrep

import tweepy
from instabot import Bot
from mastodon import Mastodon

MAINTENANCE = "/home/ubuntu/bndwdth.maintenance"

CSV_FILE = "/home/ubuntu/bndwdth.csv"

SPEED_MAX_DOWN = 1184890880         # 1,130 Mbps
SPEED_MIN_DOWN = SPEED_MAX_DOWN / 2 #   565 Mbps
SPEED_MAX_UP   = 54525952           #    52 Mbps
SPEED_MIN_UP   = SPEED_MAX_UP / 2   #    26 Mbps

DOWN_LINE_COLOUR = "blue"
UP_LINE_COLOUR   = "magenta"
PING_LINE_COLOUR = "orange"
TTFB_LINE_COLOUR = "red"

logger = logging.getLogger('bndwdth-log')
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.handlers.SysLogHandler(address = '/dev/log'))

def notice(s):
    logger.debug(s)
    print(s)

def mean(numbers):
    return float(sum(numbers)) / max(len(numbers), 3)

#
## Get bandwidth in human readable form
#

def ghr(b, u=True):
    return get_human_readable(b, u)

def get_human_readable(bps, units=True):
    K = 1024
    M = K * K
    G = M * K
    Bps = bps / 8
    if bps < K:
        return "%d%s" % (bps, "bps" if units else "")
    if bps < M:
        return "%d%s" % ((bps / K), "Kbps" if units else "")
    return "%d%s" % ((bps / M), "Mbps" if units else "")

def sufficiently_fast_up(speed):
    return speed["up"] > SPEED_MIN_UP

def sufficiently_fast_down(speed):
    return speed["down"] > SPEED_MIN_DOWN

def host_missing(speed):
    return speed["up"] == 0 and speed["down"] == 0

def insufficiently_fast(speed):
    return speed["up"] < SPEED_MIN_UP and speed["down"] < SPEED_MIN_DOWN

def sufficiently_fast(speed):
    return sufficiently_fast_up(speed) and sufficiently_fast_down(speed)

def extra_fast(speed):
    return speed["up"] > SPEED_MAX_UP and speed["down"] > SPEED_MAX_DOWN

#
## Generate a graph of the previous 24 hours bandwidth
#

def parse_bandwidth(duration=24):
    # datetime,down speed,up speed,ping,latency/ttfb,host up,latency up,host down,latency down
    date = []
    down = []
    up = []
    ping = []
    ttfb = []
    with open(CSV_FILE, "r") as c:
        reader = csv.reader(c)
        now = datetime.now()
        for row in reader:
            d = dp.parse(row[0])
            if now - d < timedelta(hours=duration):
                date.append(d)
                if len(row) > 1:
                    down.append(int(row[1]))
                    up.append(int(row[2]))
                    if len(row) > 3:
                        ping.append(float(row[3]))
                        if len(row) > 4:
                            if len(row) > 8:
                                ttfb.append(min(float(row[6]), float(row[8])))
                            else:
                                ttfb.append(float(row[4]))
                        else:
                            ttfb.append(0)
                    else:
                        ping.append(0)
                        ttfb.append(0)
                else:
                    down.append(0)
                    up.append(0)
                    ping.append(0)
                    ttfb.append(0)
    return { "date": date, "down": down, "up": up, "ping": ping, "ttfb": ttfb }

def average_bandwidth(data, avg=2):
    date = []
    down = []
    up = []
    ping = []
    ttfb = []
    for d in range(int(len(data["date"]) / avg)):
        t = d * avg
        date.append(data["date"][t])

        ad = 0
        for i in range(avg):
            ad = ad + data["down"][t + i]
        ad = ad / avg
        down.append(ad)

        au = 0
        for i in range(avg):
            au = au + data["up"][t + i]
        au = au / avg
        up.append(au)

        ap = 0
        for i in range(avg):
            ap = ap + data["ping"][t + i]
        ap = ap / avg
        ping.append(ap)

        ac = 0
        for i in range(avg):
            ac = ac + data["ttfb"][t + i]
        ac = ac / avg
        ttfb.append(ac)

    return { "date": date, "down": down, "up": up, "ping": ping, "ttfb": ttfb }

def generate_graph(data, scale="log", xkcd=False):
    if xkcd:
        plt.xkcd()
    fig, speed_plot = plt.subplots()
    duration_plot = speed_plot.twinx()

    timestamp = []
    down = []
    up = []
    ping = []
    ttfb = []

    data_points = len(data["date"])

    for i in range(data_points):
        timestamp.append(float(data["date"][i].timestamp()))
        down.append(float(data["down"][i]) / 1000000)
        up.append(float(data["up"][i]) / 1000000)
        ping.append(float(data["ping"][i]))
        ttfb.append(float(data["ttfb"][i]))

    now = datetime.now()
    plt.suptitle("Bandwidth [" + now.strftime("%Y-%m-%d %-I%p") + "]")
    speed_plot.set_xlabel("Date / Time")
    speed_plot.set_ylabel("Bandwidth (Mbps)")
    duration_plot.set_ylabel("Delay (ms)")

    # calculate (and draw) a pretty looking line
    timestamp_min = np.array(timestamp).min()
    timestamp_max = np.array(timestamp).max()

    time_continuous = np.linspace(int(timestamp_min), int(timestamp_max), int(timestamp_max - timestamp_min))
    time = np.linspace(int(timestamp_min), int(timestamp_max), len(timestamp))

    down_continuous = splev(time_continuous, splrep(time, down))
    up_continuous = splev(time_continuous, splrep(time, up))
    ping_continuous = splev(time_continuous, splrep(time, ping))
    ttfb_continuous = splev(time_continuous, splrep(time, ttfb))

    time_scale = []
    time_continuous_scale = []
    for i in time:
        time_scale.append(datetime.fromtimestamp(i))
    down_expected = []
    up_expected = []
    for i in time_continuous:
        time_continuous_scale.append(datetime.fromtimestamp(i))
        down_expected.append(float(ghr(SPEED_MAX_DOWN, False)))
        up_expected.append(float(ghr(SPEED_MAX_UP, False)))

    ln =      speed_plot.plot(time_continuous_scale, down_continuous, color=DOWN_LINE_COLOUR, markersize=2, linewidth=1, label="Down")
    ln = ln + speed_plot.plot(time_continuous_scale,   up_continuous, color=  UP_LINE_COLOUR, markersize=2, linewidth=1, label="Up")
    ln = ln + speed_plot.plot(time_continuous_scale, ping_continuous, color=PING_LINE_COLOUR, markersize=2, linewidth=1, label="Ping")
    ln = ln + speed_plot.plot(time_continuous_scale, ttfb_continuous, color=TTFB_LINE_COLOUR, markersize=2, linewidth=1, label="TTFB")
    ln = ln + speed_plot.plot(time_continuous_scale,   down_expected, color=DOWN_LINE_COLOUR, markersize=2, linewidth=1, linestyle="dashed", label="Expected: " + str(int(SPEED_MAX_DOWN / 1024 / 1024)) + "Mbps")
    ln = ln + speed_plot.plot(time_continuous_scale,     up_expected, color=  UP_LINE_COLOUR, markersize=2, linewidth=1, linestyle="dashed", label="Expected: " + str(int(SPEED_MAX_UP   / 1024 / 1024)) + "Mbps")

    # finish up
    plt.legend(ln, [l.get_label() for l in ln], loc= "upper center", ncol=4, bbox_to_anchor=(0.5, 1.06), framealpha=1).get_frame().set_linewidth(0.4)

    speed_plot.xaxis.set_major_formatter(mdates.DateFormatter("%m/%d %H:%M"))

    plt.xlim([datetime.fromtimestamp(time_continuous[0]), datetime.now()])

    speed_plot.set_yscale(scale)
    speed_plot.set_ylim(1, 2000)
    speed_plot.yaxis.set_major_formatter(FuncFormatter(lambda y, p: format(int(y), ',')))

    duration_plot.set_yscale(scale)
    duration_plot.set_ylim(1, max(data["ping"] + data["ttfb"]) * 1.5)
    duration_plot.yaxis.set_major_formatter(FuncFormatter(lambda y, p: format(int(y), ',')))
    fig.autofmt_xdate()
    fig.set_size_inches(10, 8)

    out = "%s-bw-%s.png" % (datetime.now().strftime("%Y%m%d%H%M"), scale)
    plt.savefig(out, dpi=320)

    return out

#
## Send the latest results to Twitter
#

TWITTER_ACCESS_TOKEN        = ""
TWITTER_ACCESS_TOKEN_SECRET = ""
TWITTER_CONSUMER_KEY        = ""
TWITTER_CONSUMER_SECRET     = ""

def get_twitter():
    twitter_auth_keys = {
        "consumer_key"        : TWITTER_CONSUMER_KEY,
        "consumer_secret"     : TWITTER_CONSUMER_SECRET,
        "access_token"        : TWITTER_ACCESS_TOKEN,
        "access_token_secret" : TWITTER_ACCESS_TOKEN_SECRET
        }
    auth = tweepy.OAuthHandler(
        twitter_auth_keys['consumer_key'],
        twitter_auth_keys['consumer_secret']
        )
    auth.set_access_token(
        twitter_auth_keys['access_token'],
        twitter_auth_keys['access_token_secret']
        )
    return tweepy.API(auth)

MASTODON_CLIENT_KEY    = ""
MASTODON_CLIENT_SECRET = ""
MASTODON_ACCESS_TOKEN  = ""
MASTODON_HOST          = ""
MASTODON_USERNAME      = ""
MASTODON_PASSWORD      = ""

def get_mastodon():
    mastodon = Mastodon(client_id = MASTODON_CLIENT_KEY,
                    client_secret = MASTODON_CLIENT_SECRET,
                    access_token  = MASTODON_ACCESS_TOKEN,
                    api_base_url  = MASTODON_HOST)
    mastodon.log_in(username  = MASTODON_USERNAME,
                    password  = MASTODON_PASSWORD,
                    scopes    = ["read", "write"])
    return mastodon

def alert_network_down(graph):
    t = get_twitter()
    s = "@virginmedia it looks like there might be network trouble in #Northants! My #awscloud EC2 instance has lost its connection with my home PC."
    if t is not None:
        t.statuses.update(status=s)
    m = get_mastodon()
    if m is not None:
        m.toot(s)

def send_update(speed, graph):
    # upload graph to twitter/mastodon
    t = get_twitter()
    m = get_mastodon()
    twit_img_ids = []
    mast_img_ids = []
    if t is not None:
        twit_img_ids.append(t.media_upload(graph).media_id)
    if m is not None:
        mast_img_ids.append(m.media_post(graph).id)

    # send tweet
    t = get_twitter()
    s = "Average #bandwidth in #Northants for the last hour was %s down and %s up" % (ghr(speed["down"]), ghr(speed["up"]))
    if os.path.exists(MAINTENANCE):
        s = "Down for scheduled maintenance"
    else:
        if insufficiently_fast(speed):
            s = "@virginmedia Why's my hourly average #bandwidth in #Northants only %s down / %s up? Not even close to %s/%s" % (ghr(speed["down"]), ghr(speed["up"]), ghr(SPEED_MAX_DOWN, False), ghr(SPEED_MAX_UP, False))
        elif not sufficiently_fast_down(speed):
            s = "@virginmedia Why's my hourly average download #bandwidth in #Northants only %s? Nowhere close to %s! (%s up)" % (ghr(speed["down"]), ghr(SPEED_MAX_DOWN, False), ghr(speed["up"]))
        elif not sufficiently_fast_up(speed):
            s = "@virginmedia Why's my hourly average upload #bandwidth in #Northants only %s? Not even close to %s! (%s down)" % (ghr(speed["up"]), ghr(SPEED_MAX_UP, False), ghr(speed["down"]))
        elif extra_fast(speed):
            s = "@virginmedia Excellent: my broadband #bandwidth in #Northants is %s down and %s up! #M600Fibre #superfast" % (ghr(speed["down"]), ghr(speed["up"]))

    if t is not None:
        t.update_status(status=s, media_ids=twit_img_ids)
    if m is not None:
        m.status_post(status=s, media_ids=mast_img_ids)
    notice(s)

if __name__ == "__main__":
    s = parse_bandwidth(duration=1)
    s = average_bandwidth(s, len(s["date"]))
    s["up"] = s["up"][0]
    s["down"] = s["down"][0]

    graph = generate_graph(average_bandwidth(parse_bandwidth(), 4), scale="log")
    alert = Path("bndwdth-alert")
    if int(s["up"]) > 0 and int(s["down"]) > 0:
        if alert.is_file():
            alert.unlink()
        send_update(s, graph)
    else:
        if not alert.is_file():
            alert.open("a")
            alert_network_down(graph)

    g = Path(graph)
    if g.is_file():
        g.unlink()