#!/usr/bin/env python3

import time
import os
import subprocess
import argparse
import json

import scripts.swf_to_batsim_split_by_user as split_user
from scripts.util import *


def prepare_input_data(expe_num, start_date):
    """Cut the original trace to extract 72h starting from this start date"""
    end_date = start_date + 72*3600
    to_keep = f"submit_time >= {start_date} and submit_time <= {end_date}"

    if not os.path.exists(f'{WL_DIR}/expe{expe_num}'):
        os.makedirs(f'{WL_DIR}/expe{expe_num}')
    split_user.generate_workload(
        input_swf=f'{WL_DIR}/MC_selection_article.swf',
        output_folder=f'{WL_DIR}/expe{expe_num}',
        keep_only=to_keep,
        job_grain=10,
        job_walltime_factor=8)


def run_expe(expe_num, user_category, window_size, clean_log, nb_cores=16):
    """Run batmen with given behavior and demand response window.
    Expe_num should be a small integer (eg < 100)"""
    # Useful vars and output folder
    EXPE_DIR = f"{ROOT_DIR}/out/expe{expe_num}/{user_category}_window{window_size}"
    create_dir_rec_if_needed(EXPE_DIR)
    create_dir_rec_if_needed(f"{EXPE_DIR}/cmd")
    EXPE_FILE = f"{EXPE_DIR}/cmd/robinfile.yaml"
    wl_folder = f'{WL_DIR}/expe{expe_num}'
    pf = f"{ROOT_DIR}/platform/average_metacentrum.xml"
    wl = f"{WL_DIR}/empty_workload.json"
    uf = f"{EXPE_DIR}/cmd/user_description_file.json"

    # Demand response window, from 16 to (16 + window_size) on day2
    dm_window = [(24+16)*3600, (int) ((24+16+window_size)*3600)]

    # User description file
    def user_description(user):
        return {
            "name": user,
            "category": user_category,
            "param": {"input_json": f"{wl_folder}/{user}.json"}
        }
    user_names = [user_file.split('.')[0] for user_file in os.listdir(wl_folder)]
    data = {}
    data["dm_window"] = dm_window
    data["log_user_stats"] = True
    data["log_folder"] = EXPE_DIR
    data["users"] = [user_description(user) for user in user_names]
    with open(uf, 'w') as user_description_file:
        json.dump(data, user_description_file)

    # Generate and run robin instance
    socket_batsim = f"tcp://localhost:280{expe_num:02d}"
    socket_batsched = f"tcp://*:280{expe_num:02d}"
    quiet_batsim, quiet_batmen = "", ""
    if clean_log:
        quiet_batsim = "-q "
        quiet_batmen = "--verbosity=silent "
    batcmd = gen_batsim_cmd(
        pf, wl, EXPE_DIR, 
        f"--socket-endpoint={socket_batsim} {quiet_batsim} --energy --enable-compute-sharing --enable-dynamic-jobs --acknowledge-dynamic-jobs --enable-profile-reuse")
    schedcmd = f"batmen {quiet_batmen} --socket-endpoint={socket_batsched} --nb_cores={nb_cores} -v bin_packing_energy --variant_options_filepath={uf}"
    instance = RobinInstance(output_dir=EXPE_DIR,
                            batcmd=batcmd,
                            schedcmd=schedcmd,
                            simulation_timeout=604800, ready_timeout=10,
                            success_timeout=3600, failure_timeout=5
                            )
    instance.to_file(EXPE_FILE)
    print(f"Run robin {EXPE_FILE}")
    ret = run_robin(EXPE_FILE)
    print(f"Robin {EXPE_FILE} finished")

    # Remove the log files that can quickly become heavy...
    # if clean_log:
    #     os.remove(f"{EXPE_DIR}/log/batsim.log")
    #     os.remove(f"{EXPE_DIR}/log/sched.err.log")
    #     os.remove(f"{EXPE_DIR}/log/sched.out.log")



def start_instance(expe_num, prepare_workload=True, start_date=0, clean_log=False, only_degrad_temp=False):
    # Prepare workload
    if prepare_workload:
        prepare_input_data(expe_num, start_date)

    # Create expe folder
    create_dir_rec_if_needed(f"{ROOT_DIR}/out/expe{expe_num}")

    # Run with Rigid behavior (the demand response window has no influence here)
    run_expe(expe_num=expe_num, 
            user_category="replay_user_rigid",
            window_size=1, clean_log=clean_log)

    # 4*2 = 8 expe
    if only_degrad_temp:
        # bhvs = ["dm_user_degrad_temp"]
        bhvs = []
    else:
        bhvs = ["dm_user_reconfig","dm_user_degrad_space", "dm_user_degrad_temp",
                     "dm_user_renounce","dm_user_delay"]
    for behavior in bhvs:
        for window_size in [1, 4]:
            run_expe(expe_num, behavior, window_size, clean_log=clean_log)


    ###### Output data treatment ######
    # Produce the utilisation viz?
    return expe_num

def main():
    parser = argparse.ArgumentParser(
        description='One expe instance. To launch for example with `oarsub -l walltime=2 "./1_one_instance arg1 arg2 arg3"`')
    parser.add_argument('expe_num', type=int, help='The expe ID')
    # parser.add_argument('start_date', type=int, 
    #                     help='Start of the 3-day window (in seconds since the start of the original trace)')
    args = parser.parse_args()
    
    start_instance(args.expe_num, prepare_workload=False, clean_log=True, only_degrad_temp=True)



if __name__ == "__main__":
    main()