#!/usr/bin/env python3
from optparse import OptionParser
import os
import sys
import json
import subprocess
import time
import shlex
import shutil

def wizard(params,server_info):
    # get name of test
    params["test_name"] = input("Provide a test name: ")
    params["room_temp_c"], params["room_temp_f"] = get_room_temperature()
    params["max_duration"] = get_test_duration()
    params["polling_interval"], tmap_result = get_tmap_interval()
    params["fio_flag"] = get_fio_setting()
    params["fio_write"] = get_fio_write()
    params["graph_flag"] = get_graph_flag()
    params["graph_settings"] = get_graph_settings(params,server_info,tmap_result) if params["graph_flag"] else {}
    params["rsync_settings"] = get_rsync_params()

def get_room_temperature():
    temp_type = ""
    room_temp = ""
    while(not room_temp.isnumeric()):
        room_temp = input("Provide room temperature reading (number only): ")
        if not room_temp.isnumeric():
            print("invalid room temperature provided, try again.")
    while(temp_type not in ["c","f"]):
        temp_type = input("Temperature reporting (c=celcius, f=farenheit): ")
        if temp_type not in ["c","f"]:
            print(f"Invalid selection \"{temp_type}\"")
    
    rtc = 0
    rtf = 0

    if temp_type == "c":
        rtc = int(room_temp)
        rtf = int((int(room_temp) * (9/5)) + 32)
    elif temp_type == "f":
        rtf = int(room_temp)
        rtc = int((int(room_temp) - 32) * (5/9))

    return rtc, rtf

def get_test_duration():
    valid = False
    dur_s = 0
    dur_str = None
    while(not valid):
        dur_str = input("maximum test duration (HH:MM): ")
        dur_args = dur_str.split(":")
        if len(dur_args) == 2 and dur_args[0].isnumeric() and dur_args[1].isnumeric():
            dur_s = int(dur_args[0])*3600 + int(dur_args[1])*60
            valid = True
        else:
            print("Invalid test duration provided. Try again.")
            valid = False
    return dur_s

def get_tmap_interval():
    start_time = time.time()
    print("determining ideal measurement interval, please wait...")
    tmap_result = run_tmap(start_time)
    measurement_duration = tmap_result["duration"]
    measurement_interval = ((measurement_duration + int((measurement_duration*0.25))) * 3) + 5
    print(f"Ideal interval between measurements calculated to be: {measurement_interval} s")
    selection = ""
    while(selection not in ["y","Y"]):
        selection = input("Would you like to use this interval between measurements? (y/n): ")
        if selection in ["n","N"]:
            user_interval = input("Provide an interval between measurements (in seconds): ")
            if user_interval.isnumeric() and int(user_interval) > (measurement_duration + int((measurement_duration*0.25)) + 5):
                measurement_interval = int(user_interval)
            elif user_interval.isnumeric():
                print("Warning: measurement interval provided will not allow enough time to perform a fio job between readings.")
                measurement_interval = int(user_interval)
            else:
                print("Invalid measurement interval provided. Try again")
    return measurement_interval, tmap_result

def run_tmap(start_time):
    # run tmap and append the time that the measurement started as well as the duration of the measurement
    tmap_start_time = time.time()
    tmap = subprocess.Popen(
        shlex.split("./tmap"), stdout=subprocess.PIPE, universal_newlines=True)
    tmap_output = tmap.communicate()
    tmap_json = json.loads(tmap_output[0])
    tmap_end_time = time.time()
    tmap_json["time"] = int(tmap_start_time - start_time)
    tmap_json["duration"] = int(tmap_end_time - tmap_start_time)
    return tmap_json

def get_fio_setting():
    selection = ""
    flag = True
    while selection not in ["y","Y","n","N"]:
        selection = input("Would you like to read from all disks between measurements using fio? (y/n): ")
        if selection not in ["y","Y","n","N"]:
            print("Invalid selection. Try again.")
    if selection in ["y","Y"]:
        flag = True
    if selection in ["n","N"]:
        flag = False
    return flag

def get_fio_write():
    selection = ""
    flag = True
    while selection not in ["y","Y","n","N"]:
        selection = input("Would you like to perform WRITES on all storage disks between measurements using fio? (y/n): ")
        if selection not in ["y","Y","n","N"]:
            print("Invalid selection. Try again.")
    if selection in ["y","Y"]:
        confirm = input("Are you sure? Data can and will be corrupted on your storage disks if you write. (type 'yesiamsure' to confirm): ")
        if confirm == "yesiamsure":
            flag = True
        else:
            print("yesiamsure not entered. Defaulting to read only.")
            flag = False
    if selection in ["n","N"]:
        flag = False
    return flag

