#!/usr/bin/env python3
"""scrd: the scr service daemon. Started by systemd; launches, watches and restarts services.
Every service (except the systemd runner) runs in a tmux session "scr-<name>" on $SCR_RUN/tmux/<user>.sock,
so `scr attach` and `scr send` work.

Two unix sockets (JSON lines, one request and one reply per connection):
  root.sock  mode 600, root only: every command.
  user.sock  any local user: status/check, plus logs/start/stop/restart for services running as the caller (SO_PEERCRED).
Config: $SCR_DIR/services.json. Logs: $SCR_LOG/<name>.log. Sockets: $SCR_RUN.
"""
import asyncio
import json
import os
import pwd
import shlex
import signal
import socket
import struct
import subprocess
import sys
import time
import urllib.request

CONF_DIR = os.environ.get("SCR_DIR", "/etc/scr")
RUN_DIR = os.environ.get("SCR_RUN", "/run/scr")
LOG_DIR = os.environ.get("SCR_LOG", "/var/log/scr")
CONF = os.path.join(CONF_DIR, "services.json")

USER_CMDS = {"status", "check", "logs", "start", "stop", "restart", "where", "send"}
OWN_ONLY = {"logs", "start", "stop", "restart", "where", "send"}  # on the user socket, only for services running as the caller
BACKOFF_MAX = 60  # seconds
MAX_FAILS = 5  # crashes in a row before "failed"
STABLE_AFTER = 30  # a run this long resets the crash counter
FIRST_ID = 1  # new IDs start at 0001; register with --id to choose one


class Service:
    def __init__(self, conf: dict):
        self.conf = conf
        self.pid = None  # pid of the process in the tmux pane
        self.state = "stopped"  # stopped | starting | running | backoff | failed
        self.started = 0.0
        self.restarts = 0
        self.fails = 0
        self.want = False  # should be running
        self.task = None  # supervise() task

    @property
    def name(self) -> str:
        return self.conf["name"]


services: dict = {}


def load() -> None:
    try:
        data = json.load(open(CONF))
    except FileNotFoundError:
        data = []
    for conf in data:
        services[conf["name"]] = Service(conf)


def save() -> None:
    os.makedirs(CONF_DIR, exist_ok=True)
    tmp = CONF + ".tmp"
    with open(tmp, "w") as f:
        json.dump([s.conf for s in services.values()], f, indent=2)
    os.replace(tmp, CONF)


def find(ref: str):
    if ref in services:
        return services[ref]
    for s in services.values():
        if s.conf["id"] == ref:
            return s
    return None


def command(conf: dict) -> list:
    runner, target = conf["runner"], conf["target"]
    extra = shlex.split(conf.get("args") or "")
    if runner == "uv":
        return ["uv", "run", target] + extra
    if runner == "npm":
        return ["npm", "run", conf["task"]] + (["--"] + extra if extra else [])
    if runner == "python":
        return [os.path.join(conf["venv"], "bin", "python") if conf.get("venv") else "python3", target] + extra
    if runner == "node":
        return ["node", target] + extra
    if runner == "java":
        return ["java"] + (["-Xmx" + conf["mem"]] if conf.get("mem") else []) + ["-jar", target] + extra
    if runner == "cargo":
        return ["cargo", "run", "--release"] + (["--"] + extra if extra else [])
    return ["/bin/sh", "-c", target]


def workdir(conf: dict) -> str:
    target = conf["target"]
    if conf.get("cwd"):
        return conf["cwd"]
    if conf["runner"] == "exec":
        return os.path.expanduser("~" + conf["user"])
    path = target if os.path.isdir(target) else os.path.dirname(os.path.abspath(target))
    return path or "/"


