# controller_async.py
import asyncio
import sys
import argparse
import json
import os

from datetime import datetime
from aioconsole import ainput

# Usage:

# source ~/venv/ve-srslog/bin/activate
# python3 log-parser-controller.py -e 194_co 192_cu 195_du

save_period = 30  # Interval to send 'save' commands to clientes [seconds]

# Send one command to one client
async def send_command_to_addr(addr_ip, addr_port, cmd):

    try:
        reader, writer = await asyncio.open_connection(addr_ip, addr_port)
        writer.write(cmd.encode())
        await writer.drain()
        response = await reader.read(100)
        print(f"- Response: {response.decode().strip()} :: {addr_ip}:{addr_port}")
        writer.close()
        await writer.wait_closed()
    except Exception as e:
        print(f"Error in {addr_ip}:{addr_port}: {e}")

# Send one command to a list of clients
# - dest = 'parser' / 'sim'
async def send_command_to_clients(cmd, dest):

    dic_port = {'parser': 'port-parser', 'sim': 'port-sim'}

    key = dic_port[dest]

    print("Sending '" + cmd + "' command to all '" + dest + "' clients...")

    for c in j_cfg_a:
        print("-", j_cfg_a[c]['ip'] + ':' + str(j_cfg_a[c][key]))

    tasks = [asyncio.create_task(send_command_to_addr(j_cfg_a[c]['ip'], j_cfg_a[c][key], cmd)) for c in j_cfg_a]
    await asyncio.gather(*tasks)
    print("... commands sent succesfully\n")

# Sleep for a given period or until a key is pressed
async def sleep_key(period: float) -> bool:

    async def read_until_q():
        while True:
            s = (await ainput("")).strip().lower()
            if s == "q":
                return True

    try:
        return await asyncio.wait_for(read_until_q(), timeout=period)

    except asyncio.TimeoutError:
        return False

# Send commands concurrently to all clients: save / enable / quit
async def main():

    while True:

        # 3GPP TS 23.501 – System Architecture for the 5G System (5GS)
        # The 5G System (5GS) comprises the 5G Core (5GC) and the NG-RAN (Next Generation-Radio Access Network)

        print("NOTE: After running the 'rfp' or 'rsp' commands press 'q' to quit the controller and stop all remote parsers\n")

        # rfp = run five (5GS) + parser / rsp = run simulator + parser / rs = run (only) simulator
        cmd = input("Enter command (run 5GS + parser [rfp] / run sim + parser [rsp] / run sim [rs] / quit [q]): ").strip().lower()

        print("")

        if cmd in ["quit", "q"]:
            break

        elif cmd == "enable":
            await send_command_to_clients(cmd,'parser')

        elif cmd == "save":

            cmd = "save dt=" + datetime.now().isoformat()
            await send_command_to_clients(cmd, 'parser')

        elif cmd == "rs":

            cmd = "sim dt=" + sim_log_ref + " path=" + sim_log_path + "<END>"
            await send_command_to_clients(cmd, 'sim')

        elif cmd == "rfp":

            await send_command_to_clients("enable", 'parser')

            if await sleep_key(save_period):
                await send_command_to_clients("quit", 'parser')
                print("... controller quit (after enable command)")
                sys.exit()

            while True:
                cmd = "save dt=" + datetime.now().isoformat()

                await send_command_to_clients(cmd, 'parser')

                if await sleep_key(save_period):
                    await send_command_to_clients("quit", 'parser')
                    print("... controller quit (after save command)")
                    sys.exit()

        elif cmd == "rsp":

            dt_now_sim = datetime.now()

            cmd = "sim dt=" + sim_log_ref + " path=" + sim_log_path + "<END>"
            await send_command_to_clients(cmd, 'sim')

            await send_command_to_clients("enable", 'parser')
            await asyncio.sleep(save_period)

            while True:
                cmd = "save dt=" + datetime.now().isoformat()
                await send_command_to_clients(cmd, 'parser')
                await sleep_key(save_period)

                dt_now = datetime.now()

                if  (dt_now - dt_now_sim).total_seconds() > sim_iter_period:
                    dt_now_sim = dt_now
                    cmd = "sim dt=" + sim_log_ref + " path=" + sim_log_path + "<END>"
                    await send_command_to_clients(cmd, 'sim')

        else:
            print("ERROR: command not supported")

if __name__ == "__main__":

    try:

        print("--------------------------")
        print("log-parser Controller 1.01")
        print("--------------------------")
        print("")

        user = os.getlogin()

        if sys.platform == 'linux':
            config_file_path = r'/home/' + user + '/config/srs-log/'
        else:
            config_file_path = r'C:/Zona/Profesional/O-RAN/Org/Config srsRAN/Deneb-4 - Core/home · epc · config/srs-log/'

        config_file_name = "config-parser-controller.json"

        print("- Configuration path:", config_file_path)
        print("- Configuration file:", config_file_name)
        print("")

        config_file = config_file_path + config_file_name

        print("Loading configuration file '" + config_file_name + "' ...")

        try:
            with open(config_file, 'r') as j_file:
                j_cfg = json.load(j_file)
        except FileNotFoundError:
            print("Error: configuration file '" + config_file_name + "' does not exist")
            sys.exit()

        print("... configuration file '" + config_file_name + "' loaded")
        print("")

        parser = argparse.ArgumentParser(description='log-parser-controller')
        parser.add_argument('-e', '--elements', nargs='+', help='Network elements to initiate: local_cu, local_du, 194_co, 192_cu, 192_du, 198_cu, 198_du, 193_du, 195_du', required=True)
        parser.add_argument('-d', '--dataset', help='Dataset used for simulations: 1, 2, 3, 4, ...', type=int, required=False, default=1)

        args = vars(parser.parse_args(args=None if sys.argv[1:] else ['--help']))

        if len(sys.argv) == 1:
            parser.print_help(sys.stderr)
            sys.exit(1)

        arg_elements = args['elements']
        arg_dataset = args['dataset']

        dataset = str(arg_dataset).rjust(2,'0')

        print("Elements: " + ', '.join(arg_elements))
        print("Dataset: " + dataset)
        print("")
        
        j_cfg_a_all = j_cfg["addreses"]

        j_cfg_a = {}

        for c in j_cfg_a_all:
            if j_cfg_a_all[c]['alias'] in arg_elements:
                j_cfg_a[c] = j_cfg_a_all[c]

        print("Selected network clients:")
        print("")

        print("log-parser-main      log-parser-sim       Alias    Machine")
        for c in j_cfg_a:
            print(j_cfg_a[c]['ip'] + ':' + str(j_cfg_a[c]['port-parser']) +
                  " | "
                  + j_cfg_a[c]['ip'] + ':' + str(j_cfg_a[c]['port-sim']) +
                  " | "
                  + j_cfg_a[c]['alias'] +
                  " | "
                  + c)
        print("")

        sim_log_path = j_cfg["sim_input_path_" + dataset]  # Path to log to simulate
        sim_log_ref = j_cfg["sim_ref_dt_log_" + dataset]  # Temporal reference to initiate the simulation
        sim_iter_period = j_cfg["sim_iterat_per_" + dataset]  # Interval between complete simulation iterations

        print("Simulation log path:", sim_log_path)
        print("Simulation log date-time reference:", sim_log_ref)
        print("Simulation iteration period:", sim_iter_period, "seconds")
        print("")

        asyncio.run(main())

    except KeyboardInterrupt:
        sys.stdout.write('\x1b[2K\n')
        print("\n... controller gracefully stopped")