#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p python3
"""
usage: sane-sysload [options...] formatstr

pretty-prints a battery estimate (icon to indicate state, and a duration estimate)

options:
  --debug: output additional information, to stderr
  --minute-suffix <string>:  use the provided string as a minutes suffix
  --hour-suffix <string>:  use the provided string as an hours suffix
  --icon-suffix <string>:  use the provided string as an icon suffix
  --percent-suffix <string>:  use the provided string when displaying percents

formatstr is a Python format string.
variables available for formatting:
- {bat_icon}
- {bat_pct}
- {bat_time}
- {cpu_icon}
- {cpu_pct}
- {mem_icon}
- {mem_pct}
and some presets, encapsulating the above:
- {all}
- {bat}
- {cpu}
- {mem}
"""

import argparse
import logging
import time

from dataclasses import dataclass
from enum import Enum

logger = logging.getLogger(__name__)

# these icons may only render in nerdfonts
ICON_BAT_CHG = ["󰢟", "󱊤", "󱊥", "󰂅"]
ICON_BAT_DIS = ["󰂎", "󱊡", "󱊢", "󱊣"]
ICON_CPU=""
ICON_MEM="☵"
SUFFIX_ICON = " "  # thin space
SUFFIX_PERCENT = "%"
# SUFFIX_ICON=" "

# render time like: 2ʰ08ᵐ
# unicode sub/super-scripts: <https://en.wikipedia.org/wiki/Unicode_subscripts_and_superscripts>
# SUFFIX_HR="ʰ"
# SUFFIX_MIN="ᵐ"

# render time like: 2ₕ08ₘ
# SUFFIX_HR="ₕ"
# SUFFIX_MIN="ₘ"

# render time like: 2h08m
# SUFFIX_HR="H"
# SUFFIX_MIN="M"

# render time like: 2:08
# SUFFIX_HR=":"
# SUFFIX_MIN=

# render time like: 2꞉08⧗
SUFFIX_HR="꞉"
SUFFIX_MIN="⧗"
# variants:
# SUFFIX_HR=":"
# SUFFIX_MIN="⧖"
# SUFFIX_MIN="⌛"

# render time like: 2'08"
# SUFFIX_HR="'"
# SUFFIX_MIN='"'

class ChargeDirection(Enum):
    Charging = "Charging"
    Discharging = "Discharging"

@dataclass
class Formatter:
    suffix_icon: str = SUFFIX_ICON
    suffix_percent: str = SUFFIX_PERCENT
    suffix_hr: str = SUFFIX_HR
    suffix_min: str = SUFFIX_MIN

    def render_charge_icon(self, direction: ChargeDirection, percentage: float) -> str:
        return f"{self._choose_icon(direction, percentage)}{self.suffix_icon}"

    def render_cpu_icon(self) -> str:
        return f"{ICON_CPU}{self.suffix_icon}"

    def render_mem_icon(self) -> str:
        return f"{ICON_MEM}{self.suffix_icon}"

    def render_hours_minutes(self, minutes: int) -> str:
        hr = minutes // 60
        min = minutes % 60
        return f"{hr}{self.suffix_hr}{min:02}{self.suffix_min}"

    def render_percent(self, pct: int) -> str:
        return f"{pct}{self.suffix_percent}"

    def _choose_icon(self, direction: ChargeDirection, percentage: float) -> str:
        level = percentage / 25
        level = max(0, min(3, level))
        level = int(round(level))

        logger.debug(f"render_charge_icon: direction={direction} level={level}")

        if direction == ChargeDirection.Charging:
            return ICON_BAT_CHG[level]
        elif direction == ChargeDirection.Discharging:
            return ICON_BAT_DIS[level]

        raise RuntimeError(f"invalid ChargeDirection {direction}")