async def health(conf: dict) -> bool:
    spec = conf.get("check")
    if not spec:
        return True
    try:
        if spec.startswith("port:"):
            _, w = await asyncio.wait_for(asyncio.open_connection("127.0.0.1", int(spec[5:])), 2)
            w.close()
            return True
        if spec.startswith("unit:"):
            return await asyncio.get_running_loop().run_in_executor(None, unit_active, spec[5:])
        if spec.startswith("proc:"):  # proc:<user>:<command line pattern>
            _, user, pattern = spec.split(":", 2)
            rc = await asyncio.get_running_loop().run_in_executor(
                None, lambda: subprocess.run(["pgrep", "-u", user, "-f", "--", pattern], stdout=subprocess.DEVNULL).returncode)
            return rc == 0
        if spec.startswith("http"):
            def get():
                with urllib.request.urlopen(spec, timeout=2) as r:
                    return r.status < 500
            return await asyncio.get_running_loop().run_in_executor(None, get)
    except (OSError, ValueError, asyncio.TimeoutError):
        return False
    return False


RUNNERS = ("exec", "uv", "npm", "python", "node", "java", "cargo", "systemd", "external", "derive")


def unit_active(unit: str) -> bool:
    return subprocess.run(["systemctl", "is-active", "--quiet", unit]).returncode == 0


def unit_pid(unit: str):
    out = subprocess.run(["systemctl", "show", "-p", "MainPID", "--value", unit], capture_output=True, text=True).stdout
    return int(out.strip() or 0) or None


async def systemctl(*args: str) -> int:
    p = await asyncio.create_subprocess_exec("systemctl", *args)
    return await p.wait()


def drop_to(user: str):
    info = pwd.getpwnam(user)

    def fn():
        os.setsid()
        if os.getuid() == 0:
            os.initgroups(user, info.pw_gid)
            os.setgid(info.pw_gid)
            os.setuid(info.pw_uid)
    return fn, info


def session(conf: dict) -> str:
    return "scr-" + conf["name"]


def tmux_sock(user: str) -> str:
    return os.path.join(RUN_DIR, "tmux", user, "tmux.sock")


async def tmux(conf: dict, *args: str):
    """Run a tmux command on the service user's tmux server, as that user. Returns (returncode, stdout)."""
    fn, info = drop_to(conf["user"])
    env = {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin"), "HOME": info.pw_dir,
           "USER": conf["user"], "LOGNAME": conf["user"], "TERM": "xterm-256color"}
    p = await asyncio.create_subprocess_exec("tmux", "-S", tmux_sock(conf["user"]), *args, stdout=asyncio.subprocess.PIPE,
                                             stderr=asyncio.subprocess.PIPE, preexec_fn=fn, env=env)
    out, _ = await p.communicate()
    return p.returncode, out.decode(errors="replace")


async def pane_state(s: Service):
    """None if the session is gone, else (pane_dead, pane_pid)."""
    rc, out = await tmux(s.conf, "list-panes", "-t", session(s.conf), "-F", "#{pane_dead} #{pane_pid}")
    parts = out.split()
    if rc != 0 or len(parts) < 2:
        return None
    return parts[0] == "1", int(parts[1])


def log_path(conf: dict) -> str:
    return os.path.join(LOG_DIR, conf["name"] + ".log")


def clean(out: str) -> str:
    """Drop the blank padding lines tmux adds and its own 'Pane is dead' notice."""
    lines = [l.rstrip() for l in out.splitlines()]
    return "\n".join(l for l in lines if l and not l.startswith("Pane is dead")) + "\n"


async def dump_log(s: Service) -> None:
    """Append what the pane printed to the log file, then clear it so it is not written twice."""
    rc, out = await tmux(s.conf, "capture-pane", "-p", "-t", session(s.conf), "-S", "-20000")
    if rc == 0:
        with open(log_path(s.conf), "a") as f:
            f.write(clean(out))
        await tmux(s.conf, "clear-history", "-t", session(s.conf))


