#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import os
import re
import subprocess
import time
import json
import curses


DefaultConfigPath = "/usr/local/etc/sadm.conf"


# ================================================================================================================


class TFormat:
    def __init__(self, S=""):
        self.Str   = S
        self.Color = 0
        self.Prop  = 0
        self.Algn  = "<"
        self.Wdth  = 0
        self.Pcut  = 0
        self.Pre   = ""
    def NoColor(self):      self.Color = 0;     return self
    def White(self):        self.Color = 29;    return self
    def Black(self):        self.Color = 30;    return self
    def Red(self):          self.Color = 31;    return self
    def Green(self):        self.Color = 32;    return self
    def Brown(self):        self.Color = 33;    return self
    def Blue(self):         self.Color = 34;    return self
    def Purple(self):       self.Color = 35;    return self
    def Cyan(self):         self.Color = 36;    return self
    def LightGray(self):    self.Color = 37;    return self
    def Normal(self):       self.Prop  = 0;     return self
    def Bold(self):         self.Prop  = 1;     return self
    def Dark(self):         self.Prop  = 2;     return self
    def Italic(self):       self.Prop  = 3;     return self
    def Underlined(self):   self.Prop  = 4;     return self
    def Blink(self):        self.Prop  = 5;     return self
    def Right(self):        self.Algn  = ">";   return self
    def Left(self):         self.Algn  = "<";   return self
    def Width(self, W):     self.Wdth  = W;     return self
    def Cut(self, C):       self.Pcut  = C;     return self
    def Prefix(self, P):    self.Pre   = P;     return self
    def __call__(self, S):  self.Str   = S;     return self
    def __str__(self):
        Str = self.Str
        if len(self.Pre) > 0:
            Str = self.Pre + Str
        if self.Pcut != 0 and len(Str) > self.Pcut:
            Str = Str[:self.Pcut - 2] + ".."
        if self.Wdth != 0:
            Str = "{S:{A}{W}}".format(S=Str, A=self.Algn, W=self.Wdth)
        if self.Color > 0:
            Str = "{R}{P};{C}m{S}{R}0m".format(R="\033[", P=self.Prop, C=self.Color, S=Str)
        return Str


class TScreen:
    def __init__(self):
        self.Line = 0
        self.Screen = curses.initscr()
        curses.noecho()
        curses.cbreak()

    def __call__(self, *args, sep=" "):
        self.Screen.addstr(self.Line, 0, sep.join(map(str, args)))
        self.Line += 1

    def Update(self):
        self.Line = 0
        self.Screen.refresh()

    def __del__(self):
        curses.echo()
        curses.nocbreak()
        curses.endwin()


# ================================================================================================================


class Exec:
    def __init__(self, Command, Env={"LANG": "C"}, ErrToOut=False, Pipe=False):
        self.Pid = -1
        self.Status = 1
        self.Out = ""
        self.Err = ""
        try:
            self.P = subprocess.Popen(
                Command,
                env=Env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT if ErrToOut else subprocess.PIPE,
                close_fds=True,
                preexec_fn=lambda: os.setpgid(os.getpid(), os.getpid())
            )
            self.Pid = self.P.pid
            if Pipe:
                self.Stdout = self.P.stdout
                self.Stderr = self.P.stderr
            else:
                try:
                    self.Out = self.P.stdout.read().strip().decode("utf-8")
                    self.Err = "" if ErrToOut else self.P.stderr.read().strip().decode("utf-8")
                except KeyboardInterrupt:
                    self.Stop()
                self.P.wait()
                self.Status = self.P.returncode
        except OSError as e:
            print("Got exception running {}: {}".format(Command, e))
            exit(1)

    def IsRunning(self):
        Poll = self.P.poll()
        if Poll is None:
            return True
        self.Status = Poll
        return False

    def Stop(self):
        if self.Pid > 0:
            try:
                os.killpg(self.Pid, 2)
                self.P.wait()
                self.Status = self.P.returncode
            except ProcessLookupError:
                pass

    def Wait(self):
        if self.Pid > 0:
            try:
                self.P.wait()
                self.Status = self.P.returncode
            except KeyboardInterrupt:
                self.Stop()


def GetFileData(FileName, Data=""):
    try:
        File = open(FileName, 'r')
        Data = File.read()
        File.close()
    except (IOError, OSError) as e:
        if e.errno == 2 or e.errno == 3:
            pass
        else:
            raise
    return Data


def PID(Pid):
    try:
        Pid = int(Pid)
        return None if Pid == 0 else Pid
    except ValueError:
        return None


def PidCmdLine(Pid):
    try:
        Line = GetFileData("/proc/{}/cmdline".format(Pid))
        if len(Line) > 0:
            return Line.split('\x00')[0]
    except (ValueError, IndexError, KeyError):
        pass
    return None


