#!/usr/bin/env python3
"""
Battery tray icon for Argon ONE UP on Wayland (labwc).
Reads SOC from /sys/class/power_supply/BAT0/capacity and charging
status from AC0/online. Updates icon every 5 seconds.
"""

import sys
import os
import time
from PyQt5.QtWidgets import QApplication, QSystemTrayIcon, QMenu, QAction
from PyQt5.QtGui import QIcon, QPainter, QColor, QFont, QPixmap
from PyQt5.QtCore import QTimer, Qt

SYSFS_BAT_CAPACITY = "/sys/class/power_supply/BAT0/capacity"
SYSFS_BAT_STATUS = "/sys/class/power_supply/BAT0/status"
SYSFS_AC_ONLINE = "/sys/class/power_supply/AC0/online"
UPDATE_INTERVAL_MS = 5000


def read_sysfs(path, default=None):
    try:
        with open(path) as f:
            return f.read().strip()
    except (FileNotFoundError, PermissionError, OSError):
        return default


def make_battery_icon(percent, charging):
    """Draw a battery icon with percentage text."""
    size = 64
    pix = QPixmap(size, size)
    pix.fill(Qt.transparent)
    p = QPainter(pix)
    p.setRenderHint(QPainter.Antialiasing)

    # Battery outline
    bx, by, bw, bh = 8, 14, 48, 36
    # Terminal nub
    nx, ny, nw, nh = bx + bw, 24, 6, 16

    p.setPen(QColor(200, 200, 200))
    p.setBrush(Qt.NoBrush)
    p.drawRoundedRect(bx, by, bw, bh, 4, 4)
    p.drawRoundedRect(nx, ny, nw, nh, 2, 2)

    # Fill level
    margin = 3
    fill_max_w = bw - 2 * margin
    fill_w = max(1, int(fill_max_w * percent / 100))
    fill_x = bx + margin
    fill_y = by + margin
    fill_h = bh - 2 * margin

    if charging:
        fill_color = QColor(100, 200, 255)  # blue when charging
    elif percent <= 10:
        fill_color = QColor(255, 60, 60)    # red when critical
    elif percent <= 25:
        fill_color = QColor(255, 180, 0)    # orange when low
    else:
        fill_color = QColor(80, 220, 80)    # green normal

    p.setPen(Qt.NoPen)
    p.setBrush(fill_color)
    p.drawRoundedRect(fill_x, fill_y, fill_w, fill_h, 2, 2)

    # Percentage text
    font = QFont("Sans", 11, QFont.Bold)
    p.setFont(font)
    p.setPen(QColor(255, 255, 255))
    text = f"{percent}%"
    p.drawText(bx, by, bw, bh, Qt.AlignCenter, text)

    # Charging indicator (lightning bolt)
    if charging:
        font2 = QFont("Sans", 9)
        p.setFont(font2)
        p.setPen(QColor(255, 255, 100))
        p.drawText(0, 0, size, 16, Qt.AlignCenter, "⚡")

    p.end()
    return QIcon(pix)


class BatteryTray:
    def __init__(self):
        self.app = QApplication(sys.argv)
        self.app.setQuitOnLastWindowClosed(False)

        # At login the panel may not have started its system tray yet.
        # Wait for it instead of giving up, so autostart works after a reboot.
        deadline = time.monotonic() + 120
        while not QSystemTrayIcon.isSystemTrayAvailable():
            if time.monotonic() > deadline:
                print("System tray not available after 120s", file=sys.stderr)
                sys.exit(1)
            time.sleep(2)

        self.tray = QSystemTrayIcon()

        # Context menu
        menu = QMenu()
        self.status_action = QAction("Battery: --")
        self.status_action.setEnabled(False)
        menu.addAction(self.status_action)
        menu.addSeparator()
        quit_action = QAction("Quit")
        quit_action.triggered.connect(self.app.quit)
        menu.addAction(quit_action)
        self.tray.setContextMenu(menu)

        # Initial update
        self.update()

        # Timer for periodic updates
        self.timer = QTimer()
        self.timer.timeout.connect(self.update)
        self.timer.start(UPDATE_INTERVAL_MS)

        self.tray.show()

    def update(self):
        cap_str = read_sysfs(SYSFS_BAT_CAPACITY)
        status = read_sysfs(SYSFS_BAT_STATUS, "Unknown")
        ac_str = read_sysfs(SYSFS_AC_ONLINE, "0")

        if cap_str is not None:
            try:
                percent = int(cap_str)
            except ValueError:
                percent = 0
            charging = (ac_str == "1" or status == "Charging")
            icon = make_battery_icon(percent, charging)
            self.tray.setIcon(icon)

            if charging:
                state = "Charging"
            elif ac_str == "1":
                state = "On AC"
            else:
                state = "Discharging"
            tooltip = f"Battery: {percent}% ({state})"
            self.tray.setToolTip(tooltip)
            self.status_action.setText(tooltip)
        else:
            # No battery sysfs - module not loaded
            icon = make_battery_icon(0, False)
            self.tray.setIcon(icon)
            self.tray.setToolTip("Battery: N/A (module not loaded?)")
            self.status_action.setText("Battery: N/A (module not loaded?)")

    def run(self):
        return self.app.exec()


def main():
    bt = BatteryTray()
    sys.exit(bt.run())


if __name__ == "__main__":
    main()