from configuration import read_config
# region Logger
import logging
from debug import setup_logging, catch_errors

log = logger = logging.getLogger("default")
setup_logging()
# endregion
from statistics import median
from threading import Thread
import speedtest
from dbo import Entry
import routeros_api

config = read_config()
secrets = read_config('secrets')

import time


def mbits(bits):
    return round(bits / 1000000, 2)


@catch_errors
def gather_data():
    log.debug("Gathering data...")
    downloads = []
    uploads = []
    dates = []
    for entry in Entry.select():
        downloads.append(entry.download)
        uploads.append(entry.upload)
        dates.append(entry.date_created)

    return dates, downloads, uploads
import matplotlib
import matplotlib.pyplot as plt

@catch_errors
def generate_plot_image(dates, downloads, uploads, name="speed", description="Speed Graph"):
    log.debug("Genering image output for {}...".format(name))


    dates = matplotlib.dates.date2num(dates)
    fig = plt.figure(figsize=(15, 3))
    plt.plot_date(dates, downloads, linestyle='solid', color="red", linewidth=2, marker=None)
    plt.plot_date(dates, uploads, linestyle='dashed', color="green", linewidth=0.5, marker=None)
    plt.ylabel(description + " (Mbps)")
    plt.tight_layout()
    plt.savefig(read_config()['output_image_path'].format(name))


@catch_errors
def generate_diff(dates, downloads, uploads):
    dl = []
    up = []
    for i, date in enumerate(dates):
        curr = downloads[i] if downloads[i] else 0
        prev = downloads[i-1] if downloads[i-1] and i > 0 else 0
        dl.append(curr - prev)
        curr = downloads[i] if downloads[i] else 0
        prev = downloads[i - 1] if downloads[i - 1] and i > 0 else 0
        up.append(curr - prev)
    generate_plot_image(dates, dl, up, "diff", "Difference Graph")


def generate_updown_plot_simple(downloads, uploads, name, description):
    fig = plt.figure(figsize=(10, 3))
    plt.plot(downloads, linestyle='solid', color="red", linewidth=2, marker=None)
    plt.plot(uploads, linestyle='dashed', color="green", linewidth=0.5, marker=None)
    plt.ylabel(description + " (Mbps)")
    plt.tight_layout()
    plt.savefig(read_config()['output_image_path'].format(name))

@catch_errors
def gather_week_median_data(dates, downloads, uploads):
    log.debug("Gather week median data...")
    dl = []
    up = []

    for hour in range(24):
        dl.append([])
        up.append([])

    for i, date in enumerate(dates):
        hour = date.hour
        dl[hour].append(downloads[i] if downloads[i] else 0)
        up[hour].append(uploads[i] if downloads[i] else 0)

    for i, hour in enumerate(dl):
        dl[i] = median(dl[i]) if len(dl[i]) > 0 else 0

    for i, hour in enumerate(up):
        up[i] = median(up[i]) if len(up[i]) > 0 else 0

    return dl, up

@catch_errors
def generate_week_median(dates, downloads, uploads):
    dl, up = gather_week_median_data(dates, downloads, uploads)
    
    generate_updown_plot_simple(dl, up, "week_median", "Week median")

@catch_errors
def generate_week_median_diff(dates, downloads, uploads):
    dl, up = gather_week_median_data(dates, downloads, uploads)

    downs = []
    ups = []

    for i, down in enumerate(dl):
        if i > 0:
            downs.append(dl[i]-dl[i-1])
            ups.append(up[i]-up[i-1])
        else:
            downs.append(down)
            ups.append(up[i])

    generate_updown_plot_simple(downs, ups, "week_median_diff", "Week median diff")


@catch_errors
def generate_txt_output(dates, downloads, uploads):
    log.debug("Genering txt output...")
    txt = "Date:			Down;			Up;\n"
    for i, date in enumerate(dates):
        download = downloads[i]
        upload = uploads[i]

        txt += f"{date}:			{download} Mbps;			{upload} Mbps\n"
    with open(read_config()['output_txt_path'], "w+") as f:
        f.write(txt)


@catch_errors
def ros_fastrack_enable(enable):
    '''
    Enable or disable fasttrack
    :param enable: True to enable, False to disable
    '''

    if enable:
        disabled = 'false'
    else:
        disabled = 'true'

    log.info(f"Changing fasttrack disabled to {disabled}")

    connection = routeros_api.RouterOsApiPool(config['ros_ip'], username=secrets["ros_login"],
                                              password=secrets["ros_password"], plaintext_login=True)
    api = connection.get_api()

    filter = api.get_resource('/ip/firewall/filter')

    for i in filter.get():
        if config['ros_fasttrack_comment'] in i['comment']:
            filter.set(id=i['id'], disabled=disabled)
    connection.disconnect()