def GetProcDict():
    ProcDict = {None: {
        "Command": "",
        "Time":    "",
        "Cpu":     "",
        "Mem":     "",
        "Children": []
    }}
    PProcDict = {}
    for Line in Exec(["/bin/ps", "h", "-Awwo", "pid,ppid,etime,pcpu,rsz,args"]).Out.split('\n'):
        ProcArray = Line.split(None, 5)
        Pid  = int(ProcArray[0])
        PPid = int(ProcArray[1])
        ProcDict[Pid] = {
            "Command": ProcArray[5].strip(),
            "Time":    ProcArray[2],
            "Cpu":     "{:5.1f}".format(float(ProcArray[3])),
            "Mem":     "{:7.1f}".format(float(ProcArray[4])/1024.0)
        }
        if PPid not in PProcDict:
            PProcDict[PPid] = []
        PProcDict[PPid].append(Pid)
    for Pid in ProcDict:
        ProcDict[Pid]["Children"] = PProcDict.get(Pid, [])
    return ProcDict


def SetServiceStatus(ServiceList):
    Cmd = [
        "/bin/systemctl",
        "show",
        "--no-pager",
        "--property=ActiveState,SubState,MainPID,ExecMainCode,ExecMainStatus"
    ] + [
        ServiceDict["Service"] for ServiceDict in ServiceList
    ]
    for Idx, Service in enumerate(Exec(Cmd).Out.split('\n\n')):
        StatusDict = {}
        for L in Service.split('\n'):
            L = L.split('=')
            if len(L) > 1:
                StatusDict[L[0]] = L[1].strip()
        ServiceDict = ServiceList[Idx]
        ServiceDict["Pid"] = PID(StatusDict.get("MainPID"))
        ServiceDict["Status"] = StatusDict.get("ActiveState")
        ServiceDict["SubStatus"] = StatusDict.get("SubState")


# ================================================================================================================


