#!/usr/bin/env python3

# vi: filetype=python

import gi
gi.require_version("Gtk", "4.0")
from gi.repository import Gtk, Gio, GLib, GObject

import subprocess
import shlex
import threading
import uuid
import argparse
import logging

# CLI / Logging
parser = argparse.ArgumentParser(description="Systemd Scope Manager (GTK4 ColumnView)")
parser.add_argument("--log-level", "-l", choices=["DEBUG","INFO","WARNING","ERROR","CRITICAL"], default="INFO")
parser.add_argument("--values-mode", "-m", choices=["raw","human"], default="human",
                    help="Display memory / CPU values raw or human‑readable")
args = parser.parse_args()

logging.basicConfig(level=getattr(logging, args.log_level.upper()),
                    format="%(asctime)s %(name)s %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

# Helpers
def human_bytes(value: int) -> str:
    for unit in ("B","KB","MB","GB","TB"):
        if value < 1024:
            return f"{value:.1f}{unit}"
        value /= 1024
    return f"{value:.1f}PB"

def human_time(nsec: int) -> str:
    sec = nsec / 1_000_000_000
    if sec < 60:
        return f"{sec:.1f}s"
    minutes = sec / 60
    if minutes < 60:
        return f"{minutes:.1f}m"
    return f"{minutes/60:.1f}h"

def run_systemctl_show(unit: str) -> dict:
    cmd = [
        "systemctl", "--user", "show", unit,
        "--property=MemoryCurrent,MemorySwapCurrent,CPUUsageNSec,ActiveState,Restart"
    ]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, check=True)
        props = {}
        for line in res.stdout.splitlines():
            if "=" in line:
                k, v = line.split("=", 1)
                props[k.strip()] = v.strip()
        return props
    except Exception:
        logger.exception("systemctl show failed for %s", unit)
        return {}

# Data row
class ScopeRow(GObject.Object):
    __gtype_name__ = "ScopeRow"
    unit = GObject.Property(type=str)
    cli  = GObject.Property(type=str)
    mem  = GObject.Property(type=str)
    swap = GObject.Property(type=str)
    cpu  = GObject.Property(type=str)
    state= GObject.Property(type=str)

    def __init__(self, unit, cli, mem, swap, cpu, state):
        super().__init__()
        self.unit  = unit
        self.cli   = cli
        self.mem   = mem
        self.swap  = swap
        self.cpu   = cpu
        self.state = state