def get_graph_flag():
    selection = ""
    flag = True
    while selection not in ["y","Y","n","N"]:
        selection = input("Would you like a graph of the test results? (y/n): ")
        if selection not in ["y","Y","n","N"]:
            print("Invalid selection. Try again.")
    if selection in ["y","Y"]:
        flag = True
    if selection in ["n","N"]:
        flag = False
    return flag

def get_graph_settings(params,server_info,tmap_result):
    graph_vars = {
        "title": ""
    }
    graph_vars["title"] = "{model} - {dc} Drives - {rt}°C".format(model=server_info["Model"],dc=len(tmap_result["DISKS"]),rt=params["room_temp_c"])
    print("Auto-generated title for graph: {t}".format(t=graph_vars["title"]))
    confirm_title = input("Would you like to use this Auto-generated title? (y/n): ")
    if confirm_title not in ["Y", "y"]:
        title_done = False
        while not title_done:
            temp_title = input("Enter a title for graph: ")
            accept_title = input("Use this title for graph? \"{tt} - {rt}°C\" (y/n): ".format(tt=temp_title,rt=params["room_temp_c"]))
            if accept_title in ["Y","y"]:
                graph_vars["title"] = "{tt} - {rt}°C".format(tt=temp_title,rt=params["room_temp_c"])
                title_done = True
            else:
                print("Trying again..")
    return graph_vars

def get_rsync_params():
    rsync_settings = {
        "rsync_flag": False,
        "rsync_user": "",
        "rsync_host": "",
        "rsync_dir": "",
        "rsync_command": ""
    }

    response = input("Would you like to send the results to a host using rsync? (y/n): ")
    if response in ["y","Y"]:
        confirm = "n"
        while(confirm not in ["y","Y"]):
            rsync_settings["rsync_flag"] = True
            rsync_settings["rsync_user"] = input("enter user for rsync target (ex: root): ")
            rsync_settings["rsync_host"] = input("enter ip/hostname to rsync to (ex: server): ")
            rsync_settings["rsync_dir"] = input("enter the directory on target machine. (ex: /home/user): ")
            if rsync_settings["rsync_dir"] and rsync_settings["rsync_dir"][-1] == "/":
                rsync_settings["rsync_dir"] = rsync_settings["rsync_dir"][:-1]
            print("rsync command: rsync -avh <files> {u}@{h}:{d}/".format(u=rsync_settings["rsync_user"], h=rsync_settings["rsync_host"], d=rsync_settings["rsync_dir"]))
            confirm = input("Use of the command shown above? (y,n): ")
            if confirm in ["y","Y"]:
                rsync_settings["rsync_command"] = "rsync -avh <files> {u}@{h}:{d}/".format(u=rsync_settings["rsync_user"], h=rsync_settings["rsync_host"], d=rsync_settings["rsync_dir"])
            else:
                cancel = input("abandon using rsync? (y/n): ")
                if cancel in ["y","Y"]:
                    confirm = "y"
    if len(rsync_settings["rsync_command"]) == 0:
        rsync_settings["rsync_flag"] = False
    
    return rsync_settings
        

def get_server_info():
    # get server_info.json file and read it in
    si_path = "/etc/45drives/server_info/server_info.json"
    if not os.path.exists(si_path) or not os.path.isfile(si_path):
        print(f"Required File: '{si_path}' was not found. Ensure that 45drives-tools is installed and run 'dmap'. Then try again.")
        sys.exit(1)
    server_info = None
    with open(si_path,"r") as si_f:
        try:
            server_info = json.load(si_f)
        except:
            print("Failed to load json data from '{si_path}'")
            sys.exit(1)
    return server_info

def run_tplot(params):
    #./tplot -o output/$TEST_NAME.csv -d 120 -i 10 -s 3 -f
    steady_state = 3 + int((60/params["polling_interval"]))

    command = "./tplot -o output/{tn}.csv -d {dur} -i {pi} -s {ss} {fio} {write}".format(
        tn=params["test_name"],
        dur=params["max_duration"],
        pi=params["polling_interval"],
        ss=steady_state,
        fio = "-f" if params["fio_flag"] else "",
        write = "-w" if params["fio_flag"] and params["fio_write"] else ""
        )
    return os.system(command)


def run_csv_converter(params):
    #./csv_converter -i output/$TEST_NAME.csv -o chart_csv/$TEST_NAME.csv
    print("Converting {tn}.csv".format(tn=params["test_name"]))
    command = "./csv_converter -i output/{tn}.csv -o chart_csv/{tn}.csv".format(tn=params["test_name"])
    return os.system(command)

def run_make_graph(params):
    #./make_graph -i chart_csv/$TEST_NAME.csv -o graphs/$TEST_NAME.png -t "$GRAPH_TITLE"
    print("Making graph")
    command = "./make_graph -i chart_csv/{tn}.csv -o graphs/{tn}.png -t '{gt}'".format(tn=params["test_name"],gt=params["graph_settings"]["title"])
    return os.system(command)