class MemInfo:
    """
    reads values from /proc/meminfo
    """
    def __init__(self):
        try:
            lines = open("/proc/meminfo").readlines()
        except Exception as e:
            logger.info(f"failed to open /proc/meminfo: {e}")
            lines = []

        # lines are like:
        # MemTotal:       16262708 kB
        # HugePages_Total:       0
        self.entries = {}
        for l in lines:
            if ":" not in l: continue
            key_len = l.index(":")
            key, value_str = l[:key_len].strip(), l[key_len+1:].strip()

            unit_str = ""
            if " " in value_str:
                value_len = value_str.index(" ")
                value_str, unit_str = value_str[:value_len].strip(), value_str[value_len+1:].strip()

            try:
                value = int(value_str)
            except:
                logger.info(f"unexpected /proc/meminfo line: {l}")
                continue

            if unit_str == "kB":
                value = value * 1024
            self.entries[key] = value

    def get(self, entry):
        v = self.entries.get(entry)
        logger.debug(f"/proc/meminfo: {entry}={v}")
        return v

class ProcStat:
    """
    reads vaues from /proc/stat, mostly CPU-related.
    """
    # /proc/stat format is documented here: <https://www.linuxhowtos.org/System/procstat.htm>
    # these are AGGREGATES SINCE SYSTEM BOOT
    # to measure current CPU usage, need to take multiple samples
    # cpu  <user> <system> <nice> <idle> <iowait> <irg> <softirq> 0 0 0
    # (what are the last three fields?)
    # where:
    # measurements are in units of jiffies or USER_HZ
    # user: normal processes executing in user mode
    # nice: niced processes executing in user mode
    # system: processes executing in kernel mode
    # idle: twiddling thumbs
    # iowait: waiting for I/O to complete
    # irq: servicing interrupts
    # softirq: servicing softirqs
    def __init__(self, entries=None):
        if entries is not None:
            self.entries = entries
            return

        # else, read from procfs...
        try:
            lines = open("/proc/stat").readlines()
        except Exception as e:
            logger.info(f"failed to open /proc/stat: {e}")
            lines = []

        self.entries = {}
        for l in lines:
            pieces = l.strip().split(" ")
            name, values = pieces[0], [p for p in pieces[1:] if p]
            if name:
                self.entries[name] = [int(v) for v in values]

    @staticmethod
    def sample(seconds: float = 1.0):
        sample1 = ProcStat()
        time.sleep(seconds)
        sample2 = ProcStat()
        return sample2 - sample1

    def __sub__(self, other: 'ProcStat') -> 'ProcStat':
        entries = {}
        for k in self.entries:
            entries[k] = [i - j for i, j in zip(self.entries[k], other.entries[k])]
        return ProcStat(entries)

    @property
    def cpu_user(self) -> int:
        return self.entries["cpu"][0]
    @property
    def cpu_system(self) -> int:
        return self.entries["cpu"][1]
    @property
    def cpu_nice(self) -> int:
        return self.entries["cpu"][2]
    @property
    def cpu_idle(self) -> int:
        return self.entries["cpu"][3]
    @property
    def cpu_iowait(self) -> int:
        return self.entries["cpu"][4]
    @property
    def cpu_irq(self) -> int:
        return self.entries["cpu"][5]
    @property
    def cpu_softirq(self) -> int:
        return self.entries["cpu"][6]

    @property
    def cpu_total(self) -> int:
        # TODO: not sure if i'm supposed to include irq stuff here?
        return self.cpu_user + self.cpu_system + self.cpu_nice + self.cpu_idle + self.cpu_iowait + self.cpu_irq + self.cpu_softirq