async def spawn(s: Service) -> None:
    conf = s.conf
    sockdir = os.path.dirname(tmux_sock(conf["user"]))  # owned by the service user so its tmux server can create the socket
    os.makedirs(sockdir, mode=0o700, exist_ok=True)
    if os.getuid() == 0:
        info = pwd.getpwnam(conf["user"])
        os.chown(sockdir, info.pw_uid, info.pw_gid)
    os.makedirs(LOG_DIR, exist_ok=True)
    with open(log_path(conf), "a") as f:
        f.write("--- scr start %s ---\n" % time.strftime("%F %T"))
    await tmux(conf, "start-server", ";", "set", "-g", "exit-empty", "off", ";", "set", "-g", "remain-on-exit", "on",
               ";", "set", "-g", "history-limit", "20000")
    name = session(conf)
    if await pane_state(s) is not None:  # a leftover (dead) pane: reuse the session
        rc, out = await tmux(conf, "respawn-pane", "-k", "-t", name, "-c", workdir(conf), *command(conf))
    else:
        rc, out = await tmux(conf, "new-session", "-d", "-s", name, "-x", "200", "-y", "50", "-c", workdir(conf),
                             *command(conf))
    if rc != 0:
        raise RuntimeError("tmux: " + out.strip())
    st = await pane_state(s)
    s.pid = st[1] if st else None
    s.started = time.time()


async def wait_ready(s: Service, timeout: float = 30) -> bool:
    """A dependency counts as ready once it runs and (if it has a check) the check passes."""
    end = time.time() + timeout
    while time.time() < end:
        if s.state == "failed":
            return False
        if s.state == "running" and await health(s.conf):
            return True
        await asyncio.sleep(0.5)
    return False


async def supervise(s: Service) -> None:
    delay = 1
    for dep in s.conf.get("after", []):
        d = find(dep)
        if d is None:
            s.state = "failed"
            return
        if not d.want:
            await start(d)
        s.state = "starting"
        if not await wait_ready(d):
            s.state = "failed"
            return
    if s.conf["runner"] == "systemd":  # systemd restarts it itself; we only start it and watch
        s.state = "starting"
        if await systemctl("start", s.conf["target"]) != 0:
            s.state = "failed"
            return
        s.started = time.time()
        while s.want:
            s.state = "running" if unit_active(s.conf["target"]) else "failed"
            await asyncio.sleep(2)
        s.state = "stopped"
        return
    while s.want:
        s.state = "starting"
        try:
            await spawn(s)
        except Exception as e:  # bad user, missing binary...
            print("spawn %s: %s" % (s.name, e), file=sys.stderr)
            s.fails += 1
        else:
            s.state = "running"
            while s.want:
                await asyncio.sleep(1)
                st = await pane_state(s)
                if st is None or st[0]:  # session killed or process exited
                    break
            if not s.want:
                break
            await dump_log(s)
            if time.time() - s.started > STABLE_AFTER:
                s.fails = 0
                delay = 1
            s.fails += 1
        if s.fails >= MAX_FAILS:
            s.state = "failed"
            return
        s.state = "backoff"
        s.restarts += 1
        await asyncio.sleep(delay)
        delay = min(delay * 2, BACKOFF_MAX)
    s.state = "stopped"


async def start(s: Service) -> None:
    if s.task and not s.task.done():
        return
    s.want, s.fails = True, 0
    s.task = asyncio.create_task(supervise(s))


async def stop(s: Service) -> None:
    s.want = False
    for other in services.values():  # dependents go down first
        if s.conf["id"] in other.conf.get("after", []) or s.name in other.conf.get("after", []):
            if other.want:
                await stop(other)
    if s.conf["runner"] == "systemd":
        await systemctl("stop", s.conf["target"])
    if s.conf["runner"] != "systemd":
        st = await pane_state(s)
        if st and not st[0]:
            for sig, wait in ((signal.SIGTERM, int(s.conf.get("stop_timeout") or 15) * 2), (signal.SIGKILL, 4)):
                try:
                    os.killpg(st[1], sig)
                except ProcessLookupError:
                    break
                for _ in range(wait):
                    await asyncio.sleep(0.5)
                    st = await pane_state(s)
                    if st is None or st[0]:
                        break
                else:
                    continue
                break
        if st is not None:
            await dump_log(s)
            await tmux(s.conf, "kill-session", "-t", session(s.conf))
        s.pid = None
    if s.task and not s.task.done():  # the process is gone; end the supervisor (it may be sleeping in backoff)
        s.task.cancel()
        try:
            await s.task
        except asyncio.CancelledError:
            pass
    s.state = "stopped"