@catch_errors
def ros_dynamic_speed(upload, download):
    '''
    Adjust router os queue speed.
    :param upload: Upload speed from speedtest
    :param download: Download speed from speedtest
    '''
    log.debug(f"Set Dynamic Speed to: DOWN {mbits(download)} mbps; UP {mbits(upload)} mbps")
    connection = routeros_api.RouterOsApiPool(config['ros_ip'], username=secrets["ros_login"],
                                              password=secrets["ros_password"], plaintext_login=True)
    api = connection.get_api()

    list_queues = api.get_resource('/queue/simple')

    for queue in list_queues.get():
        if queue['name'] in config['ros_queues']:
            log.debug(
                f"Adjust Queue {queue['name']}: max_limit {int(upload)}/{int(download)}")

            if config["ros_du_invert"] == True:
                # Inverting upload and download values, because master queue is most likely applied to the bridge
                list_queues.set(id=queue['id'], max_limit=f"{int(download)}/{int(upload)}")
            else:
                # Not inverting, use this in case master queue is applied to something like LTE interface
                list_queues.set(id=queue['id'], max_limit=f"{int(upload)}/{int(download)}")
    connection.disconnect()


wan_upload = None
wan_download = None
results_dict = None
test_started = False
speedtest_failed = False
downloading = True  # True for download, False for upload
threads = None
servers = []


def reset_globals():
    global wan_upload
    global wan_download
    global results_dict
    global test_started
    global speedtest_failed
    global downloading

    wan_upload = None
    wan_download = None
    results_dict = None
    test_started = False
    speedtest_failed = False
    downloading = True


def threaded_speedtest():
    global test_started
    global downloading
    global results_dict
    global speedtest_failed
    try:
        s = speedtest.Speedtest()
        s.get_servers(servers)
        s.get_best_server()
        print(f"Running test...")
        test_started = True
        downloading = True
        s.download(threads=threads)
        downloading = False
        s.upload(threads=threads, pre_allocate=False)

        results_dict = s.results.dict()
        log.info(
            f"Speedtest.net result: DOWN {mbits(results_dict['download'])} mbps; UP {mbits(results_dict['upload'])} mbps;")
    except Exception as e:
        speedtest_failed = True
        log.error(f"Speedtest.net failed: {e}", exc_info=True)
    return results_dict


def threaded_wan_speed():
    global test_started
    global results_dict
    global downloading
    print("Waiting for test to start...")
    while not test_started:
        if speedtest_failed:
            return
        time.sleep(1)
    print("Allow warm-up...")
    time.sleep(2)
    uploads = []
    downloads = []
    print("Monitoring...")
    while results_dict == None:
        connection = routeros_api.RouterOsApiPool(config['ros_ip'], username=secrets["ros_login"],
                                                  password=secrets["ros_password"], plaintext_login=True)
        api = connection.get_api()

        traffic = api.get_resource('/').call('interface/monitor-traffic',
                                             {'interface': config['ros_wan_interface'], 'once': ' '})[0]

        traffic['rx-bits-per-second'] = int(traffic['rx-bits-per-second'])
        traffic['tx-bits-per-second'] = int(traffic['tx-bits-per-second'])

        if downloading:
            downloads.append(traffic['rx-bits-per-second'])
            print(f"DL: {mbits(traffic['rx-bits-per-second'])} mbps;")
        else:
            uploads.append(traffic['tx-bits-per-second'])
            print(f"UP: {mbits(traffic['tx-bits-per-second'])} mbps;")

        time.sleep(1)
    global wan_download
    global wan_upload
    wan_download = median(downloads)
    wan_upload = median(uploads)
    log.info(f"Monitor result: {mbits(wan_download)} mbps;  {mbits(wan_upload)} mbps;")


def test_speed():
    global wan_download
    tries = 2
    for i in range(tries):
        reset_globals()
        sws = Thread(target=threaded_wan_speed)
        st = Thread(target=threaded_speedtest)
        st.start()
        sws.start()
        st.join()
        sws.join()
        if wan_download >= config['ros_doublecheck_speed']:
            break
        elif i < tries - 1:
            log.warning("Speed is below ros_doublecheck_speed limit, retrying in a few seconds.")
            time.sleep(10)
    return


def generate_database_reports():
    dates, downloads, uploads = gather_data()

    generate_txt_output(dates, downloads, uploads)

    generate_plot_image(dates, downloads, uploads)

    generate_week_median(dates, downloads, uploads)

    generate_week_median_diff(dates, downloads, uploads)

    generate_diff(dates, downloads, uploads)


import sys
from netutils import test_intertnet_connection


def on_fail_or_no_connection():
    ros_fastrack_enable(False)

    entry = Entry()
    entry.upload = None
    entry.download = None
    entry.save()

    generate_database_reports()

    log.warning("No internet connection! Exiting.")
    sys.exit()


if __name__ == "__main__":
    '''
    This script will run a few speed tests, calculate average upload and download speeds and record them into database.
    Once finished it will also generate an image with graph plotted.
    '''

    # generate_plot_image(*gather_data())
    # import sys
    # sys.exit()
    from random import uniform

    try:
        log.debug("Test internet connection...")

        if config["ros_dynamic_speed"]:
            ros_fastrack_enable(True)
            time.sleep(5)

        test_speed()

        if speedtest_failed:
            on_fail_or_no_connection()

        entry = Entry()
        entry.upload = mbits(wan_upload)
        entry.download = mbits(wan_download)
        entry.save()

        if wan_download < config['ros_minimum_speed']:
            wan_download = config['ros_minimum_speed']

        if config["ros_dynamic_speed"]:
            ros_dynamic_speed(wan_upload, wan_download)
            ros_fastrack_enable(False)

        generate_database_reports()


    except:
        log.error("Error!", exc_info=True)

        if config["ros_dynamic_speed"]:
            ros_fastrack_enable(False)