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

log = logger = logging.getLogger("default")
setup_logging()
# endregion

from dbo import Entry


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

def generate_plot_image(dates, downloads, uploads):
    log.debug("Genering image output...")
    import matplotlib
    import matplotlib.pyplot as plt

    dates = matplotlib.dates.date2num(dates)
    plt.plot_date(dates, downloads, fmt="b-")
    plt.ylabel('Download Speed Mbps')
    plt.tight_layout()
    plt.savefig(read_config()['output_image_path'])

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)

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.
    '''
    from random import uniform
    try:
        import speedtest

        servers = []
        threads = None

        for i in range (0, 3):
            try:
                log.debug("Initializing speedtest...")
                s = speedtest.Speedtest()


                log.debug(f"Running test...")
                s.get_servers(servers)
                s.get_best_server()
                s.download(threads=threads)
                s.upload(threads=threads, pre_allocate=False)

                results_dict = s.results.dict()
                download = round(results_dict['download']/1000000, 2)
                upload = round(results_dict['upload']/1000000, 2)
                # download = uniform(0,2)
                # upload = uniform(0,2)
                break
            except:
                log.error(f"Test failed, try {i+1}/3", exc_info=True)

        log.debug(f"{download}mbps, {upload}mbps")

        entry = Entry()
        entry.upload = upload
        entry.download = download
        entry.save()
    except:
        log.error("Data record error.", exc_info=True)

    try:
        dates, downloads, uploads = gather_data()

        try:
            generate_txt_output(dates, downloads, uploads)
        except:
            log.error("Unable to save text file.", exc_info=True)

        try:
            generate_plot_image(dates, downloads, uploads)
        except:
            log.error("Unable to save plot file.", exc_info=True)

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