#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021-2025 The WfCommons Team.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

import os
import pathlib
import subprocess
import time
import sys
import signal
import queue
import argparse
import re
import json
import logging
import pandas as pd

from io import StringIO
from filelock import FileLock
from pathos.helpers import mp as multiprocessing
from typing import List, Optional


# Configure logging
logging.basicConfig(
    level=logging.INFO,  # Change this to control the verbosity
    format="[WfBench][%(asctime)s][%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
    handlers=[logging.StreamHandler()]
)


this_dir = pathlib.Path(__file__).resolve().parent


def log_info(msg: str):
    """
    Log an info message to stderr

    :param msg:
    :type msg: str
    """
    logging.info(msg)

def log_debug(msg: str):
    """
    Log a debug message to stderr

    :param msg:
    :type msg: str
    """
    logging.debug(msg)

def log_error(msg: str):
    """
    Log a error message to stderr

    :param msg:
    :type msg: str
    """
    logging.error(msg)


def lock_core(path_locked: pathlib.Path,
              path_cores: pathlib.Path) -> int:
    """
    Lock cores in use.

    :param path_locked:
    :type path_locked: pathlib.Path
    :param path_cores:
    :type path_cores: pathlib.Path

    :return:
    :rtype: int
    """
    all_cores = set(range(os.cpu_count()))
    path_locked.touch(exist_ok=True)
    path_cores.touch(exist_ok=True)

    while True:
        with FileLock(path_locked) as lock:
            try:
                lock.acquire()
                taken_cores = {
                    int(line) for line in path_cores.read_text().splitlines() if line.strip()
                }
                available = all_cores - taken_cores
                if available:
                    core = available.pop()
                    taken_cores.add(core)
                    path_cores.write_text("\n".join(map(str, taken_cores)))
                    return core

                log_debug("All Cores are taken")
            finally:
                lock.release()
        time.sleep(1)


def unlock_core(path_locked: pathlib.Path,
                path_cores: pathlib.Path,
                core: int) -> None:
    """
    Unlock cores after execution is done.

    :param path_locked:
    :type path_locked: pathlib.Path
    :param path_cores:
    :type path_cores: pathlib.Path
    :param core:
    :type core: int
    """
    with FileLock(path_locked) as lock:
        lock.acquire()
        try:
            taken_cores = {
                int(line) for line in path_cores.read_text().splitlines()
                if int(line) != core
            }
            path_cores.write_text("\n".join(map(str, taken_cores)))
        finally:
            lock.release()

def monitor_progress(proc, cpu_queue):
    """Monitor progress from the CPU benchmark process."""
    for line in iter(proc.stdout.readline, ""):  # No decode needed
        line = line.strip()
        if line.startswith("Progress:"):
            try:
                progress = float(line.split()[1].strip('%'))
                cpu_queue.put(progress)
            except (ValueError, IndexError):
                pass