def run_rsync(params):
    os.makedirs("./tbench_results/{tn}".format(tn=params["test_name"]),exist_ok=True)
    shutil.copyfile("./output/{tn}.csv".format(tn=params["test_name"]), "./tbench_results/{tn}/{tn}_raw.csv".format(tn=params["test_name"]))
    shutil.copyfile("./chart_csv/{tn}.csv".format(tn=params["test_name"]), "./tbench_results/{tn}/{tn}_chart.csv".format(tn=params["test_name"]))
    if params["graph_flag"]:
        shutil.copyfile("./graphs/{tn}.png".format(tn=params["test_name"]), "./tbench_results/{tn}/{tn}.png".format(tn=params["test_name"]))

    command = "rsync -avh ./tbench_results/{tn} {u}@{h}:{d}/".format(tn=params["test_name"],u=params["rsync_settings"]["rsync_user"],h=params["rsync_settings"]["rsync_host"],d=params["rsync_settings"]["rsync_dir"])
    return os.system(command)

def run_rsync_test(params):
    os.makedirs("./tbench_results/{tn}".format(tn=params["test_name"]),exist_ok=True)
    with open('./tbench_results/{tn}/{tn}.json'.format(tn=params["test_name"]), 'w') as outfile:
        outfile.write(json.dumps(params))
    command = "rsync -avh ./tbench_results/{tn} {u}@{h}:{d}/".format(tn=params["test_name"],u=params["rsync_settings"]["rsync_user"],h=params["rsync_settings"]["rsync_host"],d=params["rsync_settings"]["rsync_dir"])
    print("testing out rsync.")
    if os.system(command) != 0:
        print("rsync test failed. You can run rsync yourself afterward using this command: \n{c}".format(c=command))
    else:
        print("rsync test passed.")
    return os.system(command)

def show_welcome():
    msg = [
    "                                                            ",
    "   /$$     /$$                                     /$$      ",
    "  | $$    | $$                                    | $$      ",
    " /$$$$$$  | $$$$$$$   /$$$$$$  /$$$$$$$   /$$$$$$$| $$$$$$$ ",
    "|_  $$_/  | $$__  $$ /$$__  $$| $$__  $$ /$$_____/| $$__  $$",
    "  | $$    | $$  \ $$| $$$$$$$$| $$  \ $$| $$      | $$  \ $$",
    "  | $$ /$$| $$  | $$| $$_____/| $$  | $$| $$      | $$  | $$",
    "  |  $$$$/| $$$$$$$/|  $$$$$$$| $$  | $$|  $$$$$$$| $$  | $$",
    "   \___/  |_______/  \_______/|__/  |__/ \_______/|__/  |__/",
    "                                                            "]
    for line in msg:
        print(line)

def main():

    show_welcome()
    params = {
        "test_name": "",
        "room_temp_c": "",
        "room_temp_f": "",
        "max_duration": 3600,
        "polling_interval": 60,
        "fio_flag": True,
        "write_flag": False,
        "graph_flag": True,
        "graph_settings": {},
        "rsync_settings": {},
        "script_commands": {}
    }
    server_info = get_server_info()
    wizard(params,server_info)
    os.makedirs("./tbench_results/{tn}".format(tn=params["test_name"]),exist_ok=True)

    if params["rsync_settings"]["rsync_flag"]:
        params["script_commands"]["rsync"] = command = "rsync -avh ./tbench_results/{tn} {u}@{h}:{d}/".format(tn=params["test_name"],u=params["rsync_settings"]["rsync_user"],h=params["rsync_settings"]["rsync_host"],d=params["rsync_settings"]["rsync_dir"])

    params["script_commands"]["csv_converter"] = "./csv_converter -i output/{tn}.csv -o chart_csv/{tn}.csv".format(tn=params["test_name"])
    
    if params["graph_flag"]:
        params["script_commands"]["make_graph"] = "./make_graph -i chart_csv/{tn}.csv -o graphs/{tn}.png -t '{gt}'".format(tn=params["test_name"],gt=params["graph_settings"]["title"])

    with open('./tbench_results/{tn}/{tn}.json'.format(tn=params["test_name"]), 'w') as outfile:
        outfile.write(json.dumps(params))
        
    if params["rsync_settings"]["rsync_flag"]:
        run_rsync_test(params)

    print(json.dumps(params,indent=4))
    rv = run_tplot(params)
    if rv == 0:
        rv = run_csv_converter(params)
    if rv == 0 and params["graph_flag"]:
        rv = run_make_graph(params)
    
    if rv == 0 and params["rsync_settings"]["rsync_flag"]:
        rv = run_rsync(params)

    if rv == 0:
        print("thermal benchmark complete.")
        print("raw data is here: output/{tn}.csv".format(tn=params["test_name"]))
        if params["graph_flag"]:
            print("graph is here: graphs/{tn}.png".format(tn=params["test_name"]))


if __name__ == "__main__":
    main()