class TServices:
    def __init__(self, ConfigPath=DefaultConfigPath, PatternList=[]):
        ServiceExistDict = {}
        def _ServiceExists(Service):
            return ServiceExistDict.setdefault(
                Service,
                   os.path.exists("/etc/systemd/system/{}.service".format(Service))
                or os.path.exists("/lib/systemd/system/{}.service".format(Service))
                or os.path.exists("/etc/init/{}.conf".format(Service))
                or os.path.exists("/etc/init.d/{}".format(Service))
            )
        try:
            AllServices = [{
                "Service":          S["Service"],
                "Process":          S.get("Process", None),
                "ProcessFilter":    S.get("ProcessFilter", None),
                "Pid":              None,
                "Status":           None,
                "SubStatus":        None
            } if isinstance(S, dict) else {
                "Service":          S,
                "Process":          None,
                "ProcessFilter":    None,
                "Pid":              None,
                "Status":           None,
                "SubStatus":        None
            } for S in json.loads(GetFileData(ConfigPath)).get("Services", [])]
        except Exception as e:
            print("Failed to load config from {}: {}".format(ConfigPath, e))
            exit(1)
        if not os.path.exists("/bin/systemctl"):
            print("Systemctl not found")
            exit(1)
        self.MyUID = os.getuid()
        self.SudoExists = os.path.isfile("/usr/bin/sudo")
        self.ExactError = None
        self.LooseError = None
        Size = len(AllServices)
        ExactList = [False]*Size
        LooseList = [False]*Size
        for Pattern in PatternList:
            ModPattern = Pattern.replace("%", ".*")
            ReExact = re.compile("^{}$".format(ModPattern))
            ReLoose = re.compile(ModPattern)
            ExactFound = []
            LooseFound = []
            for Idx, ServiceDict in enumerate(AllServices):
                Service = ServiceDict["Service"]
                if (Pattern == Service or ReExact.search(Service)) and _ServiceExists(Service):
                    ExactList[Idx] = True
                    ExactFound.append(Idx)
                if (Pattern in Service or ReLoose.search(Service)) and _ServiceExists(Service):
                    LooseList[Idx] = True
                    LooseFound.append(Idx)
            if not ExactFound:
                # if only one service matches, use it as exact match
                if len(LooseFound) == 1:
                    ExactList[LooseFound[0]] = True
                elif not self.ExactError:
                    self.ExactError = "No service match '{}' exactly!".format(Pattern)
                    if LooseFound:
                        self.ExactError += " Loose match: {}".format(", ".join([AllServices[Idx]["Service"] for Idx in LooseFound]))
            if not self.LooseError and not LooseFound:
                self.LooseError = "No service match '{}'!".format(Pattern)
        if len(PatternList) == 0:
            LooseList = [_ServiceExists(D["Service"]) for D in AllServices]
            self.ExactError = "No pattern specified!"
        self.ExactList = [AllServices[Idx] for Idx, Ok in enumerate(ExactList) if Ok]
        self.LooseList = [AllServices[Idx] for Idx, Ok in enumerate(LooseList) if Ok]

    def List(self):
        if self.LooseError:
            print(self.LooseError)
            exit(0)
        print(" ".join([S["Service"] for S in self.LooseList]))

    def Status(self, ServiceList=[], Width=None, TopMode=False):
        if self.LooseError:
            print(self.LooseError)
            exit(0)
        if len(ServiceList) == 0:
            ServiceList = self.LooseList
        IsAtty  = sys.stdout.isatty()
        Width   = Width if Width else os.get_terminal_size()[0] if IsAtty else 1024
        WSrv    = 24 if Width >= 132 else 22
        WSta    = 16 if Width >= 132 else 13
        if not IsAtty:
            TopMode = False
        Print   = TScreen() if TopMode else print
        FSrv    = TFormat().Width(WSrv).Cut(WSrv).Left()
        FSta    = TFormat().Width(WSta).Cut(WSta).Left()
        FPid    = TFormat().Width(7).Right()
        FTime   = TFormat().Width(12).Right()
        FCpu    = TFormat().Width(7 if Width >= 132 else 5).Right()
        FMem    = TFormat().Width(9 if Width >= 132 else 7).Right()
        FCmd    = TFormat().Prefix(" ").Cut(max(Width - (82 if Width >= 132 else 73), 50))
        while True:
            if IsAtty:
                if TopMode:
                    Print(FSrv("SERVICE"),
                          FSta("STATUS"),
                          FPid("PID"),
                          FTime("TIME"),
                          FCpu("CPU"),
                          FMem("MEM"),
                          FCmd("COMMAND"))
                else:
                    Print(FSrv("SERVICE").Bold().Black(),
                          FSta("STATUS").Bold().Black(),
                          FPid("PID").Bold().Black(),
                          FTime("TIME").Bold().Black(),
                          FCpu("CPU").Bold().Black(),
                          FMem("MEM").Bold().Black(),
                          FCmd("COMMAND").Bold().Black())
                    FSrv.White()
                    FSta.Normal().White()
                    FPid.Blue()
                    FTime.Normal().White()
                    FCpu.Red()
                    FMem.Green()
                    FCmd.Normal().White()
            ProcDict = GetProcDict()
            SetServiceStatus(ServiceList)
            for ServiceDict in ServiceList:
                Status = ServiceDict["Status"] + (("/" + ServiceDict["SubStatus"]) if ServiceDict["SubStatus"] else "")
                Pid = ServiceDict["Pid"]
                Proc = ProcDict.get(Pid, ProcDict[None])
                Print(FSrv(ServiceDict["Service"]),
                      FSta(Status),
                      FPid(str(Pid) if Pid is not None else ""),
                      FTime(Proc["Time"]),
                      FCpu(Proc["Cpu"]),
                      FMem(Proc["Mem"]),
                      FCmd(Proc["Command"]))
                for ChildPid in Proc["Children"]:
                    Proc = ProcDict.get(ChildPid, ProcDict[None])
                    Print(FSrv("-"),
                          FSta("-"),
                          FPid(str(ChildPid)),
                          FTime(Proc["Time"]),
                          FCpu(Proc["Cpu"]),
                          FMem(Proc["Mem"]),
                          FCmd(Proc["Command"]))
            if TopMode:
                Print.Update()
                time.sleep(1)
            else:
                break

    def Stop(self):
        if self.ExactError:
            print(self.ExactError)
            exit(0)
        SetServiceStatus(self.ExactList)
        for ServiceDict in self.ExactList:
            Pid = ServiceDict["Pid"]
            Service = ServiceDict["Service"]
            Status = ServiceDict["Status"]
            if Status in ["active", "activating"]:
                ProcName = PidCmdLine(Pid)
                print("Trying to stop {} (process {})".format(Service, ProcName))
                if self.MyUID != 0 and not self.SudoExists:
                    print("Must use sudo to stop services, but no sudo is available!")
                    exit(1)
                os.system("/usr/bin/sudo -n /bin/systemctl stop {} >/dev/null 2>&1 &".format(Service))
                if ProcName is not None:
                    for c in range(20):
                        time.sleep(0.5)
                        if PidCmdLine(Pid) != ProcName:
                            break
                        print("Waiting for {:.1f} seconds ({})...".format(c/2.0, Service))
                    else:
                        print("Service {} is too slow to stop without help, kill it!".format(Service))
                        os.system("/usr/bin/sudo -n /bin/kill -9 {}".format(Pid))
                        time.sleep(0.5)
            else:
                print("Need to stop {}, but it is in state {}/{}!".format(Service, Status, ServiceDict["SubStatus"]))

    def Start(self):
        if self.ExactError:
            print(self.ExactError)
            exit(0)
        SetServiceStatus(self.ExactList)
        for ServiceDict in self.ExactList[::-1]:
            Service = ServiceDict["Service"]
            Status = ServiceDict["Status"]
            if Status in ["inactive", "deactivating", "failed"]:
                print("Trying to start {}".format(Service))
                if self.MyUID != 0 and not self.SudoExists:
                    print("Must use sudo to start services, but no sudo is available!")
                    exit(1)
                os.system("/usr/bin/sudo -n /bin/systemctl start {} >/dev/null 2>&1 &".format(Service))
                time.sleep(0.5)
                ServiceList = [ServiceDict]
                SetServiceStatus(ServiceList)
                self.Status(ServiceList=ServiceList)
            else:
                print("Need to start {}, but it is in state {}/{}!".format(Service, Status, ServiceDict["SubStatus"]))

    def Restart(self):
        if self.ExactError:
            print(self.ExactError)
            exit(0)
        self.Stop()
        SetServiceStatus(self.ExactList)
        self.Start()

    def Logs(self, Since, Until, Follow=False):
        if self.LooseError:
            print(self.LooseError)
            exit(0)
        Services = [S["Service"] for S in self.LooseList]
        if len(Services) != 1:
            print("Can show logs only for one service! (Matching services: {})".format(", ".join(Services)))
            exit(0)
        Cmd = [
            "/usr/bin/sudo", "/bin/journalctl",
            "-q",
            "-o", "short-precise",
            "-u", Services[0],
            "-S", Since
        ]
        if Until != "inf":
            Cmd += ["-U", Until]
        if Follow:
            Cmd.append("-f")
        R = Exec(Cmd, ErrToOut=True, Pipe=True)
        try:
            while R.IsRunning():
                print(R.Stdout.readline().decode("utf-8"), end="")
        except KeyboardInterrupt:
            os.system("/usr/bin/sudo -n /bin/kill -2 {} >/dev/null 2>&1".format(R.Pid))
        R.Wait()