def cpu_mem_benchmark(cpu_queue: multiprocessing.Queue,
                      cpu_threads: Optional[int] = 5,
                      mem_threads: Optional[int] = 5,
                      cpu_work: Optional[int] = 100,
                      core: Optional[int] = None,
                      total_mem: Optional[int] = None) -> List:
    """
    Run CPU and memory benchmark.

    :param cpu_queue: Queue to push CPU benchmark progress as a float.
    :type cpu_queue: multiprocessing.Queue
    :param cpu_threads: Number of threads for CPU benchmark.
    :type cpu_threads: Optional[int]
    :param mem_threads: Number of threads for memory benchmark.
    :type mem_threads: Optional[int]
    :param cpu_work: Total work units for CPU benchmark.
    :type cpu_work: Optional[int]
    :param core: Core to pin the benchmark processes to.
    :type core: Optional[int]
    :param total_mem: Total memory to use for memory benchmark.
    :type total_mem: Optional[float]

    :return: Lists of CPU and memory subprocesses.
    :rtype: List
    """
    total_mem = f"{total_mem}B" if total_mem else f"{100.0 / os.cpu_count()}%"
    cpu_work_per_thread = int(cpu_work / cpu_threads)

    cpu_procs = []
    mem_procs = []
    cpu_prog = [f"{this_dir.joinpath('cpu-benchmark')}", f"{cpu_work_per_thread}"]
    mem_prog = ["stress-ng", "--vm", f"{mem_threads}",
                "--vm-bytes", f"{total_mem}", "--vm-keep"]

    for i in range(cpu_threads):
        cpu_proc = subprocess.Popen(cpu_prog, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

        # NOTE: might be a good idea to use psutil to set the affinity (works across platforms)
        if core:
            os.sched_setaffinity(cpu_proc.pid, {core})
        cpu_procs.append(cpu_proc)

        # Start a thread to monitor the progress of each CPU benchmark process
        monitor_thread = multiprocessing.Process(target=monitor_progress, args=(cpu_proc, cpu_queue))
        monitor_thread.start()

    if mem_threads > 0:
        # NOTE: add a check to use creationflags=subprocess.CREATE_NEW_PROCESS_GROUP for Windows
        mem_proc = subprocess.Popen(mem_prog, preexec_fn=os.setsid)
        if core:
            os.sched_setaffinity(mem_proc.pid, {core})
        mem_procs.append(mem_proc)

    return cpu_procs, mem_procs


def io_read_benchmark_user_input_data_size(inputs,
                                           rundir=None,
                                           memory_limit=None):
    if memory_limit is None:
        memory_limit = -1
    memory_limit = int(memory_limit)
    log_debug("Starting IO Read Benchmark...")
    for file, size in inputs.items():
        with open(rundir.joinpath(file), "rb") as fp:
            log_debug(f"Reading '{file}'")
            chunk_size = min(size, memory_limit)
            while fp.read(chunk_size):
                pass
    log_debug("Completed IO Read Benchmark!")


def io_write_benchmark_user_input_data_size(outputs,
                                            rundir=None,
                                            memory_limit=None):
    if memory_limit is None:
        memory_limit = sys.maxsize
    memory_limit = int(memory_limit)
    for file_name, file_size in outputs.items():
        log_debug(f"Writing output file '{file_name}'")
        file_size_todo = file_size
        while file_size_todo > 0:
            with open(rundir.joinpath(file_name), "ab") as fp:
                chunk_size = min(file_size_todo, memory_limit)
                file_size_todo -= fp.write(os.urandom(int(chunk_size)))


def io_alternate(inputs, outputs, cpu_queue: multiprocessing.Queue, memory_limit=None, rundir=None, event=None):
    """Alternate between reading and writing to a file, ensuring read only happens after write."""

    if memory_limit is None:
        memory_limit = 10 * 1024 * 1024  # sys.maxsize
    memory_limit = int(memory_limit)

    # queue will have messages in the form (cpu_percent_completed)
    # Get the last message and trash the rest

    # Create empty files
    for name in outputs:
        open(rundir.joinpath(name), "wb").close()

    io_completed = 0
    bytes_read = {
        name: 0
        for name in inputs
    }
    bytes_written = {
        name: 0
        for name in outputs
    }

    # get size of inputs
    inputs = {
        name: os.path.getsize(rundir.joinpath(name))
        for name in inputs
    }

    while io_completed < 100:
        cpu_percent = max(io_completed, cpu_queue.get())
        while True: # Get the last message
            try:
                cpu_percent = max(io_completed, cpu_queue.get_nowait())
            except queue.Empty:
                break

        log_debug(f"CPU Percent: {cpu_percent}")
        if cpu_percent:
            bytes_to_read = {
                name: int(size * (cpu_percent / 100) - bytes_read[name])
                for name, size in inputs.items()
            }
            bytes_to_write = {
                name: int(size * (cpu_percent / 100) - bytes_written[name])
                for name, size in outputs.items()
            }
            io_read_benchmark_user_input_data_size(bytes_to_read, rundir, memory_limit=memory_limit)
            io_write_benchmark_user_input_data_size(bytes_to_write, rundir, memory_limit=memory_limit)

            bytes_read = {
                name: bytes_read[name] + bytes_to_read[name]
                for name in bytes_to_read
            }
            bytes_written = {
                name: bytes_written[name] + bytes_to_write[name]
                for name in bytes_to_write
            }

            log_debug(f"Bytes Read: {bytes_read}")
            log_debug(f"Bytes Written: {bytes_written}")

            io_completed = cpu_percent

            if io_completed >= 100:
                break

def get_available_gpus():
    proc = subprocess.Popen(["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, _ = proc.communicate()
    df = pd.read_csv(StringIO(stdout.decode("utf-8")), sep=" ")
    return df[df["utilization.gpu"] <= 5].index.to_list()


def gpu_benchmark(time: int = 100,
                  work: int = 100, 
                  device: int = 0): #work, device
    
    gpu_prog = [f"CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES={device} {this_dir.joinpath('./gpu_benchmark')} {work} {time}"]
    log_debug(f"Running GPU Benchmark: {gpu_prog}")
    subprocess.Popen(gpu_prog, shell=True)


def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser.add_argument("--name", default=None, required=True, help="Task name.")
    parser.add_argument("--rundir", help="Run directory.")
    parser.add_argument("--percent-cpu", default=0.5, type=float,
                        help="percentage related to the number of CPU threads.")
    parser.add_argument("--path-lock", default=None, help="Path to lock file.")
    parser.add_argument("--path-cores", default=None, help="Path to cores file.")
    parser.add_argument("--cpu-work", default=None, help="Amount of CPU work.")
    parser.add_argument("--gpu-work", default=None, help="Amount of GPU work.")
    parser.add_argument("--time", default=None, help="Time limit (in seconds) to complete the task (overrides CPU and GPU works).")
    parser.add_argument("--mem", type=float, default=None, help="Max amount (in MB) of memory consumption.")
    parser.add_argument("--output-files", help="output file names with sizes in bytes as a JSON dictionary "
                                               "(e.g., --output-files {\\\"file1\\\": 1024, \\\"file2\\\": 2048}).")
    parser.add_argument("--input-files", help="input files names as a JSON array "
                                              "(e.g., --input-files [\\\"file3\\\", \\\"file4\\\"]).")
    parser.add_argument("--debug", action="store_true", help="Enable debug messages.")
    parser.add_argument("--with-flowcept", action="store_true", default=False, help="Enable Flowcept monitoring.")
    parser.add_argument("--workflow_id", default=None, help="Id to group tasks in a workflow.")

    return parser
    

def begin_flowcept(args):
    log_info("Running with Flowcept.")
    from flowcept import Flowcept, FlowceptTask
    # TODO: parametrize to allow storing individual tasks
    f = Flowcept(workflow_id=args.workflow_id,
                 bundle_exec_id=args.workflow_id,
                 start_persistence=False, save_workflow=False)
    f.start()
    t = FlowceptTask(task_id=f"{args.workflow_id}_{args.name}", workflow_id=args.workflow_id, used={**args.__dict__})
    return f, t


def end_flowcept(flowcept, flowcept_task):
    flowcept_task.end()
    flowcept.stop()


def main():
    """Main program."""
    parser = get_parser()
    args = parser.parse_args()
    core = None

    if args.with_flowcept:
        flowcept, flowcept_task = begin_flowcept(args)

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    if args.rundir:
        rundir = pathlib.Path(args.rundir)
    else:
        rundir = pathlib.Path(os.getcwd())

    if args.path_lock and args.path_cores:
        path_locked = pathlib.Path(args.path_lock)
        path_cores = pathlib.Path(args.path_cores)
        core = lock_core(path_locked, path_cores)

    log_info(f"Starting {args.name} Benchmark")

    mem_bytes = args.mem * 1024 * 1024 if args.mem else None

    procs = []
    io_proc = None
    outputs_dict = {}

    cpu_queue = multiprocessing.Queue()

    log_debug(f"Working directory: {os.getcwd()}")

    # Deal with input/output files if any
    cleaned_input = "{}" if args.input_files is None else re.sub(r'\\+', '', args.input_files)
    cleaned_output = "{}" if args.output_files is None else re.sub(r'\\+', '', args.output_files)
    # print("CLEANED INPUT", cleaned_input)
    # print("CLEANED OUTPUT", cleaned_output)

    if cleaned_input or cleaned_output:
        log_debug("Starting IO benchmark...")

        # Attempt to parse the cleaned string
        try:
            outputs_dict = json.loads(cleaned_output)
        except json.JSONDecodeError as e:
            log_error(f"Failed to decode --output-files JSON string argument: {e}")
            sys.exit(1)

        try:
            inputs_array = json.loads(cleaned_input)
        except json.JSONDecodeError as e:
            log_error(f"Failed to decode --input-files JSON string argument: {e}")
            sys.exit(1)

        # print("OUTPUT", outputs_dict)
        # print("INPUTS", inputs_array)

        # Create a multiprocessing event that in the first run is set to True
        write_done_event = multiprocessing.Event() 
        # Set this to True to allow the first read to happen
        write_done_event.set()
        # Print the value of the event
        # print("Event Value:", write_done_event.is_set())

        io_proc = multiprocessing.Process(
            target=io_alternate,
            args=(inputs_array, outputs_dict, cpu_queue, mem_bytes, rundir, write_done_event)
        )
        io_proc.start()
        procs.append(io_proc)

    if args.gpu_work:
        log_info(f"Starting GPU Benchmark for {args.name}...")
        available_gpus = get_available_gpus() #checking for available GPUs

        if not available_gpus:
            log_error("No GPU available")
            sys.exit(1)
        else:
            device = available_gpus[0]
            log_debug(f"Running on GPU {device}")

            if args.time:
                log_debug(f" Time:{args.time}, Work:{args.gpu_work}, Device:{device}")
                gpu_benchmark(time=int(args.time), work=int(args.gpu_work), device=device)
            else:
                gpu_benchmark(work=int(args.gpu_work), device=device)
            
    if args.cpu_work:
        log_info(f"Starting CPU and Memory Benchmarks for {args.name}...")
        if core:
            log_debug(f"{args.name} acquired core {core}")

        mem_threads=int(10 - 10 * args.percent_cpu)
        cpu_procs, mem_procs = cpu_mem_benchmark(cpu_queue=cpu_queue, 
                                                 cpu_threads=int(10 * args.percent_cpu),
                                                 mem_threads=mem_threads,
                                                 cpu_work=sys.maxsize if args.time else int(args.cpu_work),
                                                 core=core,
                                                 total_mem=mem_bytes)
                    
        procs.extend(cpu_procs)
        if args.time:
            time.sleep(int(args.time))
            for proc in procs:
                if isinstance(proc, multiprocessing.Process):
                    if proc.is_alive():
                        proc.terminate()
                elif isinstance(proc, subprocess.Popen):
                    proc.terminate()
        else:
            for proc in procs:
                if isinstance(proc, subprocess.Popen):
                    proc.wait()
        if io_proc is not None and io_proc.is_alive():
            # io_proc.terminate()
            io_proc.join()

        for mem_proc in mem_procs:
            try:
                os.kill(mem_proc.pid, signal.SIGKILL)  # Force kill if SIGTERM fails
            except subprocess.TimeoutExpired:
                log_debug("Memory process did not terminate; force-killing.")
        # As a fallback, use pkill if any remaining instances are stuck
        subprocess.Popen(["pkill", "-f", "stress-ng"]).wait()

        log_debug("Completed CPU and Memory Benchmarks!")

    # NOTE: If you would like to run only IO add time.sleep(2)
    # Check if all procs are done, if not, kill them
    log_debug("Checking if all processes are done...")
    for proc in procs:
        if isinstance(proc, multiprocessing.Process):
            if proc.is_alive():
                proc.terminate()
                proc.join()
        if isinstance(proc, subprocess.Popen):
            proc.wait()

    if core:
        unlock_core(path_locked, path_cores, core)

    if args.with_flowcept:
        end_flowcept(flowcept, flowcept_task)

    log_info(f"Benchmark {args.name} completed!")

if __name__ == "__main__":
    main()
