#!/usr/bin/env python3
# sbb-remote - sb-builder helper on the build VM
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

"""Helper run over ssh by the sb-builder GUI.

Every command prints one JSON object on stdout, except job-follow, which
streams raw log bytes. Jobs run detached (setsid, stdin=/dev/null), one at a
time under a flock, from a fixed table of kinds: the GUI sends a kind and
its arguments, this file validates them and builds the argv itself.
"""

import fcntl
import glob
import json
import os
import re
import shutil
import signal
import sqlite3
import subprocess
import sys
import time
import tomllib

VERSION = "0.1.0"
DEFAULTS = {
    "state_dir": "/var/lib/sb-builder",
    "repo_dir": "/var/lib/sbopkg/SBo-danix",
    "upstream_branch": "current",
    "slackrepo_db": "/var/lib/slackrepo/SBo-danix/database_SBo-danix_x86_64.sqlite3",
    "sbpull_script": "/usr/local/share/slackrepo_setup/sbpull.bash",
    "overlays": ["personal", "pentesting"],
}
FLAGS = ("--yes", "--no")
JOB_ID = re.compile(r"^\d{8}-\d{6}-[a-z]+(-\d+)?$")
PKG = re.compile(r"^[^/]+/([^/]+)/\1\.SlackBuild$")
ACTIVE = ("starting", "running")
START_GRACE = 10  # seconds a job may sit in "starting" before it counts as lost


class Error(Exception):
    pass


def load_config(path=None):
    path = path or os.environ.get("SBB_CONFIG", "/etc/sb-builder/config.toml")
    cfg = dict(DEFAULTS)
    try:
        with open(path, "rb") as f:
            data = tomllib.load(f)
    except FileNotFoundError:
        return cfg
    except tomllib.TOMLDecodeError as e:
        raise Error(f"config: {e}")
    for k, v in data.items():
        if k not in DEFAULTS:
            raise Error(f"config: {k}: unknown key")
        if type(v) is not type(DEFAULTS[k]):
            raise Error(f"config: {k}: expected {type(DEFAULTS[k]).__name__}")
        cfg[k] = v
    if not all(isinstance(o, str) for o in cfg["overlays"]):
        raise Error("config: overlays: expected list of strings")
    return cfg


# ── job kinds: (cfg, args) -> argv, raising Error on bad args ───────────────

def _flag(args):
    if len(args) != 1 or args[0] not in FLAGS:
        raise Error("expected exactly one of --yes, --no")
    return args[0]


def kind_sbpull(cfg, args):
    if not args or args[0] not in cfg["overlays"]:
        raise Error(f"sbpull: overlay must be one of {cfg['overlays']}")
    return ["bash", "-c", '. "$1" && sbpull "$2" "$3"', "_",
            cfg["sbpull_script"], args[0], _flag(args[1:])]


def kind_sbshadow(cfg, args):
    return ["sbshadow", _flag(args)]


KINDS = {"sbpull": kind_sbpull, "sbshadow": kind_sbshadow}


# ── job state ───────────────────────────────────────────────────────────────

def job_dir(cfg, jid):
    if not JOB_ID.match(jid or ""):
        raise Error(f"bad job id: {jid}")
    d = os.path.join(cfg["state_dir"], "jobs", jid)
    if not os.path.isdir(d):
        raise Error(f"no such job: {jid}")
    return d


def read_meta(d):
    with open(os.path.join(d, "meta.json")) as f:
        return json.load(f)


def write_meta(d, meta):
    tmp = os.path.join(d, "meta.json.tmp")
    with open(tmp, "w") as f:
        json.dump(meta, f)
    os.replace(tmp, os.path.join(d, "meta.json"))


def lock_file(cfg):
    return open(os.path.join(cfg["state_dir"], "lock"), "a")