async def describe(s: Service) -> dict:
    wrapped = s.conf["runner"] == "systemd"
    if s.conf["runner"] in ("external", "derive"):  # not run by scr: state is just its health check
        up = await health(s.conf)
        if s.conf["runner"] == "derive":  # started by another service: only up while that one is
            parent = find(s.conf["target"])
            up = up and parent is not None and (await describe(parent))["state"] == "running"
        return {"id": s.conf["id"], "name": s.name, "label": s.conf.get("label"), "state": "running" if up else "down", "enabled": False, "pid": None,
                "uptime": 0, "restarts": 0, "healthy": up, "runner": s.conf["runner"], "user": s.conf["user"], "after": [],
                "check": s.conf.get("check")}
    state = s.state
    if wrapped:  # reflect the real unit even if scr never started it
        state = "running" if unit_active(s.conf["target"]) else ("failed" if s.want else "stopped")
    running = state == "running" and (unit_active(s.conf["target"]) if wrapped else s.pid is not None)
    return {"id": s.conf["id"], "name": s.name, "label": s.conf.get("label"), "state": state, "enabled": s.conf.get("enabled", False),
            "pid": (unit_pid(s.conf["target"]) if wrapped else s.pid) if running else None, "uptime": int(time.time() - s.started) if running and s.started else 0,
            "restarts": s.restarts, "healthy": await health(s.conf) if running else False,
            "runner": s.conf["runner"], "user": s.conf["user"], "after": s.conf.get("after", []),
            "check": s.conf.get("check")}


def tail(name: str, n: int) -> str:
    try:
        with open(os.path.join(LOG_DIR, name + ".log"), "rb") as f:
            return b"".join(f.readlines()[-n:]).decode(errors="replace")
    except FileNotFoundError:
        return ""


def next_id() -> str:
    used = [int(s.conf["id"]) for s in services.values()]
    return "%04d" % (max([FIRST_ID - 1] + used) + 1)


