#!/usr/bin/env python3
"""
Battery tray icon for Argon ONE UP on Wayland (labwc).

Reads the battery state of charge and AC status directly from the battery
management IC over I2C (bus 1, address 0x64) -- the same registers the
oneUpPower kernel module uses. This means the icon works for an ordinary
user (member of the ``i2c`` group) without sudo and without the kernel
module being loaded. Updates every 5 seconds.
"""

import sys
import time

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

# I2C layout of the Argon ONE UP battery IC (see battery/oneUpPower.c).
I2C_BUS = 1
BATTERY_ADDR = 0x64
SOC_REG = 0x04          # state of charge, percent
CURRENT_HIGH_REG = 0x0E  # bit 7 set => running on battery (not plugged in)
UPDATE_INTERVAL_MS = 5000


def read_battery(bus):
    """Return (percent, plugged_in) read over I2C, or (None, None) on error.

    force=True bypasses the I2C_SLAVE busy check so the read still works if
    the oneUpPower kernel module has claimed the address.
    """
    try:
        soc = bus.read_byte_data(BATTERY_ADDR, SOC_REG, force=True)
        current_high = bus.read_byte_data(BATTERY_ADDR, CURRENT_HIGH_REG,
                                          force=True)
    except OSError:
        return None, None

    percent = max(0, min(100, soc))
    plugged_in = (current_high & 0x80) == 0
    return percent, plugged_in


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.bus = smbus2.SMBus(I2C_BUS)
        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):
        percent, plugged_in = read_battery(self.bus)

        if percent is not None:
            charging = plugged_in and percent <= 95
            icon = make_battery_icon(percent, charging)
            self.tray.setIcon(icon)

            if charging:
                state = "Charging"
            elif plugged_in:
                state = "On AC"
            else:
                state = "Discharging"
            tooltip = f"Battery: {percent}% ({state})"
            self.tray.setToolTip(tooltip)
            self.status_action.setText(tooltip)
        else:
            # I2C read failed - IC not reachable
            icon = make_battery_icon(0, False)
            self.tray.setIcon(icon)
            self.tray.setToolTip("Battery: N/A (I2C read failed)")
            self.status_action.setText("Battery: N/A (I2C read failed)")

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


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


if __name__ == "__main__":
    main()