#!/usr/bin/env python

# Copyright (C) 2017-2025 Alex Manuskin, Maor Veitsman
#
# 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 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
"""This module reads intel power measurements"""

from __future__ import annotations

import glob
import logging
import os
import re
from collections import namedtuple
from multiprocessing import cpu_count

from s_tui.helper_functions import cat
from s_tui.sources.msr import msr_available, read_msr

INTER_RAPL_DIR = "/sys/class/powercap/intel-rapl/"
AMD_ENERGY_DIR_GLOB = "/sys/devices/platform/amd_energy.0/hwmon/hwmon*/"
MICRO_JOULE_IN_JOULE = 1000000.0


UNIT_MSR = 0xC0010299
CORE_MSR = 0xC001029A
PACKAGE_MSR = 0xC001029B
ENERGY_UNIT_MASK = 0x1F00


RaplStats = namedtuple("rapl", ["label", "current", "max"])


class RaplReader:
    def __init__(self) -> None:
        basenames = glob.glob("/sys/class/powercap/intel-rapl:*/")
        self.basenames = sorted(set(basenames))

    def read_power(self) -> list[RaplStats]:
        """Read power stats and return dictionary"""

        pjoin = os.path.join
        ret = []
        for path in self.basenames:
            name = None
            try:
                name = cat(pjoin(path, "name"), fallback=None, binary=False)
            except (OSError, ValueError) as err:
                logging.warning("ignoring %r for file %r", (err, path), RuntimeWarning)
                continue
            if name:
                try:
                    current = cat(pjoin(path, "energy_uj"))
                    max_reading = 0.0
                    ret.append(RaplStats(name, float(current), max_reading))
                except (OSError, ValueError) as err:
                    logging.warning(
                        "ignoring %r for file %r", (err, path), RuntimeWarning
                    )
        return ret

    @staticmethod
    def available() -> bool:
        return os.path.exists("/sys/class/powercap/intel-rapl")


class AMDEnergyReader:
    def __init__(self) -> None:
        self.inputs = list(
            zip(
                (
                    cat(filename, binary=False)
                    for filename in sorted(
                        glob.glob(AMD_ENERGY_DIR_GLOB + "energy*_label")
                    )
                ),
                sorted(glob.glob(AMD_ENERGY_DIR_GLOB + "energy*_input")),
            )
        )

        # How many socket does the system have?
        socket_number = sum(1 for label, _ in self.inputs if "socket" in label)
        self.inputs.sort(key=lambda x: self.get_input_position(x[0], socket_number))

    @staticmethod
    def match_label(label: str) -> re.Match[str] | None:
        return re.search(r"E(core|socket)([0-9]+)", label)

    @staticmethod
    def get_input_position(label: str, socket_number: int) -> int:
        m = AMDEnergyReader.match_label(label)
        assert m is not None, f"Unexpected label format: {label}"
        num = int(m.group(2))
        if "socket" in label:
            return num
        else:
            return num + socket_number

    def read_power(self) -> list[RaplStats]:
        ret = []
        for label, inp in self.inputs:
            value = cat(inp)
            ret.append(RaplStats(label, float(value), 0.0))
        return ret

    @staticmethod
    def available() -> bool:
        return os.path.exists("/sys/devices/platform/amd_energy.0")


class AMDRaplMsrReader:
    def __init__(self) -> None:
        self.core_cpus: dict[int, int] = {}
        self.package_cpus: dict[int, int] = {}
        for i in range(cpu_count()):
            curr_core_id = int(
                cat(
                    f"/sys/devices/system/cpu/cpu{i}/topology/core_id",
                    binary=False,
                )
            )
            if curr_core_id not in self.core_cpus:
                self.core_cpus[curr_core_id] = i

            curr_package_id = int(
                cat(
                    f"/sys/devices/system/cpu/cpu{i}/topology/physical_package_id",
                    binary=False,
                )
            )
            if curr_package_id not in self.package_cpus:
                self.package_cpus[curr_package_id] = i

    def read_power(self) -> list[RaplStats]:
        ret = []
        first_cpu = next(iter(self.package_cpus.values()))
        unit_msr = read_msr(first_cpu, UNIT_MSR)
        energy_factor = 0.5 ** ((unit_msr & ENERGY_UNIT_MASK) >> 8)

        for i, cpu in self.package_cpus.items():
            package_msr = read_msr(cpu, PACKAGE_MSR)
            ret.append(
                RaplStats(
                    "Package " + str(i + 1),
                    package_msr * energy_factor * MICRO_JOULE_IN_JOULE,
                    0.0,
                )
            )

        for i, cpu in self.core_cpus.items():
            core_msr = read_msr(cpu, CORE_MSR)
            ret.append(
                RaplStats(
                    "Core " + str(i + 1),
                    core_msr * energy_factor * MICRO_JOULE_IN_JOULE,
                    0.0,
                )
            )

        return ret

    @staticmethod
    def available() -> bool:
        try:
            cpuinfo = cat("/proc/cpuinfo", binary=False)
            # The reader only supports family 17h CPUs
            m = re.search(r"vendor_id[\s]+: ([A-Za-z]+)", cpuinfo)

            if not m or m is None:
                return False

            if m.group(1) != "AuthenticAMD":
                return False

            m = re.search(r"cpu family[\s]+: ([0-9]+)", cpuinfo)
            if not m:
                return False
            if int(m[1]) != 0x17:
                return False
        except (FileNotFoundError, PermissionError):
            return False

        # with open("/proc/cpuinfo", "rb") as cpuinfo:
        #     all_info = cpuinfo.readlines()
        #     for line in all_info:
        #         if b"vendor_id" in line:
        #             print("Vendor id", line)
        #             if b"AuthenticAMD" not in line:
        #                 return False

        #     for line in all_info:
        #         if b"cpu family" in line:
        #             print("cpu family", line)
        #             m = re.search("cpu family[\s]+: ([0-9]+)", cpuinfo)
        #             if int(m[1]) != 0x17:
        #                 return False

        return msr_available()


def get_power_reader() -> RaplReader | AMDEnergyReader | AMDRaplMsrReader | None:
    for ReaderType in (RaplReader, AMDEnergyReader, AMDRaplMsrReader):
        if ReaderType.available():
            return ReaderType()
    return None