async def handle(req: dict, caller_uid: int, root: bool) -> dict:
    cmd, a = req.get("cmd"), req.get("args", {})
    if cmd not in USER_CMDS and not root:
        return {"ok": False, "error": "needs the root socket"}
    if cmd == "status":
        return {"ok": True, "services": [await describe(s) for s in sorted(services.values(), key=lambda x: x.conf["id"])]}
    if cmd == "register":
        name = a["name"]
        if name in services:
            return {"ok": False, "error": "already registered"}
        runner = a.get("runner", "exec")
        if runner not in RUNNERS or (runner == "npm" and not a.get("task")):
            return {"ok": False, "error": "runner must be one of %s (npm needs --task)" % ", ".join(RUNNERS)}
        try:
            pwd.getpwnam(a.get("user", "root"))
        except KeyError:
            return {"ok": False, "error": "unknown user"}
        if runner == "derive" and find(a["target"]) is None:
            return {"ok": False, "error": "derive needs an existing parent service (name or ID) as the target"}
        sid = a.get("id") or next_id()
        if any(x.conf["id"] == sid for x in services.values()):
            return {"ok": False, "error": "id in use"}
        conf = {"id": sid, "name": name, "target": a["target"], "runner": runner, "user": a.get("user", "root"),
                "label": a.get("label"), "stop_timeout": a.get("stop_timeout"), "task": a.get("task"), "args": a.get("args"), "venv": a.get("venv"), "mem": a.get("mem"), "after": a.get("after", []), "check": a.get("check"), "cwd": a.get("cwd"),
                "enabled": False}
        services[name] = Service(conf)
        save()
        return {"ok": True, "id": sid}
    s = find(a.get("service", ""))
    if s is None:
        return {"ok": False, "error": "no such service"}
    if not root and cmd in OWN_ONLY and pwd.getpwuid(caller_uid).pw_name != s.conf["user"]:
        return {"ok": False, "error": "that service does not run as you"}
    if cmd == "check":
        return {"ok": True, **await describe(s)}
    if s.conf["runner"] in ("external", "derive") and cmd in ("start", "stop", "restart", "enable", "disable", "logs", "where", "send"):
        return {"ok": False, "error": ("started by %s, do not start it separately" % s.conf["target"]) if s.conf["runner"] == "derive"
                else "external service: scr only checks it, it is managed elsewhere"}
    if cmd in ("where", "send") and s.conf["runner"] == "systemd":
        return {"ok": False, "error": "systemd-wrapped services have no tmux session"}
    if cmd == "where":
        if await pane_state(s) is None:
            return {"ok": False, "error": "not running (no session)"}
        return {"ok": True, "sock": tmux_sock(s.conf["user"]), "session": session(s.conf)}
    if cmd == "send":
        if await pane_state(s) is None:
            return {"ok": False, "error": "not running (no session)"}
        await tmux(s.conf, "send-keys", "-t", session(s.conf), "-l", "--", a["text"])
        if not a.get("no_enter"):
            await tmux(s.conf, "send-keys", "-t", session(s.conf), "Enter")
        return {"ok": True, "id": s.conf["id"], "name": s.name}
    if cmd == "logs":
        if s.conf["runner"] == "systemd":
            out = subprocess.run(["journalctl", "-u", s.conf["target"], "-n", str(int(a.get("lines", 50))), "--no-pager"],
                                 capture_output=True, text=True)
            return {"ok": True, "text": out.stdout or out.stderr}
        n = int(a.get("lines", 50))
        rc, out = await tmux(s.conf, "capture-pane", "-p", "-t", session(s.conf), "-S", str(-n)) if s.pid else (1, "")
        return {"ok": True, "text": (tail(s.name, n) + clean(out)) if rc == 0 else tail(s.name, n)}
    if cmd == "start":
        await start(s)
    elif cmd == "stop":
        await stop(s)
    elif cmd == "restart":
        await stop(s)
        await start(s)
    elif cmd == "enable":
        s.conf["enabled"] = True
        save()
        await start(s)
    elif cmd == "disable":
        s.conf["enabled"] = False
        save()
        await stop(s)
    elif cmd == "remove":
        await stop(s)
        del services[s.name]
        save()
    return {"ok": True, "id": s.conf["id"], "name": s.name}


async def serve(path: str, mode: int, root: bool) -> None:
    if os.path.exists(path):
        os.unlink(path)

    async def conn(r, w):
        try:
            sock = w.get_extra_info("socket")
            uid = struct.unpack("3i", sock.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i")))[1]
            if root and uid != 0 and uid != os.getuid():
                reply = {"ok": False, "error": "root socket is root only"}
            else:
                reply = await handle(json.loads(await asyncio.wait_for(r.readline(), 5)), uid, root)
        except Exception as e:
            reply = {"ok": False, "error": "%s: %s" % (type(e).__name__, e)}
        try:
            w.write((json.dumps(reply) + "\n").encode())
            await w.drain()
        except ConnectionError:  # client left without waiting
            pass
        w.close()

    server = await asyncio.start_unix_server(conn, path)
    os.chmod(path, mode)
    await server.serve_forever()


async def main() -> None:
    os.makedirs(RUN_DIR, mode=0o755, exist_ok=True)
    load()
    loop = asyncio.get_running_loop()
    done = asyncio.Event()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, done.set)
    servers = [asyncio.create_task(serve(os.path.join(RUN_DIR, "root.sock"), 0o600, True)),
               asyncio.create_task(serve(os.path.join(RUN_DIR, "user.sock"), 0o666, False))]
    for s in services.values():
        if s.conf.get("enabled") and s.conf["runner"] not in ("external", "derive"):
            await start(s)
    await done.wait()
    for s in list(services.values()):
        if s.want:
            await stop(s)
    for user in {s.conf["user"] for s in services.values() if s.conf["runner"] != "systemd"}:
        await tmux({"user": user}, "kill-server")
    for t in servers:
        t.cancel()


if __name__ == "__main__":
    asyncio.run(main())