def refresh(cfg, d, meta):
    """Mark an active job lost when no runner holds the lock.

    _run writes the final status before it releases the lock, so a free lock
    plus a status that is still active (re-read under the lock) means the
    runner died: killed, crashed, or the VM rebooted.
    """
    if meta["status"] not in ACTIVE:
        return meta
    if meta["status"] == "starting" and time.time() - meta["created"] < START_GRACE:
        return meta
    with lock_file(cfg) as lk:
        try:
            fcntl.flock(lk, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            return meta
        meta = read_meta(d)
        if meta["status"] in ACTIVE:
            meta["status"] = "lost"
            write_meta(d, meta)
    return meta


def all_jobs(cfg):
    out = []
    for d in sorted(glob.glob(os.path.join(cfg["state_dir"], "jobs", "*")), reverse=True):
        try:
            out.append(refresh(cfg, d, read_meta(d)))
        except (OSError, ValueError):
            pass  # a dir being created or removed right now
    return out


# ── commands ────────────────────────────────────────────────────────────────

def cmd_version(cfg):
    return {"version": VERSION, "overlays": cfg["overlays"]}


def cmd_jobs(cfg):
    # ponytail: logs and job dirs are kept forever; prune once disk usage shows on the dashboard
    return {"jobs": all_jobs(cfg)}


def cmd_job_start(cfg, kind=None, *args):
    if kind not in KINDS:
        raise Error(f"unknown kind: {kind}")
    argv = KINDS[kind](cfg, list(args))
    base = time.strftime("%Y%m%d-%H%M%S-") + kind
    jid, n = base, 1
    while True:
        d = os.path.join(cfg["state_dir"], "jobs", jid)
        try:
            os.mkdir(d)
            break
        except FileExistsError:
            n += 1
            jid = f"{base}-{n}"
    write_meta(d, {"id": jid, "kind": kind, "args": list(args), "argv": argv,
                   "status": "starting", "created": time.time()})
    open(os.path.join(d, "log"), "wb").close()
    subprocess.Popen([sys.executable, os.path.abspath(__file__), "_run", jid],
                     stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                     stderr=subprocess.DEVNULL, start_new_session=True)
    for _ in range(60):
        status = read_meta(d)["status"]
        if status == "busy":
            shutil.rmtree(d)
            return {"busy": next((m["id"] for m in all_jobs(cfg) if m["status"] == "running"), None)}
        if status != "starting":
            break
        time.sleep(0.05)
    return {"id": jid}


def cmd_run(cfg, jid=None):
    d = job_dir(cfg, jid)
    meta = read_meta(d)
    with lock_file(cfg) as lk:
        try:
            fcntl.flock(lk, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            meta["status"] = "busy"
            write_meta(d, meta)
            return None
        with open(os.path.join(d, "log"), "ab") as log:
            meta["start"] = time.time()
            try:
                p = subprocess.Popen(meta["argv"], stdin=subprocess.DEVNULL, stdout=log,
                                     stderr=subprocess.STDOUT, start_new_session=True)
            except OSError as e:
                log.write(f"sbb-remote: {e}\n".encode())
                rc = 127
            else:
                meta.update(pid=p.pid, status="running")
                write_meta(d, meta)
                rc = p.wait()
        meta.update(end=time.time(), exit=rc,
                    status="ok" if rc == 0 else "killed" if rc == -signal.SIGTERM else "failed")
        write_meta(d, meta)
    return None


def cmd_job_follow(cfg, jid=None, offset="0", out=None):
    d = job_dir(cfg, jid)
    out = out or sys.stdout.buffer
    with open(os.path.join(d, "log"), "rb") as f:
        f.seek(int(offset))
        while True:
            # status first, then read: the last read after the job ends drains the log
            done = refresh(cfg, d, read_meta(d))["status"] not in ACTIVE
            chunk = f.read()
            if chunk:
                out.write(chunk)
                out.flush()
            elif done:
                return None
            else:
                time.sleep(0.5)


def cmd_job_kill(cfg, jid=None):
    meta = read_meta(job_dir(cfg, jid))
    if meta["status"] != "running":
        raise Error(f"job not running: {jid}")
    try:
        os.killpg(meta["pid"], signal.SIGTERM)
    except ProcessLookupError:
        pass  # ended on its own meanwhile
    return {"ok": True}


# ── stats ───────────────────────────────────────────────────────────────────

def _try(fn):
    try:
        return fn()
    except Exception:
        return None


def _git(cfg, *args):
    return subprocess.run(["git", "-C", cfg["repo_dir"], *args],
                          capture_output=True, text=True, check=True).stdout


def _pkgs(root, pattern):
    """Names of the <pkg>/<pkg>.SlackBuild dirs matching pattern under root."""
    out = []
    for d in glob.glob(os.path.join(root, pattern, "")):
        name = os.path.basename(os.path.dirname(d))
        if os.path.isfile(os.path.join(d, name + ".SlackBuild")):
            out.append(name)
    return out


def _packages(cfg):
    r = cfg["repo_dir"]
    mine = {o: _pkgs(r, o + "/*") for o in cfg["overlays"]}
    tree = _git(cfg, "ls-tree", "-r", "--name-only", cfg["upstream_branch"]).splitlines()
    upstream = {m.group(1) for m in map(PKG.match, tree) if m}
    shadowed = upstream & {p for names in mine.values() for p in names}
    return {"total": len(_pkgs(r, "*/*")), **{o: len(n) for o, n in mine.items()},
            "shadowed": len(shadowed)}


def _slackrepo(cfg):
    db = sqlite3.connect(f"file:{cfg['slackrepo_db']}?mode=ro", uri=True)
    try:
        results = dict(db.execute("select result, count(*) from buildresults group by result"))
        built = db.execute("select count(distinct itemid) from revisions where dep = '/'").fetchone()[0]
    finally:
        db.close()
    return {"results": results, "built": built}


def _uptime():
    with open("/proc/uptime") as f:
        return float(f.read().split()[0])


def cmd_stats(cfg):
    dirs = {"repo": cfg["repo_dir"], "state": cfg["state_dir"]}
    return {
        "load": _try(os.getloadavg),
        "uptime": _try(_uptime),
        "disk_free": {k: _try(lambda p=p: shutil.disk_usage(p).free) for k, p in dirs.items()},
        "packages": _try(lambda: _packages(cfg)),
        "sync": {o: _try(lambda o=o: int(_git(cfg, "log", "-1", "--format=%ct", "--", o + "/")))
                 for o in cfg["overlays"]},
        "slackrepo": _try(lambda: _slackrepo(cfg)),
    }


COMMANDS = {
    "version": cmd_version,
    "stats": cmd_stats,
    "jobs": cmd_jobs,
    "job-start": cmd_job_start,
    "job-follow": cmd_job_follow,
    "job-kill": cmd_job_kill,
    "_run": cmd_run,
}


def main(argv):
    try:
        if not argv or argv[0] not in COMMANDS:
            names = ",".join(c for c in COMMANDS if not c.startswith("_"))
            raise Error(f"usage: sbb-remote {{{names}}} [args]")
        cfg = load_config()
        os.makedirs(os.path.join(cfg["state_dir"], "jobs"), exist_ok=True)
        result = COMMANDS[argv[0]](cfg, *argv[1:])
    except BrokenPipeError:
        # the GUI hung up on job-follow; silence the flush at exit too
        os.dup2(os.open(os.devnull, os.O_WRONLY), sys.stdout.fileno())
        return 0
    except Exception as e:
        print(json.dumps({"error": str(e) if isinstance(e, Error) else f"{type(e).__name__}: {e}"}))
        return 1
    if result is not None:
        print(json.dumps(result))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))