class PowerSupply:
    """
    reads values from /sys/class/power_supply/$dev/ API
    """
    def __init__(self, sysfs_node: str):
        self.sysfs_node = sysfs_node
        self._cached_reads = {}

    def try_read(self, rel_path: str) -> str | None:
        if rel_path not in self._cached_reads:
            self._cached_reads[rel_path] = self.try_read_uncached(rel_path)

        return self._cached_reads[rel_path]

    def try_read_uncached(self, rel_path: str) -> str | None:
        try:
            v = open(f"{self.sysfs_node}/{rel_path}").read()
            logger.debug(f"{self.sysfs_node}/{rel_path}: {v}")
            return v
        except:
            return None

    def try_read_int(self, rel_path: str) -> int | None:
        s = self.try_read(rel_path)
        if s is None: return None
        return int(s)

    @property
    def capacity(self) -> int | None:
        return self.try_read_int("capacity")  # percent

    @property
    def charge_full_design(self) -> int | None:
        return self.try_read_int("charge_full_design")  # micro-Ah

    @property
    def current_now(self) -> int | None:
        return self.try_read_int("current_now")  # micro-A

    @property
    def energy_full(self) -> int | None:
        return self.try_read_int("energy_full")  # micro-Wh

    # @property
    # def energy_now(self) -> int | None:
    #     return self.try_read_int("energy_now")  # micro-Wh

    @property
    def power_now(self) -> int | None:
        return self.try_read_int("power_now")  # micro-W (?)

    @property
    def status(self) -> str | None:
        return self.try_read("status")  # "Charging"/"Discharging"/"Not charging"

    # @property
    # def voltage_now(self) -> int | None:
    #     return self.try_read_int("voltage_now")  # micro-V


class BatteryInfo:
    """
    higher-level battery info derived from the underlying power supply
    """
    percent_charged: int  #< always available
    minutes_to_charged: int | None = None
    minutes_to_discharged: int | None = None

    def __init__(self, capacity: int, charges_per_hour: float | None, status: str | None):
        self.percent_charged = capacity

        # correct some batteries which flip signs in places
        if status.lower().strip() == "charging" and charges_per_hour:
            charges_per_hour = abs(charges_per_hour)
            logger.debug("status==charging => forcing charges_per_hour positive")
        elif status.lower().strip() == "discharging" and charges_per_hour:
            charges_per_hour = -abs(charges_per_hour)
            logger.debug("status==discharging => forcing charges_per_hour negative")

        if charges_per_hour is not None and charges_per_hour < 0:
            self.minutes_to_discharged = int(
                60
                * self.percent_charged/100
                / -charges_per_hour
            )
        if charges_per_hour is not None and charges_per_hour > 0:
            self.minutes_to_charged = int(
                60
                * (100-self.percent_charged)/100
                / charges_per_hour
            )


def try_battery_path(p: str) -> BatteryInfo | None:
    """
    try to read battery information from some p = "/sys/class/power_supply/$node" path
    """
    ps = PowerSupply(p)
    if ps.capacity is None:
        return None

    if ps.charge_full_design and ps.current_now is not None:
        # current_now is positive when charging
        charges_per_hour = ps.current_now / ps.charge_full_design
    elif ps.energy_full and ps.power_now is not None:
        # power_now is positive when discharging
        charges_per_hour = -ps.power_now / ps.energy_full
    else:
        charges_per_hour = None

    return BatteryInfo(ps.capacity, charges_per_hour, ps.status)

def try_all_batteries() -> BatteryInfo | None:
    p = try_battery_path("/sys/class/power_supply/axp20x-battery")  # Pinephone
    if p is None:
        p = try_battery_path("/sys/class/power_supply/BAT0")  # Thinkpad

    logger.debug(f"perc: {p.percent_charged if p else None}")
    logger.debug(f"charge: {p.percent_charged if p else None}")
    logger.debug(f"min-to-charge: {p.minutes_to_charged if p else None}, min-to-discharge: {p.minutes_to_discharged if p else None}")

    return p

