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


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, str(default)).ask().strip()

    if result == "":
        return default
    elif result.isdigit():
        return int(result)
    else:
        print("Please input a number.")
        exit(1)


CORES_ALL = get_cpu_cores()


def quest_basic():
    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)
    events_record = questionary.text(
        "What's the event of perf-record? (Default: {cycles,instructions}:S)\n",
        "{cycles,instructions}:S",
    ).ask()

    freq_stat = ask_number(
        "What's count deltas of perf-stat? (Default: 1000 milliseconds)\n", 1000
    )
    events_stat = questionary.text(
        "What's the event of perf-stat?\n (Default: cycles,instructions,LLC-loads,LLC-load-misses,LLC-stores,LLC-store-misses,branch-load-misses,branch-loads)\n",
        "cycles,instructions,LLC-loads,LLC-load-misses,LLC-stores,LLC-store-misses,branch-load-misses,branch-loads",
    ).ask()

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

    return workspace, freq_record, events_record, freq_stat, events_stat, annotete


def write_title(file: TextIOWrapper):
    file.write(
        """
#!/bin/bash
# The script generated by PIPA-TREE is used to collect performance data. 
# Please check whether it meets expectations before running.
# ZJU SPAIL(https://github.com/ZJU-SPAIL) reserves all rights.

# 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"""
    )
