# SPDX-License-Identifier: MIT

"""
This module contains common utility functions and classes for various amd-debug-tools.
"""

import asyncio
import importlib.metadata
import logging
import os
import platform
import time
import struct
import subprocess
import sys
from ast import literal_eval
from datetime import date, timedelta


class Colors:
    """Colors for the terminal"""

    DEBUG = "\033[90m"
    HEADER = "\033[95m"
    OK = "\033[94m"
    WARNING = "\033[32m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    UNDERLINE = "\033[4m"


def read_file(fn) -> str:
    """Read a file and return the contents"""
    with open(fn, "r", encoding="utf-8") as r:
        return r.read().strip()


def compare_file(fn, expect) -> bool:
    """Compare a file to an expected string"""
    return read_file(fn) == expect


def get_group_color(group) -> str:
    """Get the color for a group"""
    if group == "🚦":
        color = Colors.WARNING
    elif group == "🗣️":
        color = Colors.HEADER
    elif any(mk in group for mk in ["💯", "🚫"]):
        color = Colors.UNDERLINE
    elif any(mk in group for mk in ["🦟", "🖴"]):
        color = Colors.DEBUG
    elif any(mk in group for mk in ["❌", "👀"]):
        color = Colors.FAIL
    elif any(mk in group for mk in ["✅", "🔋", "🐧", "💻", "○", "💤", "🥱"]):
        color = Colors.OK
    else:
        color = group
    return color


def print_color(message, group) -> None:
    """Print a message with a color"""
    prefix = f"{group} "
    suffix = Colors.ENDC
    color = get_group_color(group)
    if color == group:
        prefix = ""
    log_txt = f"{prefix}{message}".strip()
    if any(c in color for c in [Colors.OK, Colors.HEADER, Colors.UNDERLINE]):
        logging.info(log_txt)
    elif color == Colors.WARNING:
        logging.warning(log_txt)
    elif color == Colors.FAIL:
        logging.error(log_txt)
    else:
        logging.debug(log_txt)
    if "TERM" in os.environ and os.environ["TERM"] == "dumb":
        suffix = ""
        color = ""
    print(f"{prefix}{color}{message}{suffix}")


def colorize_choices(choices, default) -> str:
    """Output a list of choices with colors, where the default is highlighted"""
    if default not in choices:
        raise ValueError(f"Default choice '{default}' not in choices")
    choices = [c for c in choices if c != default]
    choices = [f"{Colors.OK}{default}{Colors.ENDC}"] + choices
    return ", ".join(choices)


def fatal_error(message):
    """Prints a fatal error message and exits"""
    _configure_log(None)
    print_color(message, "👀")
    sys.exit(1)


def apply_prefix_wrapper(header, message):
    """Apply a prefix to wrap a newline delimitted message"""
    s = f"{header.strip()}\n"
    lines = message.strip().split("\n")
    for i, line in enumerate(lines):
        line = line.strip()
        if not line:
            continue
        if i == len(lines) - 1:
            s += f"└─ {line}\n"
            continue
        s += f"│ {line}\n"
    return s


def show_log_info():
    """Show log information"""
    logger = logging.getLogger()
    for handler in logger.handlers:
        if isinstance(handler, logging.FileHandler):
            filename = handler.baseFilename
            if filename != "/dev/null":
                print(f"Debug logs are saved to: {filename}")


def _configure_log(prefix) -> str:
    """Configure logging for the tool"""
    if len(logging.root.handlers) > 0:
        return

    if prefix:
        user = os.environ.get("SUDO_USER")
        home = os.path.expanduser(f"~{user if user else ''}")
        path = os.environ.get("XDG_DATA_HOME") or os.path.join(
            home, ".local", "share", "amd-debug-tools"
        )
        os.makedirs(path, exist_ok=True)
        log = os.path.join(
            path,
            f"{prefix}-{date.today()}.txt",
        )
        if not os.path.exists(log):
            with open(log, "w", encoding="utf-8") as f:
                f.write("")
            if "SUDO_UID" in os.environ:
                os.chown(path, int(os.environ["SUDO_UID"]), int(os.environ["SUDO_GID"]))
                os.chown(log, int(os.environ["SUDO_UID"]), int(os.environ["SUDO_GID"]))
        level = logging.DEBUG
    else:
        log = "/dev/null"
        level = logging.WARNING
    # for saving a log file for analysis
    logging.basicConfig(
        format="%(asctime)s %(levelname)s:\t%(message)s",
        filename=log,
        level=level,
    )
    return log


def check_lockdown():
    """Check if the system is in lockdown"""
    fn = os.path.join("/", "sys", "kernel", "security", "lockdown")
    if not os.path.exists(fn):
        return False
    lockdown = read_file(fn)
    if lockdown.split()[0] != "[none]":
        return lockdown
    return False


def print_temporary_message(msg) -> int:
    """Print a temporary message to the console"""
    print(msg, end="\r", flush=True)
    return len(msg)


def clear_temporary_message(msg_len) -> None:
    """Clear a temporary message from the console"""
    print(" " * msg_len, end="\r")


def run_countdown(action, t) -> bool:
    """Run a countdown timer"""
    msg = ""
    if t < 0:
        return False
    if t == 0:
        return True
    while t > 0:
        msg = f"{action} in {timedelta(seconds=t)}"
        print_temporary_message(msg)
        time.sleep(1)
        t -= 1
    clear_temporary_message(len(msg))
    return True


def get_distro() -> str:
    """Get the distribution name"""
    distro = "unknown"
    if os.path.exists("/etc/os-release"):
        with open("/etc/os-release", "r", encoding="utf-8") as f:
            for line in f:
                if line.startswith("ID="):
                    return line.split("=")[1].strip().strip('"')
    if os.path.exists("/etc/arch-release"):
        return "arch"
    elif os.path.exists("/etc/fedora-release"):
        return "fedora"
    elif os.path.exists("/etc/debian_version"):
        return "debian"

    return distro


def get_pretty_distro() -> str:
    """Get the pretty distribution name"""
    distro = "Unknown"
    if os.path.exists("/etc/os-release"):
        with open("/etc/os-release", "r", encoding="utf-8") as f:
            for line in f:
                if line.startswith("PRETTY_NAME="):
                    distro = line.split("=")[1].strip().strip('"')
                    break
    return distro


def bytes_to_gb(bytes_value):
    """Convert bytes to GB"""
    return bytes_value * 4096 / (1024 * 1024 * 1024)


def gb_to_pages(gb_value):
    """Convert GB into bytes"""
    return int(gb_value * (1024 * 1024 * 1024) / 4096)


def reboot():
    """Reboot the system"""

    async def reboot_dbus_fast():
        """Reboot using dbus-fast"""
        try:
            from dbus_fast.aio import (  # pylint: disable=import-outside-toplevel
                MessageBus,
            )
            from dbus_fast import BusType  # pylint: disable=import-outside-toplevel

            bus = await MessageBus(bus_type=BusType.SYSTEM).connect()
            introspection = await bus.introspect(
                "org.freedesktop.login1", "/org/freedesktop/login1"
            )
            proxy_obj = bus.get_proxy_object(
                "org.freedesktop.login1", "/org/freedesktop/login1", introspection
            )
            interface = proxy_obj.get_interface("org.freedesktop.login1.Manager")
            await interface.call_reboot(True)

        except ImportError:
            return False
        return True

    def reboot_dbus():
        """Reboot using python-dbus"""
        try:
            import dbus  # pylint: disable=import-outside-toplevel

            bus = dbus.SystemBus()
            obj = bus.get_object("org.freedesktop.login1", "/org/freedesktop/login1")
            intf = dbus.Interface(obj, "org.freedesktop.login1.Manager")
            intf.Reboot(True)
        except ImportError:
            return False
        return True

    loop = asyncio.get_event_loop()
    result = loop.run_until_complete(reboot_dbus_fast())
    if not result:
        return reboot_dbus()

    return True


def get_system_mem():
    """Get the total system memory in GB using /proc/meminfo"""
    with open(os.path.join("/", "proc", "meminfo"), "r", encoding="utf-8") as f:
        for line in f:
            if line.startswith("MemTotal:"):
                # MemTotal line format: "MemTotal:       16384516 kB"
                # Extract the number and convert from kB to GB
                mem_kb = int(line.split()[1])
                return mem_kb / (1024 * 1024)
    raise ValueError("Could not find MemTotal in /proc/meminfo")


def is_root() -> bool:
    """Check if the user is root"""
    return os.geteuid() == 0


def BIT(num):  # pylint: disable=invalid-name
    """Return a bit shifted value"""
    return 1 << num


def get_log_priority(num):
    """Maps an integer debug level to a priority type"""
    if num:
        try:
            num = int(num)
        except ValueError:
            return num
        if num == 7:
            return "🦟"
        elif num == 4:
            return "🚦"
        elif num <= 3:
            return "❌"
    return "○"


def minimum_kernel(major, minor) -> bool:
    """Checks if the kernel version is at least major.minor"""
    ver = platform.uname().release.split(".")
    kmajor = int(ver[0])
    kminor = int(ver[1])
    if kmajor > int(major):
        return True
    if kmajor < int(major):
        return False
    return kminor >= int(minor)


def systemd_in_use() -> bool:
    """Check if systemd is in use"""
    # Check if /proc/1/comm exists and read its contents
    init_daemon = read_file("/proc/1/comm")
    return init_daemon == "systemd"


def get_property_pyudev(properties, key, fallback=""):
    """Get a property from a udev device"""
    try:
        return properties.get(key, fallback)
    except UnicodeDecodeError:
        return ""


def find_ip_version(base_path, kind, hw_ver) -> bool:
    """Determine if an IP version is present on the system"""
    b = os.path.join(base_path, "ip_discovery", "die", "0", kind, "0")
    for key, expected_value in hw_ver.items():
        p = os.path.join(b, key)
        if not os.path.exists(p):
            return False
        v = int(read_file(p))
        if v != expected_value:
            return False
    return True


def read_msr(msr, cpu):
    """Read a Model-Specific Register (MSR) value from the CPU."""
    p = f"/dev/cpu/{cpu}/msr"
    if not os.path.exists(p) and is_root():
        os.system("modprobe msr")
    try:
        f = os.open(p, os.O_RDONLY)
    except OSError as exc:
        raise PermissionError from exc
    try:
        os.lseek(f, msr, os.SEEK_SET)
        val = struct.unpack("Q", os.read(f, 8))[0]
    except OSError as exc:
        raise PermissionError from exc
    finally:
        os.close(f)
    return val


def relaunch_sudo() -> None:
    """Relaunch the script with sudo if not already running as root"""
    if not is_root():
        logging.debug("Relaunching with sudo")
        env_vars = []
        for var in ["DISPLAY", "WAYLAND_DISPLAY", "XAUTHORITY", "XDG_RUNTIME_DIR",
                    "DBUS_SESSION_BUS_ADDRESS", "XDG_SESSION_TYPE"]:
            value = os.environ.get(var)
            if value:
                env_vars.append(f"{var}={value}")

        if env_vars:
            sudo_cmd = ["sudo"] + [f"--preserve-env={var.split('=')[0]}" for var in env_vars] + sys.argv
        else:
            sudo_cmd = ["sudo"] + sys.argv

        os.execvp("sudo", sudo_cmd)


def running_ssh():
    return "SSH_CLIENT" in os.environ or "SSH_TTY" in os.environ


def convert_string_to_bool(str_value) -> bool:
    """convert a string to a boolean value"""
    try:
        value = literal_eval(str_value)
    except (SyntaxError, ValueError):
        value = None
        sys.exit(f"Invalid entry: {str_value}")
    return bool(value)


def _git_describe() -> str:
    """Get the git description of the current commit"""
    try:
        result = subprocess.check_output(
            ["git", "log", "-1", '--format=commit %h ("%s")'],
            cwd=os.path.dirname(__file__),
            text=True,
            stderr=subprocess.DEVNULL,
        )
        return result.strip()
    except subprocess.CalledProcessError:
        return None
    except FileNotFoundError:
        return None


def version() -> str:
    """Get version of the tool"""
    ver = "unknown"
    try:
        ver = importlib.metadata.version("amd-debug-tools")
    except importlib.metadata.PackageNotFoundError:
        pass
    describe = _git_describe()
    if describe:
        ver = f"{ver} [{describe}]"
    return ver


class AmdTool:
    """Base class for AMD tools"""

    def __init__(self, prefix):
        self.log = _configure_log(prefix)
        logging.debug("command: %s (module: %s)", sys.argv, type(self).__name__)
        logging.debug("Version: %s", version())
        if os.uname().sysname != "Linux":
            raise RuntimeError("This tool only runs on Linux")