# Main Window
class ScopeManagerWindow(Gtk.ApplicationWindow):
    def __init__(self, app):
        super().__init__(application=app, title="Systemd Scope Manager")
        self.set_default_size(1000, 500)
        self.values_mode = args.values_mode
        self.scopes = {}

        self.model = Gio.ListStore(item_type=ScopeRow)
        self.sel   = Gtk.SingleSelection.new(self.model)
        self.view  = Gtk.ColumnView.new(self.sel)
        self.view.set_reorderable(True)
        self.view.set_show_row_separators(True)

        cols = [
            ("Unit",   "unit"),
            ("CLI",    "cli"),
            ("Memory", "mem"),
            ("Swap",   "swap"),
            ("CPU",    "cpu"),
            ("State",  "state"),
        ]
        for title, prop_name in cols:
            factory = Gtk.SignalListItemFactory()
            factory.connect("setup",  self._factory_setup_label)
            factory.connect("bind",   self._make_factory_bind(prop_name))
            col = Gtk.ColumnViewColumn()
            col.set_title(title)
            col.set_factory(factory)
            col.set_resizable(True)
            col.set_expand(True)
            self.view.append_column(col)

        # Actions column
        action_factory = Gtk.SignalListItemFactory()
        action_factory.connect("setup", self._factory_setup_actions)
        action_factory.connect("bind",  self._factory_bind_actions)
        act_col = Gtk.ColumnViewColumn()
        act_col.set_title("⋮")
        act_col.set_factory(action_factory)
        act_col.set_resizable(False)
        act_col.set_expand(False)
        self.view.append_column(act_col)

        # Input area
        self.cmd_entry = Gtk.Entry()
        self.cmd_entry.set_placeholder_text("Command to run in new scope")
        run_btn = Gtk.Button(label="Run")
        run_btn.connect("clicked", self.on_run_clicked)
        prop_btn = Gtk.Button(label="Edit Property")
        prop_btn.connect("clicked", self.on_edit_property)

        input_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=6)
        input_box.append(self.cmd_entry)
        input_box.append(run_btn)
        input_box.append(prop_btn)

        scrolled = Gtk.ScrolledWindow()
        scrolled.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        scrolled.set_propagate_natural_width(True)
        scrolled.set_propagate_natural_height(True)
        scrolled.set_child(self.view)

        vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
        vbox.set_vexpand(True)
        vbox.append(input_box)
        vbox.append(scrolled)
        self.set_child(vbox)

        GLib.timeout_add_seconds(2, self.refresh_scopes)

    def _factory_setup_label(self, factory, list_item):
        lbl = Gtk.Label()
        list_item.set_child(lbl)

    def _make_factory_bind(self, prop_name):
        def bind(factory, list_item):
            row = list_item.get_item()
            lbl = list_item.get_child()
            lbl.set_text(getattr(row, prop_name))
            row.connect(f"notify::{prop_name}", lambda obj, pspec: lbl.set_text(getattr(obj, prop_name)))
        return bind

    def _factory_setup_actions(self, factory, list_item):
        btn = Gtk.MenuButton()
        btn.set_icon_name("open-menu-symbolic")
        # style class if needed:
        btn.get_style_context().add_class("flat")
        list_item.set_child(btn)

    def _factory_bind_actions(self, factory, list_item):
        row = list_item.get_item()
        btn = list_item.get_child()

        menu = Gio.Menu()
        menu.append("Stop",    f"app.stop_{row.unit}")
        menu.append("Restart", f"app.restart_{row.unit}")
        menu.append("Toggle Auto‑Restart", f"app.toggle_{row.unit}")

        btn.set_menu_model(menu)

        # create actions
        self._ensure_row_actions(row)

    def _ensure_row_actions(self, row):
        unit = row.unit
        # Stop
        act_stop = Gio.SimpleAction.new(f"stop_{unit}", None)
        act_stop.connect("activate", lambda a, v, r=row: self.menu_action("Stop", r))
        self.add_action(act_stop)
        # Restart
        act_restart = Gio.SimpleAction.new(f"restart_{unit}", None)
        act_restart.connect("activate", lambda a, v, r=row: self.menu_action("Restart", r))
        self.add_action(act_restart)
        # Toggle
        act_toggle = Gio.SimpleAction.new(f"toggle_{unit}", None)
        act_toggle.connect("activate", lambda a, v, r=row: self.menu_action("Toggle Auto‑Restart", r))
        self.add_action(act_toggle)

    def on_run_clicked(self, button):
        cmd = self.cmd_entry.get_text().strip()
        if not cmd:
            return
        unit = f"app-{uuid.uuid4().int >> 64}.scope"
        argv = ["systemd-run", "--user", "--scope", "--unit", unit,
                "-p", "MemoryAccounting=yes", "-p", "CPUAccounting=yes"] + shlex.split(cmd)
        logger.info("Starting scope: %s", " ".join(argv))

        def worker():
            try:
                p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                out, err = p.communicate()
                if out:
                    logger.debug("systemd-run stdout [%s]: %s", unit, out.strip())
                if err:
                    logger.debug("systemd-run stderr [%s]: %s", unit, err.strip())
            except Exception:
                logger.exception("Failed to run systemd-run for %s", unit)

        threading.Thread(target=worker, daemon=True).start()
        self.scopes[unit] = cmd
        self.cmd_entry.set_text("")

    def on_edit_property(self, button):
        idx = self.sel.get_selected()
        if idx < 0:
            return
        row = self.model.get_item(idx)

        dialog = Gtk.Dialog(transient_for=self, modal=True, title="Set systemd property")
        content = dialog.get_content_area()
        grid = Gtk.Grid(row_spacing=6, column_spacing=6, margin=10)
        content.append(grid)

        lbl1 = Gtk.Label(label="Property (e.g. MemoryMax):")
        entry1 = Gtk.Entry()
        lbl2 = Gtk.Label(label="Value (e.g. 512M):")
        entry2 = Gtk.Entry()
        runtime_chk = Gtk.CheckButton(label="Runtime only")

        grid.attach(lbl1,    0,0,1,1)
        grid.attach(entry1, 1,0,1,1)
        grid.attach(lbl2,    0,1,1,1)
        grid.attach(entry2,  1,1,1,1)
        grid.attach(runtime_chk,0,2,2,1)

        dialog.add_button("OK",     Gtk.ResponseType.OK)
        dialog.add_button("Cancel", Gtk.ResponseType.CANCEL)
        dialog.show()

        resp = dialog.run()
        if resp == Gtk.ResponseType.OK:
            prop = entry1.get_text().strip()
            val  = entry2.get_text().strip()
            rt   = runtime_chk.get_active()
            dialog.destroy()
            self.set_property(row.unit, prop, val, rt)
        else:
            dialog.destroy()

    def set_property(self, unit, prop, val, runtime_flag):
        cmd = ["systemctl", "--user", "set-property"]
        if runtime_flag:
            cmd.append("--runtime")
        cmd += [unit, f"{prop}={val}"]
        logger.info("Setting %s=%s on %s (runtime=%s)", prop, val, unit, runtime_flag)
        try:
            res = subprocess.run(cmd, capture_output=True, text=True)
            if res.stdout:
                logger.debug("set-property stdout: %s", res.stdout.strip())
            if res.stderr:
                logger.debug("set-property stderr: %s", res.stderr.strip())
        except Exception:
            logger.exception("Failed set-property for %s", unit)

    def menu_action(self, label, row):
        unit = row.unit
        try:
            if label == "Stop":
                subprocess.run(["systemctl","--user","stop",unit])
            elif label == "Restart":
                subprocess.run(["systemctl","--user","restart",unit])
            elif label == "Toggle Auto‑Restart":
                props = run_systemctl_show(unit)
                current = props.get("Restart","no")
                new     = "no" if current != "no" else "always"
                subprocess.run(["systemctl","--user","set-property",unit,f"Restart={new}"])
        except Exception:
            logger.exception("Action %s failed on %s", label, unit)

    def refresh_scopes(self):
        self.model.remove_all()
        for unit, cli in self.scopes.items():
            props = run_systemctl_show(unit)
            mem   = int(props.get("MemoryCurrent",  "0"))
            swap  = int(props.get("MemorySwapCurrent","0"))
            cpu   = int(props.get("CPUUsageNSec",   "0"))
            state = props.get("ActiveState","unknown")

            if self.values_mode == "human":
                mem_s  = human_bytes(mem)
                swap_s = human_bytes(swap)
                cpu_s  = human_time(cpu)
            else:
                mem_s  = str(mem)
                swap_s = str(swap)
                cpu_s  = str(cpu)

            row = ScopeRow(unit, cli, mem_s, swap_s, cpu_s, state)
            self.model.append(row)

        logger.debug("Refreshed %d scopes", self.model.get_n_items())
        return True

# Application
class ScopeManagerApp(Gtk.Application):
    def __init__(self):
        super().__init__(application_id="org.systemd.ScopeManager")
        self.connect("activate", self.on_activate)

    def on_activate(self, app):
        win = ScopeManagerWindow(self)
        win.present()

def main():
    app = ScopeManagerApp()
    return app.run()

if __name__ == "__main__":
    import sys
    sys.exit(main())
