# SPDX-License-Identifier: BSD-3-Clause
# SPDX-FileCopyrightText: 2025, Alliance for Sustainable Energy, LLC

import logging
import os
import re
from typing import Iterable
import numpy as np

from .base import BaseReader
from .utils import Quantity, Energy, Joule, Unit, Watt, Power, Second

# Module-level logger
logger = logging.getLogger(__name__)


def _get_rapl_domain_name(rapl_device_path, tag_for_unnamed_device: str) -> str:
    """Generate a domain name for the RAPL device based on its name and path.

    :param rapl_device_path: Path to the RAPL device directory.
    :param tag_for_unnamed_device: Tag to use if the device name cannot be determined.
    """

    # If the path ends in a pattern like :<number>:<number>
    if re.search(r":\d+:\d+$", rapl_device_path):
        # Remove the last part after the last colon
        parent_rapl_component_dir = re.sub(r":\d+$", "", rapl_device_path)
        domain_name = (
            _get_rapl_domain_name(parent_rapl_component_dir, tag_for_unnamed_device)
            + "-"
        )
    else:
        domain_name = ""

    try:
        # Read the actual domain name from the 'name' file
        with open(os.path.join(rapl_device_path, "name"), "r") as f:
            _name = f.read().strip()

        # Replace "package-" prefix by "cpu-"
        if _name.startswith("package-"):
            _name = "cpu-" + _name.split("-")[1]

        # Append the name to the domain name
        domain_name += _name

    except (FileNotFoundError, PermissionError, OSError):
        # Fallback to extracting from path if name file is not readable
        last_digit = re.search(r"\d+$", rapl_device_path)
        if last_digit:
            domain_name += last_digit[0]
        else:
            domain_name += tag_for_unnamed_device

    return domain_name


class RAPLDevice(BaseReader):
    """Reader for RAPL (Running Average Power Limit) devices.

    :param rapl_device_path: Path to the RAPL device directory.

    .. attribute:: UNITS

        Dictionary of measurement units for physical quantities.

    .. attribute:: path

        Path to the RAPL device directory, typically under /sys/class/powercap/intel-rapl.

    .. attribute:: name

        Name of the RAPL device, read from the 'name' file.

    .. attribute:: max_energy_range

        Maximum energy range, read from the 'max_energy_range_uj' file.

    """

    def __init__(self, rapl_device_path: str) -> None:
        super().__init__((Energy,))

        self.path = rapl_device_path

        # Read name file
        self.name = None
        try:
            with open(os.path.join(self.path, "name"), "r") as f:
                self.name = f.read().strip()
        except FileNotFoundError:
            logger.warning(f"Name file not found for {self.path}")

        # Read max energy range file
        self.max_energy_range = 0
        try:
            with open(os.path.join(self.path, "max_energy_range_uj"), "r") as f:
                self.max_energy_range = int(f.read().strip())
        except FileNotFoundError:
            logger.warning(f"Max energy range file not found for {self.path}")

        # Open handle to energy file, which is used to read the energy counter
        self.energy_file = None
        try:
            self.energy_file = open(os.path.join(self.path, "energy_uj"), "r")
        except FileNotFoundError:
            logger.warning(f"Energy file not found for {self.path}")

        # Post-process device name
        self._device_name = _get_rapl_domain_name(
            self.path, tag_for_unnamed_device="unknown"
        )

    def __del__(self):
        """Ensure the energy file is closed when the object is deleted."""
        if self.energy_file:
            self.energy_file.close()

    @property
    def derived_quantities(self) -> list[type[Quantity]]:
        return [Power]

    @property
    def tags(self) -> list[str]:
        return [f"{self._device_name}[{self.get_unit(q)}]" for q in self.quantities]

    @property
    def derived_tags(self) -> list[str]:
        return [
            f"{self._device_name}[{self.get_unit(q)}]" for q in self.derived_quantities
        ]

    def get_unit(self, quantity: type[Quantity]) -> Unit:
        UNITS = {Energy: Joule("u")}
        if quantity in UNITS:
            return UNITS[quantity]
        elif quantity in self.derived_quantities:
            return Watt()
        else:
            logger.warning(
                f"Invalid quantity requested: {quantity}. "
                f"Supported quantities are: {list(UNITS.keys())}."
            )
            return Unit()

    def read_energy(self) -> int:
        """Read the energy counter for the i-th device."""
        if self.energy_file:
            try:
                self.energy_file.seek(0)
                return int(self.energy_file.read().strip())
            except ValueError as e:
                logger.error(f"Failed to read energy for {self.path}: {e}")
                return 0
        else:
            logger.error(f"Energy file is not open for {self.path}.")
            return 0

    def read(self) -> list[int]:
        return [self.read_energy()]

    def compute_derived(self, time_series, data_series, time_unit: Second = Second()):
        # Compute deltas
        energy_delta = np.subtract(data_series[1:], data_series[:-1])
        time_delta = np.subtract(time_series[1:], time_series[:-1])
        if len(energy_delta) > 0:
            energy_delta[energy_delta < 0] += self.max_energy_range

        # Compute power in Watts
        res = np.asarray([(num / den) for num, den in zip(energy_delta, time_delta)])
        res *= self.get_unit(Energy).to_si() / time_unit.to_si()  # Convert to Watts

        return res.tolist()


