import questionary
import os
from pipa.export_config.cpu_config import get_cpu_cores
from rich import print


def ask_number(question: str, default: int) -> int:
    """
    Asks the user to input a number based on the given question and default value.

    Args:
        question (str): The question to ask the user.
        default (int): The default value to return if the user doesn't input anything.

    Returns:
        int: The number inputted by the user or the default value.
    """
    result = questionary.text(question).ask().strip()
    if result == "":
        return default
    elif result.isdigit():
        return int(result)
    else:
        print("Please input a number.")
        exit(1)


cores_available = get_cpu_cores()


def generate():
    """
    Generates a shell script for collecting performance data based on user inputs.
    """
    workspace = questionary.text(
        "Where do you want to store your data? (Default: ./)\n"
    ).ask()

    if workspace == "":
        workspace = "./"

    if not os.path.exists(workspace):
        os.makedirs(workspace)

    freq_record = ask_number("What's the frequency of perf-record? (Default: 99)\n", 99)

    freq_stat = ask_number("What's the frequency of perf-stat? (Default: 99)\n", 99)

    annotete = questionary.select(
        "Whether to use perf-annotate?\n", choices=["Yes", "No"]
    ).ask()
    annotete = True if annotete == "Yes" else False

    taskset = questionary.select(
        "Whether to use taskset?\n", choices=["Yes", "No"]
    ).ask()

    if taskset == "Yes":
        cores = questionary.text(
            "Which cores do you want to use? (Default: 120-140)\n"
        ).ask()
        if cores == "":
            cores = "120-140"

        if cores.isdigit():
            core_list = cores.strip()
        elif cores.split("-").__len__() != 2:
            print("Please input a valid range, split by '-'.")
            exit(1)
        else:
            l, r = cores.split("-")

            l, r = l.strip(), r.strip()
            if not l.isdigit() or not r.isdigit():
                print("Please input a valid range, non-digit char detected.")
                exit(1)
            l, r = int(l), int(r)
            if l < cores_available[0] or r > cores_available[-1] or l >= r:
                print("Please input a valid range.")
                exit(1)
            core_list = ",".join([str(i) for i in list(range(l, r + 1))])

    command = questionary.text("What's the command of workload?\n").ask()

    if command == "":
        print("Please input a command to run workload.")
        exit(1)

    if taskset == "Yes":
        command = f"/usr/bin/taskset -c {core_list} {command}"

    with open(workspace + "/pipa-run.sh", "w") as f:
        f.write("#!/bin/bash\n")
        f.write(
            "# The script generated by PIPA-TREE is used to collect performance data. \n"
        )
        f.write("# Please check whether it meets expectations before running.\n")
        f.write("# ZJU SPAIL(https://github.com/ZJU-SPAIL) reserves all rights.\n\n")

        f.write(
            """
# Check if sar and perf are available
if ! command -v sar &> /dev/null; then
    echo "sar command not found. Please install sar."
    exit 1
fi

if ! command -v perf &> /dev/null; then
    echo "perf command not found. Please install perf."
    exit 1
fi\n\n"""
        )

        f.write("WORKSPACE=" + workspace + "\n")
        f.write("mkdir -p $WORKSPACE\n\n")

        f.write(
            "perf record -e '{cycles,instructions}:S' -a -F"
            + f" {freq_record} -o $WORKSPACE/perf.data {command}\n"
        )
        f.write("perf script -i $WORKSPACE/perf.data > $WORKSPACE/perf.script\n")
        f.write("perf report -i $WORKSPACE/perf.data > $WORKSPACE/perf.report\n\n")

        f.write("sar -o $WORKSPACE/sar.dat 1 >/dev/null 2>&1 &\n")
        f.write("sar_pid=$!\n")
        f.write(
            f"perf stat -e cycles,instructions -C {cores_available[0]}-{cores_available[-1]} -A -x , -I {freq_stat} -o $WORKSPACE/perf-stat.csv {command}\n"
        )
        f.write("kill -9 $sar_pid\n")
        f.write("sar -A -f $WORKSPACE/sar.dat >$WORKSPACE/sar.txt\n\n")

        if annotete:
            f.write(
                "perf annotate -i $WORKSPACE/perf.data > $WORKSPACE/perf.annotate\n\n"
            )

        print("Shell script generated successfully.")
        print("Please check the script in " + workspace + "/pipa-run.sh")


if __name__ == "__main__":
    generate()
