# SPDX-License-Identifier: GPL-2.0-only
# Copyright (C) 2025 John B. Wyatt IV

import sys
from typing import List
from tuna import utils

cpupower_required_kernel = "6.12"
have_cpupower = None

try:
    import raw_pylibcpupower as lcpw
    lcpw.cpufreq_get_available_frequencies(0)
    have_cpupower = True
except ImportError:
    lcpw = None
    have_cpupower = False

if have_cpupower:
    class Cpupower:
        """The Cpupower class allows you to query and change the power states of the
        cpu.

        You may query or change the cpus all at once or a list of the cpus provided to the constructor's cpulist argument.

        The bindings must be detected on the $PYTHONPATH variable.

        You must use have_cpupower variable to determine if the bindings were
        detected in your code."""

        LCPW_ERROR_TWO_CASE = 1 # enum for common error messages
        LCPW_ERROR_THREE_CASE = 2

        def __init__(self, cpu_list=None):
            if cpu_list and not cpu_list == []:
                self.__cpu_list = cpu_list
            else:
                self.__cpu_list = utils.get_all_cpu_list()

        def handle_common_lcpw_errors(self, e, error_type, idle_name):
            if e == 0:
                pass
            elif e == -1:
                print(f"Idlestate {idle_name} not available", file=sys.stderr)
            elif e == -2:
                print("Disabling is not supported by the kernel", file=sys.stderr)
            elif e == -3:
                if error_type == Cpupower.LCPW_ERROR_THREE_CASE:
                    print("No write access to disable/enable C-states: try using sudo", file=sys.stderr)
                else:
                    print(f"Not documented: {e}", file=sys.stderr)
            else:
                print(f"Not documented: {e}", file=sys.stderr)

        def get_idle_states(self, cpu):
            """
            Get the c-states of a cpu.

            You can capture the return values with:
            states_list, states_amt = get_idle_states()

            Returns
                List[String]: list of cstates
                Int: amt of cstates
            """
            ret = []
            for cstate in range(lcpw.cpuidle_state_count(cpu)):
                ret.append(lcpw.cpuidle_state_name(cpu,cstate))
            return ret, lcpw.cpuidle_state_count(cpu)

        def get_idle_info(self, cpu):
            idle_states, idle_states_amt = self.get_idle_states(cpu)
            idle_states_list = []
            for idle_state, idle_state_name in enumerate(idle_states):
                idle_states_list.append(
                    {
                        "CPU ID": cpu,
                        "Idle State Name": idle_state_name,
                        "Flags/Description": lcpw.cpuidle_state_desc(cpu, idle_state),
                        "Latency": lcpw.cpuidle_state_latency(cpu, idle_state),
                        "Usage": lcpw.cpuidle_state_usage(cpu, idle_state),
                        "Duration": lcpw.cpuidle_state_time(cpu, idle_state)
                    }
                )
            idle_info = {
                "CPUidle-driver": lcpw.cpuidle_get_driver(),
                "CPUidle-governor": lcpw.cpuidle_get_governor(),
                "idle-states-count": idle_states_amt,
                "available-idle-states": idle_states,
                "cpu-states": idle_states_list
            }
            return idle_info

        def print_idle_info(self, cpu_list):
            for cpu in cpu_list:
                idle_info = self.get_idle_info(cpu)
                print(
f"""CPUidle driver: {idle_info["CPUidle-driver"]}
CPUidle governor: {idle_info["CPUidle-governor"]}
analyzing CPU {cpu}

Number of idle states: {idle_info["idle-states-count"]}
Available idle states: {idle_info["available-idle-states"]}""")
                for state in idle_info["cpu-states"]:
                    print(
f"""{state["Idle State Name"]}
Flags/Description: {state["Flags/Description"]}
Latency: {state["Latency"]}
Usage: {state["Usage"]}
Duration: {state["Duration"]}""")

        def idle_set_handler(self, args) -> int:
            if args.idle_state_disabled_status is not None:
                cstate_index = args.idle_state_disabled_status
                cstate_list, cstate_amt = self.get_idle_states(self.__cpu_list[0]) # Assuming all cpus have the same idle state
                if cstate_index < 0 or cstate_index >= cstate_amt:
                    print(f"Invalid idle state range. Total for this cpu is {cstate_amt}", file=sys.stderr)
                    return 1
                cstate_name = cstate_list[cstate_index]
                ret = self.is_disabled_idle_state(cstate_index)
                for i,e in enumerate(ret):
                    if e == 1:
                        print(f"CPU: {self.__cpu_list[i]} Idle state \"{cstate_name}\" is disabled.")
                    elif e == 0:
                        print(f"CPU: {self.__cpu_list[i]} Idle state \"{cstate_name}\" is enabled.")
                    else:
                        self.handle_common_lcpw_errors(e, self.LCPW_ERROR_TWO_CASE, cstate_name)
            elif args.idle_info is not None:
                self.print_idle_info(self.__cpu_list)
                return 0
            elif args.disable_idle_state is not None:
                cstate_index = args.disable_idle_state
                cstate_list, cstate_amt = self.get_idle_states(self.__cpu_list[0]) # Assuming all cpus have the same idle state
                if cstate_index < 0 or cstate_index >= cstate_amt:
                    print(f"Invalid idle state range. Total for this cpu is {cstate_amt}")
                    return 1
                cstate_name = cstate_list[cstate_index]
                ret = self.disable_idle_state(cstate_index, 1)
                for e in ret:
                    self.handle_common_lcpw_errors(e, self.LCPW_ERROR_THREE_CASE, cstate_name)
            elif args.enable_idle_state is not None:
                cstate_index = args.enable_idle_state
                cstate_list, cstate_amt = self.get_idle_states(self.__cpu_list[0]) # Assuming all cpus have the same idle state
                if cstate_index < 0 or cstate_index >= cstate_amt:
                    print(f"Invalid idle state range. Total for this cpu is {cstate_amt}")
                    return 1
                cstate_name = cstate_list[cstate_index]
                ret = self.disable_idle_state(cstate_index, 0)
                for e in ret:
                    self.handle_common_lcpw_errors(e, self.LCPW_ERROR_THREE_CASE, cstate_name)
            return 0

        def disable_idle_state(self, state, disabled) -> List[int]:
            """
            Disable or enable an idle state using the object's stored list of cpus.

            Args:
                state (int): The cpu idle state index to disable or enable as an int starting from 0.
                disabled (int): set to 1 to disable or 0 to enable. Less than 0 is an error.
            """
            ret = []
            for cpu in self.__cpu_list:
                ret.append(lcpw.cpuidle_state_disable(cpu, state, disabled))
            return ret

        def is_disabled_idle_state(self, state) -> List[int]:
            """
            Query the idle state.

            Args:
                state: The cpu idle state. 1 is disabled, 0 is enabled. Less than 0 is an error.
            """
            ret = []
            for cpu in self.__cpu_list:
                ret.append(lcpw.cpuidle_is_state_disabled(cpu, state))
            return ret