class RAPLReader(BaseReader):
    """Reader for RAPL devices in the system.

    :param rapl_dir: Directory where RAPL devices are located, typically /sys/class/powercap/intel-rapl/subsystem.

    .. attribute:: rapl_dir

        Directory where RAPL devices are located.

    .. attribute:: devices

        List of RAPLDevice instances for available RAPL devices.
    """

    def __init__(self, rapl_dir="/sys/class/powercap/intel-rapl/subsystem") -> None:
        super().__init__((Energy,))

        self.rapl_dir = rapl_dir

        # Find all RAPL devices in the specified directory
        self.devices = []
        for root, dirs, _ in os.walk(rapl_dir):
            for dir_name in dirs:
                rapl_file = os.path.join(root, dir_name, "energy_uj")
                if os.path.exists(rapl_file):
                    d = RAPLDevice(os.path.join(root, dir_name))
                    if d.energy_file:
                        self.devices.append(d)

        # Order devices by their path for consistency
        self.devices.sort(key=lambda d: d.path)

        # Post-process tags
        count = 0
        self._tags = []
        for device in self.devices:
            tag = self.devices[count].tags[0]

            if "unknown" in tag:
                tag = f"unknown-{count}"
            count += 1

            self._tags.append(tag)

    @property
    def derived_quantities(self) -> list[type[Quantity]]:
        return [Power]

    @property
    def tags(self) -> list[str]:
        return [
            tag if "unknown" not in tag else f"unknown-{i}"
            for i, device in enumerate(self.devices)
            for tag in device.tags
        ]

    @property
    def derived_tags(self) -> list[str]:
        return [tag for device in self.devices for tag in device.derived_tags]

    def get_unit(self, quantity: type[Quantity]) -> Unit:
        UNITS = {Energy: Joule("u")}
        if quantity in UNITS:
            return UNITS[quantity]
        elif quantity in self.derived_quantities:
            return Watt()
        else:
            logger.warning(
                f"Invalid quantity requested: {quantity}. "
                f"Supported quantities are: {list(UNITS.keys())}."
            )
            return Unit()

    def read_energy_on_device(self, i: int) -> int:
        """Read the energy counter of the i-th device."""
        try:
            return self.devices[i].read_energy()
        except IndexError:
            logger.error(f"Device index {i} out of range.")
            return 0
        except Exception as e:
            logger.error(f"Failed to read energy for device {i}: {e}")
            return 0

    def read_energy(self) -> list[int]:
        """Read the current energy counter for all RAPL devices."""
        return [self.read_energy_on_device(i) for i in range(len(self.devices))]

    # Alias for read_energy to match the BaseReader interface
    read = read_energy

    def compute_derived(
        self, time_series: Iterable, data_series: Iterable, time_unit: Second = Second()
    ):
        data_series_t = list(zip(*data_series))  # Transpose to group by device
        res = [
            d.compute_derived(time_series, ds, time_unit)
            for d, ds in zip(self.devices, data_series_t)
        ]
        return list(zip(*res))  # Transpose to match expected output format