@dataclass
class AllInfo:
    _fmt: Formatter
    __bat: BatteryInfo | None = None
    __cpu: ProcStat | None = None
    __mem: MemInfo | None = None

    # lazy-loading
    @property
    def _bat(self):
        if self.__bat is None:
            self.__bat = try_all_batteries()
        return self.__bat
    @property
    def _cpu(self):
        if self.__cpu is None:
            self.__cpu = ProcStat.sample()
        return self.__cpu
    @property
    def _mem(self):
        if self.__mem is None:
            self.__mem = MemInfo()
        return self.__mem

    # user-facing format shorthands
    @property
    def all(self) -> str:
        return f"{self.bat_pct} {self.bat} {self.mem} {self.cpu}"
    @property
    def bat(self) -> str:
        return f"{self.bat_icon}{self.bat_time}"
    @property
    def cpu(self) -> str:
        return f"{self.cpu_icon}{self.cpu_pct}"
    @property
    def mem(self) -> str:
        return f"{self.mem_icon}{self.mem_pct}"

    # manual/low-level fields
    @property
    def mem_icon(self) -> str:
        if self._mem is None: return ""
        return self._fmt.render_mem_icon()

    @property
    def mem_pct(self) -> str:
        if self._mem is None: return ""

        total = self._mem.get("MemTotal")
        free = self._mem.get("MemAvailable")

        if total is None or free is None or free > total:
            return ""

        mem_use_pct = int((total - free) / total * 100)
        return self._fmt.render_percent(mem_use_pct)

    @property
    def bat_icon(self) -> str:
        if self._bat is None: return ""
        elif self._bat.minutes_to_charged != None:
            logger.debug("bat_icon: charging")
            return self._fmt.render_charge_icon(ChargeDirection.Charging, self._bat.percent_charged)
        else:
            logger.debug("bat_icon: discharging")
            return self._fmt.render_charge_icon(ChargeDirection.Discharging, self._bat.percent_charged)

    @property
    def bat_pct(self) -> str:
        if self._bat is None: return ""
        return self._fmt.render_percent(self._bat.percent_charged)

    @property
    def bat_time(self) -> str:
        if self._bat is None: return ""
        elif self._bat.minutes_to_charged != None:
            duration = self._bat.minutes_to_charged
        else:
            duration = self._bat.minutes_to_discharged

        if duration is not None and duration < 1440:
            return self._fmt.render_hours_minutes(duration)
        else:
            return self._fmt.render_percent(self._bat.percent_charged)

    @property
    def cpu_icon(self) -> str:
        if self._cpu is None: return ""
        return self._fmt.render_cpu_icon()

    @property
    def cpu_pct(self) -> str:
        if self._cpu is None:
            return ""

        idle = self._cpu.cpu_idle + self._cpu.cpu_iowait
        total = self._cpu.cpu_total

        cpu_use_pct = int((total - idle) / total * 100)
        return self._fmt.render_percent(cpu_use_pct)


class LazyFormatter:
    def __init__(self, obj: object, attr: str):
        self.obj = obj
        self.attr = attr

    def __repr__(self) -> str:
        return repr(getattr(self.obj, self.attr))

    def __str__(self) -> str:
        return str(getattr(self.obj, self.attr))

def main() -> None:
    logging.basicConfig()
    logging.getLogger().setLevel(logging.INFO)

    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--icon-suffix", default=SUFFIX_ICON)
    parser.add_argument("--hour-suffix", default=SUFFIX_HR)
    parser.add_argument("--minute-suffix", default=SUFFIX_MIN)
    parser.add_argument("--percent-suffix", default=SUFFIX_PERCENT)
    parser.add_argument("formatstr", nargs="*", default=["{all}"])
    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)

    f = Formatter()
    f.suffix_icon = args.icon_suffix
    f.suffix_percent = args.percent_suffix
    f.suffix_hr = args.hour_suffix
    f.suffix_min = args.minute_suffix

    info = AllInfo(f)
    formatstr = " ".join(args.formatstr)
    print(formatstr.format(
      all=LazyFormatter(info, "all"),
      bat=LazyFormatter(info, "bat"),
      bat_icon=LazyFormatter(info, "bat_icon"),
      bat_pct=LazyFormatter(info, "bat_pct"),
      bat_time=LazyFormatter(info, "bat_time"),
      cpu=LazyFormatter(info, "cpu"),
      cpu_icon=LazyFormatter(info, "cpu_icon"),
      cpu_pct=LazyFormatter(info, "cpu_pct"),
      mem=LazyFormatter(info, "mem"),
      mem_icon=LazyFormatter(info, "mem_icon"),
      mem_pct=LazyFormatter(info, "mem_pct"),
    ))

if __name__ == "__main__":
    main()