# ================================================================================================================


class Arguments:
    def __init__(self):
        self.Flags = {"-f": False, "-t": False}
        self.Opts = {"-w": 0, "-S": "-30m", "-U": "inf"}
        self.Cmd = "status"
        self.Patterns = []
        self.Commands = {"list", "status", "stop", "start", "restart", "logs"}

    def Help(self, Err=""):
        Err = "{}\n\n".format(Err) if Err else ""
        print("""{}usage: {} [-h] [-t] [-w WIDTH] [-S SINCE] [-U UNTIL] [-f] [CMD] [PATTERN [PATTERN ...]]

CMD and PATTERN could be swapped.
Most convenient time format for SINCE and UNTIL is
  -2m10s,
  +1m20s (makes sense only for UNTIL time with -f option),
  -1d3h,
  etc.

positional arguments:
  CMD         command [{}] (default="status")
  PATTERN     service patterns

optional arguments:
  -h, --help  show this help message and exit
  -t          top mode watch for services
  -w WIDTH    set console width
  -S SINCE    show logs since that time (default="{}")
  -U UNTIL    show logs until that time (default="{}")
  -f          follow logs
""".format(
            Err,
            os.path.basename(sys.argv[0]),
            "|".join(self.Commands),
            self.Opts["-S"],
            self.Opts["-U"]
        ))
        exit(0)

    def Parse(self):
        Args = iter(sys.argv[1:])
        try:
            while True:
                Arg = next(Args)
                if Arg in ["-h", "--help"]:
                    self.Help()
                elif Arg in self.Flags:
                    self.Flags[Arg] = True
                elif Arg in self.Opts:
                    try:
                        self.Opts[Arg] = self.Opts[Arg].__class__(next(Args))
                    except ValueError as e:
                        self.Help("Bad format for '{}': {}".format(Arg, e))
                elif Arg in self.Commands:
                    self.Cmd = Arg
                else:
                    self.Patterns.append(Arg)
        except StopIteration:
            pass

    def Get(self, Arg):
        Arg = "-{}".format(Arg)
        if Arg in self.Flags:
            return self.Flags[Arg]
        if Arg in self.Opts:
            return self.Opts[Arg]
        return None


# ================================================================================================================


def Main():
    Args = Arguments()
    Args.Parse()

    Services = TServices(PatternList=Args.Patterns)

    if Args.Cmd == "list":
        Services.List()
    elif Args.Cmd == "status":
        Services.Status(Width=Args.Get("w"), TopMode=Args.Get("t"))
    elif Args.Cmd == "stop":
        Services.Stop()
    elif Args.Cmd == "start":
        Services.Start()
    elif Args.Cmd == "restart":
        Services.Restart()
    elif Args.Cmd == "logs":
        Services.Logs(Args.Get("S"), Args.Get("U"), Follow=Args.Get("f"))
    else:
        Args.Help()


if __name__ == '__main__':
    Main()
