tinytaskmanager.py

#!/usr/bin/env python3

# tinytaskmanager
# Tiny task manager for Linux, MacOS and Unix-like systems.
# Written as a single Python script.
# MIT License
# Copyright (c) 2022 Yuri Escalianti <yuriescl@gmail.com>
# Homepage: https://github.com/yuriescl/tinytaskmanager

from datetime import datetime
from fcntl import LOCK_EX, LOCK_UN, lockf
import io
from io import SEEK_END, SEEK_SET
import json
from multiprocessing.dummy import Pool as ThreadPool
import os
from os.path import abspath, join
from pathlib import Path
import re
import shlex
from shutil import get_terminal_size, rmtree
from signal import SIGINT, SIGKILL, SIGTERM, Signals, signal
from subprocess import DEVNULL, Popen, check_output
from sys import argv, exit, stderr, stdout, version_info
import tempfile
import time
from time import sleep
from typing import Dict, List, Optional, Tuple, Union

if version_info[0] < 3 or version_info[1] < 8:
    raise Exception("Python >=3.8 is required to run this program")


LOCK_FILE_NAME = "lock"
CACHE_DIR = Path.home() / ".tinytaskmanager"
LOCK_PATH = Path(CACHE_DIR / LOCK_FILE_NAME)

RESERVED_FILE_NAMES = [LOCK_FILE_NAME]

VERSION = "0.14.0"
BUSY_LOOP_INTERVAL = 0.1  # seconds
TIMESTAMP_FMT = "%Y%m%d%H%M%S"

TERMINATE = False

Task = dict


class TtmException(Exception):
    pass


class Tailer(object):
    """
    Code obtained from https://github.com/GreatFruitOmsk/tailhead/blob/master/tailhead/__init__.py
    Copyright (c) 2012 Mike Thornton
    Implements tailing and heading functionality like GNU tail and head
    commands.
    """

    LINE_TERMINATORS = (b"\r\n", b"\n", b"\r")

    def __init__(self, file, read_size=1024, end=False):
        if not isinstance(file, io.IOBase) or isinstance(file, io.TextIOBase):
            raise ValueError("io object must be in the binary mode")

        self.read_size = read_size
        self.file = file

        if end:
            self.file.seek(0, SEEK_END)

    def splitlines(self, data):
        return re.split(b"|".join(self.LINE_TERMINATORS), data)

    def read(self, read_size=-1):
        read_str = self.file.read(read_size)
        return len(read_str), read_str

    def prefix_line_terminator(self, data):
        for t in self.LINE_TERMINATORS:
            if data.startswith(t):
                return t

        return None

    def suffix_line_terminator(self, data):
        for t in self.LINE_TERMINATORS:
            if data.endswith(t):
                return t

        return None

    def seek_next_line(self):
        where = self.file.tell()
        offset = 0

        while True:
            data_len, data = self.read(self.read_size)
            data_where = 0

            if not data_len:
                break

            # Consider the following example: Foo\r | \nBar where " | " denotes current position,
            # 'Foo\r' is the read part and '\nBar' is the remaining part.
            # We should completely consume terminator "\r\n" by reading one extra byte.
            if b"\r\n" in self.LINE_TERMINATORS and data[-1] == b"\r"[0]:
                terminator_where = self.file.tell()
                terminator_len, terminator_data = self.read(1)

                if terminator_len and terminator_data[0] == b"\n"[0]:
                    data_len += 1
                    data += b"\n"
                else:
                    self.file.seek(terminator_where)

            while data_where < data_len:
                terminator = self.prefix_line_terminator(data[data_where:])
                if terminator:
                    self.file.seek(where + offset + data_where + len(terminator))
                    return self.file.tell()
                else:
                    data_where += 1

            offset += data_len
            self.file.seek(where + offset)

        return -1

    def seek_previous_line(self):
        where = self.file.tell()
        offset = 0

        while True:
            if offset == where:
                break

            read_size = self.read_size if self.read_size <= where else where
            self.file.seek(where - offset - read_size, SEEK_SET)
            data_len, data = self.read(read_size)

            # Consider the following example: Foo\r | \nBar where " | " denotes current position,
            # '\nBar' is the read part and 'Foo\r' is the remaining part.
            # We should completely consume terminator "\r\n" by reading one extra byte.
            if b"\r\n" in self.LINE_TERMINATORS and data[0] == b"\n"[0]:
                terminator_where = self.file.tell()
                if terminator_where > data_len + 1:
                    self.file.seek(where - offset - data_len - 1, SEEK_SET)
                    terminator_len, terminator_data = self.read(1)

                    if terminator_data[0] == b"\r"[0]:
                        data_len += 1
                        data = b"\r" + data

                    self.file.seek(terminator_where)

            data_where = data_len

            while data_where > 0:
                terminator = self.suffix_line_terminator(data[:data_where])
                if terminator and offset == 0 and data_where == data_len:
                    # The last character is a line terminator that finishes current line. Ignore it.
                    data_where -= len(terminator)
                elif terminator:
                    self.file.seek(where - offset - (data_len - data_where))
                    return self.file.tell()
                else:
                    data_where -= 1

            offset += data_len

        if where == 0:
            # Nothing more to read.
            return -1
        else:
            # Very first line.
            self.file.seek(0)
            return 0

    def tail(self, lines=10):
        self.file.seek(0, SEEK_END)

        for i in range(lines):
            if self.seek_previous_line() == -1:
                break

        data = self.file.read()

        for t in self.LINE_TERMINATORS:
            if data.endswith(t):
                # Only terminators _between_ lines should be preserved.
                # Otherwise terminator of the last line will be treated as separtaing line and empty line.
                data = data[: -len(t)]
                break

        if data:
            return self.splitlines(data)
        else:
            return []

    def head(self, lines=10):
        if lines < 0:
            self.file.seek(0, SEEK_END)
            for i in range(-lines):
                if self.seek_previous_line() == -1:
                    break
        else:
            self.file.seek(0)
            for i in range(lines):
                if self.seek_next_line() == -1:
                    break

        end_pos = self.file.tell()

        self.file.seek(0)
        data = self.file.read(end_pos)

        for t in self.LINE_TERMINATORS:
            if data.endswith(t):
                # Only terminators _between_ lines should be preserved.
                # Otherwise terminator of the last line will be treated as separtaing line and empty line.
                data = data[: -len(t)]
                break

        if data:
            return self.splitlines(data)
        else:
            return []

    def follow(self):
        trailing = True

        while True:
            where = self.file.tell()

            if where > os.fstat(self.file.fileno()).st_size:
                # File was truncated.
                where = 0
                self.file.seek(where)

            line = self.file.readline()

            if line:
                if trailing and line in self.LINE_TERMINATORS:
                    # This is just the line terminator added to the end of the file
                    # before a new line, ignore.
                    trailing = False
                    continue

                terminator = self.suffix_line_terminator(line)
                if terminator:
                    line = line[: -len(terminator)]

                trailing = False
                yield line
            else:
                trailing = True
                self.file.seek(where)
                yield None


class AtomicOpen:
    """https://stackoverflow.com/a/46407326/3705710"""

    def __init__(self, path, *args, noop=False, **kwargs):
        if noop is False:
            self.file = open(path, *args, **kwargs)
            self.lock_file(self.file)
        self.noop = noop

    @staticmethod
    def lock_file(f):
        if f.writable():
            lockf(f, LOCK_EX)

    @staticmethod
    def unlock_file(f):
        if f.writable():
            lockf(f, LOCK_UN)

    def __enter__(self, *args, **kwargs):
        if self.noop is False:
            return self.file

    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
        if self.noop is False:
            self.file.flush()
            os.fsync(self.file.fileno())
            self.unlock_file(self.file)
            self.file.close()
            if exc_type is not None:
                return False
            else:
                return True


class bcolors:
    """https://stackoverflow.com/a/287944/3705710"""

    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    LIGHTGREY = "\033[0;37m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


##############
# ARG PARSING


def arg_requires_value(arg: str, option: Optional[str] = None) -> bool:
    def dashes(a: str):
        return "-" if len(a) == 1 else "--"

    if arg in ["cache-dir"]:
        return True
    if arg in ["h", "help"]:
        return False
    if option is None:
        if arg in ["version"]:
            return False
    elif option == "run":
        if arg in ["s", "shell", "split-output"]:
            return False
        if arg in ["n", "name"]:
            return True
    elif option == "start":
        pass
    elif option == "stop":
        if arg in ["k", "kill"] + signals_list():
            return False
    elif option == "rm":
        if arg in ["a", "all"]:
            return False
    elif option == "ls":
        if arg in ["a", "all"]:
            return False
    elif option == "logs":
        if arg in ["f", "follow", "head"]:
            return False
    raise TtmException(f"Unrecognized argument {dashes(arg)}{arg}")


def is_value_next(args: List[str], pos: int) -> bool:
    return pos + 1 < len(args) and not args[pos + 1].startswith("-")


def parse_args(
    args_to_parse: List[str],
) -> Tuple[Dict, Optional[str], Dict, Optional[List[str]]]:
    args = args_to_parse[1:]
    global_args: Dict[str, Union[str, bool]] = {}
    option = None
    pos = 0
    while True:
        if pos >= len(args):
            break
        current_arg = args[pos]
        if current_arg in ["run", "start", "stop", "rm", "ls", "logs"]:
            option = current_arg
            pos += 1
            break
        elif current_arg.startswith("--"):
            current_arg = current_arg[2:]
            if arg_requires_value(current_arg, option):
                if not is_value_next(args, pos):
                    raise TtmException(f"Argument --{current_arg} requires a value")
                global_args[current_arg] = args[pos + 1]
                pos += 2
                continue
            else:
                global_args[current_arg] = True
                pos += 1
                continue
        elif current_arg.startswith("-"):
            current_arg = current_arg[1:]
            if len(current_arg) == 1:
                if arg_requires_value(current_arg, option):
                    if not is_value_next(args, pos):
                        raise TtmException(f"Argument -{current_arg} requires a value")
                    global_args[current_arg] = args[pos + 1]
                    pos += 2
                    continue
                else:
                    global_args[current_arg] = True
                    pos += 1
                    continue
            else:
                for letter in current_arg:
                    if arg_requires_value(letter, option):
                        raise TtmException(
                            f"Argument -{letter} cannot be grouped with other arguments"
                        )
                    global_args[letter] = True
                    pos += 1
                    continue
        else:
            raise TtmException(f"Unrecognized option {current_arg}")
        pos += 1

    option_args: Dict[str, Union[str, bool]] = {}
    command = None

    if option is not None:
        if pos >= len(args) and option not in ["ls"]:
            raise TtmException(f"Missing arguments for option '{option}'")
        while True:
            if pos >= len(args):
                break
            current_arg = args[pos]
            if current_arg.startswith("--"):
                current_arg = current_arg[2:]
                if arg_requires_value(current_arg, option):
                    if not is_value_next(args, pos):
                        raise TtmException(f"Argument --{current_arg} requires a value")
                    option_args[current_arg] = args[pos + 1]
                    pos += 2
                    continue
                else:
                    option_args[current_arg] = True
                    pos += 1
                    continue
            elif current_arg.startswith("-"):
                current_arg = current_arg[1:]
                if len(current_arg) == 1:
                    if arg_requires_value(current_arg, option):
                        if not is_value_next(args, pos):
                            raise TtmException(
                                f"Argument -{current_arg} requires a value"
                            )
                        option_args[current_arg] = args[pos + 1]
                        pos += 2
                        continue
                    else:
                        option_args[current_arg] = True
                        pos += 1
                        continue
                else:
                    for letter in current_arg:
                        if arg_requires_value(letter, option) and not is_value_next(
                            args, pos
                        ):
                            raise TtmException(
                                f"Argument -{letter} cannot be grouped with other arguments"
                            )
                        option_args[letter] = True
                        pos += 1
                        continue
            else:
                command = args[pos:]
                break

    return global_args, option, option_args, command


##################
# FILE OPERATIONS


def init_cache_dir(cache_dir: Optional[Union[str, Path]]):
    global CACHE_DIR
    global LOCK_PATH
    if cache_dir is not None:
        CACHE_DIR = Path(cache_dir)
    os.makedirs(CACHE_DIR, exist_ok=True)
    LOCK_FILE_NAME = "lock"
    LOCK_PATH = Path(CACHE_DIR / LOCK_FILE_NAME)
    LOCK_PATH.touch(exist_ok=True)


def get_task_label(task: Task):
    if task["name"] is not None:
        return f"{task['name']}-{task['id']}"
    else:
        return task["id"]


def parse_task_id_or_name(task_name_or_id: str) -> Tuple[Optional[str], Optional[str]]:
    try:
        task_id = str(int(task_name_or_id))
        name = None
    except (ValueError, TypeError):
        task_id = None
        name = task_name_or_id
    return task_id, name


def create_pidfile(task: Task):
    if task.get("pid") is not None:
        if task.get("pidfile"):
            Path(task["pidfile"]).unlink(missing_ok=True)
        fh, task["pidfile"] = tempfile.mkstemp()
        with os.fdopen(fh, "w") as f:
            f.write(task["pid"])


def create_task_cache(task: Task, split_output=False) -> Task:
    dir_name = get_task_label(task)
    dir_path = CACHE_DIR / dir_name
    os.makedirs(dir_path, exist_ok=True)
    filepath = dir_path / "task.json"
    timestamp = datetime.now().strftime(TIMESTAMP_FMT)
    if split_output:
        stdout_path = dir_path / f"{dir_name}-{timestamp}.out"
        stderr_path = dir_path / f"{dir_name}-{timestamp}.err"
        task.update(
            {
                "stdout": str(stdout_path),
                "stderr": str(stderr_path),
                "started_at": timestamp,
            }
        )
    else:
        logs_path = dir_path / f"{dir_name}-{timestamp}.log"
        task.update(
            {
                "logs": str(logs_path),
                "started_at": timestamp,
            }
        )
    create_pidfile(task)
    with open(filepath, "w") as f:
        task_to_dump = dict(task)
        task_to_dump.pop("pid", None)
        json.dump(task_to_dump, f)
    return task


def update_task_cache(task: Task):
    dir_name = get_task_label(task)
    dir_path = CACHE_DIR / dir_name
    filepath = dir_path / "task.json"
    create_pidfile(task)
    with open(filepath, "w") as f:
        task_to_dump = dict(task)
        task_to_dump.pop("pid", None)
        json.dump(task_to_dump, f)


def is_task_running(task: Task) -> bool:
    output = check_output(["ps", "-ax", "-o", "pid,args"], start_new_session=True)
    for line in output.splitlines():
        decoded = line.decode().strip()
        ps_pid, cmdline = decoded.split(" ", 1)
        if task.get("pid") is not None and ps_pid == task["pid"]:
            return True
    return False


def get_task_from_cache_file(cache_file_path: str):
    with open(cache_file_path) as f:
        task = json.load(f)
    if task.get("pidfile") and Path(task["pidfile"]).exists():
        with open(task["pidfile"], "r") as f:
            task["pid"] = f.read()
    return task


def find_task_by_name(name: str) -> Optional[Task]:
    for filename in os.listdir(CACHE_DIR):
        if filename not in RESERVED_FILE_NAMES and filename.split("-")[0] == name:
            path = abspath(join(CACHE_DIR, filename, "task.json"))
            return get_task_from_cache_file(path)
    return None


def find_task_by_id(task_id: str) -> Optional[Dict]:
    for filename in os.listdir(CACHE_DIR):
        if filename in RESERVED_FILE_NAMES:
            continue
        try:
            filename_split = filename.split("-")
            if len(filename_split) == 1:
                filename_task_id = filename_split[0]
            else:
                filename_task_id = filename_split[1]
            if filename_task_id == task_id:
                path = abspath(join(CACHE_DIR, filename, "task.json"))
                return get_task_from_cache_file(path)
        except IndexError:
            pass
    return None


def delete_pidfile(task: Task):
    if task.get("pidfile"):
        Path(task["pidfile"]).unlink(missing_ok=True)


def remove_task_by_name(name: str):
    with AtomicOpen(LOCK_PATH):
        for filename in os.listdir(CACHE_DIR):
            filename_split = filename.split("-")
            if len(filename_split) == 1:
                continue
            else:
                filename_task_name = filename_split[0]

            if filename_task_name == name:
                task = find_task_by_name(name)
                if task is None:
                    raise TtmException("Failed to find task by name")
                if is_task_running(task):
                    raise TtmException(
                        "Cannot remove task while it's running.\n"
                        "To stop it, run:\n"
                        f"tinytaskmanager stop {name}"
                    )
                dir_path = abspath(join(CACHE_DIR, filename))
                rmtree(dir_path)
                delete_pidfile(task)
                return
        raise TtmException(f"No task with name {name}")


def remove_task_by_id(task_id: str):
    with AtomicOpen(LOCK_PATH):
        for filename in os.listdir(CACHE_DIR):
            try:
                filename_split = filename.split("-")
                if len(filename_split) == 1:
                    filename_task_id = filename_split[0]
                else:
                    filename_task_id = filename_split[1]

                if filename_task_id == task_id:
                    task = find_task_by_id(task_id)
                    if task is None:
                        raise TtmException("Failed to find task by id")
                    if is_task_running(task):
                        raise TtmException(
                            "Cannot remove task while it's running.\n"
                            "To stop it, run:\n"
                            f"tinytaskmanager stop {task_id}"
                        )
                    dir_path = abspath(join(CACHE_DIR, filename))
                    rmtree(dir_path)
                    delete_pidfile(task)
                    return
            except IndexError:
                pass
        raise TtmException(f"No task with ID {task_id}")


def generate_id():
    existing_ids = []
    for filename in os.listdir(CACHE_DIR):
        try:
            existing_ids.append(str(int(filename)))
        except ValueError:
            try:
                existing_ids.append(str(int(filename.split("-")[1])))
            except (ValueError, IndexError):
                pass
    for i in range(1, 10000):
        str_i = str(i)
        if str_i not in existing_ids:
            return str_i
    raise TtmException("Failed to generated task ID")


def get_child_pids(parent_pid: int):
    output = check_output(["ps", "-ax", "-o", "pid,ppid"], start_new_session=True)
    ppid = str(parent_pid)
    child_pids = []
    for line in output.splitlines():
        decoded = line.decode().strip()
        ps_child_pid, ps_ppid = decoded.split(None, 1)
        if ps_ppid == ppid:
            child_pids.append(int(ps_child_pid))
    return child_pids


def kill_recursively(pid: int, sig: int):
    for child_pid in get_child_pids(pid):
        kill_recursively(child_pid, sig)
    os.kill(pid, sig)


#############
# OPERATIONS


def run(
    command: List[str],
    name: Optional[str] = None,
    split_output=False,
    shell=False,
) -> Task:
    with AtomicOpen(LOCK_PATH):
        if name is not None:
            task = find_task_by_name(name)
            if task:
                if is_task_running(task):
                    raise TtmException(
                        f"Task {name} is already running with PID {task['pid']}"
                    )
                raise TtmException(
                    f"Task {name} already exists and it's not running.\n"
                    "To remove it, run:\n"
                    f"tinytaskmanager rm {name}"
                )
        task = {
            "id": generate_id(),
            "name": name,
            "cwd": os.getcwd(),
            "command": command,
            "shell": shell,
        }
        if split_output:
            task = create_task_cache(task, split_output=split_output)
            stdout_path = task["stdout"]
            stderr_path = task["stderr"]
            logs_path = ""
        else:
            task = create_task_cache(task, split_output=split_output)
            stdout_path = ""
            stderr_path = ""
            logs_path = task["logs"]
        if split_output:
            with open(stdout_path, "wb") as out:
                with open(stderr_path, "wb") as err:
                    proc = Popen(
                        build_cmd(command, shell),
                        shell=shell,
                        cwd=task["cwd"],
                        stdin=DEVNULL,
                        stdout=out,
                        stderr=err,
                        start_new_session=True,
                    )
        else:
            with open(logs_path, "wb") as output:
                proc = Popen(
                    build_cmd(command, shell),
                    shell=shell,
                    cwd=task["cwd"],
                    stdin=DEVNULL,
                    stdout=output,
                    stderr=output,
                    start_new_session=True,
                )
        task["pid"] = str(proc.pid)
        update_task_cache(task)
        return task


def start_task(task_id: Optional[str] = None, name: Optional[str] = None):
    with AtomicOpen(LOCK_PATH):
        if name is not None:
            task = find_task_by_name(name)
            if task is not None:
                if is_task_running(task):
                    raise TtmException(
                        f"Task {name} is already running with PID {task['pid']}"
                    )
            else:
                raise TtmException(f"No task with name {name}")
        elif task_id is not None:
            task = find_task_by_id(task_id)
            if task is not None:
                if is_task_running(task):
                    raise TtmException(
                        f"Task with ID {task_id} is already running with PID {task['pid']}"
                    )
            else:
                raise TtmException(f"No task with ID {task_id}")
        else:
            raise ValueError("Either task_id or name must be set")

        if task["name"] is not None:
            dir_name = f"{task['name']}-{task['id']}"
        else:
            dir_name = task["id"]
        dir_path = CACHE_DIR / dir_name
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        if task.get("stdout") is not None:
            task["stdout"] = str(dir_path / f"{dir_name}-{timestamp}.out")
            task["stderr"] = str(dir_path / f"{dir_name}-{timestamp}.err")
            with open(task["stdout"], "wb") as out:
                with open(task["stderr"], "wb") as err:
                    proc = Popen(
                        build_cmd(task["command"], task["shell"]),
                        shell=task["shell"],
                        cwd=task["cwd"],
                        stdin=DEVNULL,
                        stdout=out,
                        stderr=err,
                        start_new_session=True,
                    )
        else:
            task["logs"] = str(dir_path / f"{dir_name}-{timestamp}.log")
            with open(task["logs"], "wb") as output:
                proc = Popen(
                    build_cmd(task["command"], task["shell"]),
                    shell=task["shell"],
                    cwd=task["cwd"],
                    stdin=DEVNULL,
                    stdout=output,
                    stderr=output,
                    start_new_session=True,
                )
        task["pid"] = str(proc.pid)
        task["started_at"] = timestamp
        update_task_cache(task)


def stop_task(
    task_id: Optional[str] = None, name: Optional[str] = None, sig: int = SIGTERM
):
    with AtomicOpen(LOCK_PATH):
        if name is not None:
            task = find_task_by_name(name)
            if task is not None:
                if not is_task_running(task):
                    raise TtmException(f"Task {name} is not running")
            else:
                raise TtmException(f"No task with name {name}")
        elif task_id is not None:
            task = find_task_by_id(task_id)
            if task is not None:
                if not is_task_running(task):
                    raise TtmException(f"Task with ID {task_id} is not running")
            else:
                raise TtmException(f"No task with ID {task_id}")
        else:
            raise ValueError("Either task_id or name must be set")

    # We kill and busy wait outside the above file lock for better parallel performance
    kill_recursively(int(task["pid"]), sig)
    while True:
        # TODO add timeout
        with AtomicOpen(LOCK_PATH):
            if not is_task_running(task):
                delete_pidfile(task)
                break
        sleep(BUSY_LOOP_INTERVAL)


def remove_all_tasks():
    with AtomicOpen(LOCK_PATH):
        for filename in os.listdir(CACHE_DIR):
            if TERMINATE:
                return
            if filename in RESERVED_FILE_NAMES:
                continue
            path = abspath(join(CACHE_DIR, filename, "task.json"))
            try:
                task = get_task_from_cache_file(path)
                if is_task_running(task):
                    print_error(f"Task {task['id']}: cannot remove while it's running")
                else:
                    dir_path = abspath(join(CACHE_DIR, filename))
                    rmtree(dir_path)
            except (NotADirectoryError, FileNotFoundError, ValueError):
                pass


def rm(task_name_or_id: Optional[str], rm_all=False) -> bool:
    try:
        if rm_all:
            remove_all_tasks()
        else:
            if task_name_or_id is None:
                raise ValueError("task_name_or_id is None")
            task_id, name = parse_task_id_or_name(task_name_or_id)
            if task_id is not None:
                remove_task_by_id(task_id)
            elif name is not None:
                remove_task_by_name(name)
    except TtmException as e:
        print_error(str(e))
        return False
    return True


def logs(task_name_or_id: str, follow=False, head=False):
    def print_lines(lines):
        if isinstance(lines, list):
            for line in lines:
                stdout.buffer.write(line)
                stdout.buffer.write("\n".encode())
        elif isinstance(lines, bytes):
            stdout.buffer.write(lines)
            stdout.buffer.write("\n".encode())
        stdout.buffer.flush()

    if follow and head:
        raise TtmException("--follow and --head cannot be used together")

    task_id, name = parse_task_id_or_name(task_name_or_id)

    if task_id is not None:
        task = find_task_by_id(task_id)
        if task is None:
            raise TtmException(f"No task with ID {task_id}")
    elif name is not None:
        task = find_task_by_name(name)
        if task is None:
            raise TtmException(f"No task with name {name}")
    else:
        raise ValueError("task_id and name are None")

    logs_path = task.get("logs")
    if logs_path is None:
        raise TtmException(
            "Task was created using --split-output, use 'stdout' or 'stderr' instead of 'logs'"
        )

    line_count = 15

    if not head and follow:
        with open(logs_path, "rb") as file:
            print_grey(f"{logs_path} last {line_count} lines:")
            print_lines(Tailer(file).tail(lines=line_count))

    with open(logs_path, "rb") as file:
        if head:
            print_grey(f"{logs_path} first {line_count} lines:")
            print_lines(Tailer(file).head(lines=line_count))
        elif follow:
            print_grey(f"{logs_path} followed tail:")
            for line in Tailer(file, end=True).follow():
                if line is None:
                    time.sleep(0.01)
                    continue
                print_lines(line)
        else:
            print_grey(f"{logs_path} last {line_count} lines:")
            print_lines(Tailer(file).tail(lines=line_count))

        print_lines([])


def start(task_name_or_id: str) -> bool:
    task_id, name = parse_task_id_or_name(task_name_or_id)

    try:
        start_task(task_id=task_id, name=name)
        return True
    except TtmException as e:
        print_error(str(e))
        return False


def stop(task_name_or_id: str, sig: int):
    task_id, name = parse_task_id_or_name(task_name_or_id)

    try:
        stop_task(task_id=task_id, name=name, sig=sig)
        return True
    except TtmException as e:
        print_error(str(e))
        return False


def ls(ls_all=False, command: Optional[List[str]] = None):
    tasks = []
    with AtomicOpen(LOCK_PATH):
        for filename in os.listdir(CACHE_DIR):
            if filename in RESERVED_FILE_NAMES:
                continue
            path = abspath(join(CACHE_DIR, filename, "task.json"))
            force_list = False
            if command:
                for task_name_or_id in command:
                    task_id, name = parse_task_id_or_name(task_name_or_id)
                    filename_split = filename.split("-")
                    if task_id in filename_split or name in filename_split:
                        force_list = True
                if not force_list:
                    continue
            try:
                task = get_task_from_cache_file(path)
                task["started_at"] = datetime.strptime(
                    task["started_at"], TIMESTAMP_FMT
                )
                if is_task_running(task):
                    diff = datetime.now() - task["started_at"]
                    task["uptime"] = format_seconds(int(diff.total_seconds()))
                    tasks.append(task)
                elif ls_all or force_list:
                    task["pid"] = "-"
                    task["uptime"] = "-"
                    tasks.append(task)
            except (NotADirectoryError, FileNotFoundError, ValueError):
                pass

    name_len_max = 4
    for task in tasks:
        if task["name"] is not None and len(task["name"]) > name_len_max:
            name_len_max = len(task["name"])

    tasks = sorted(tasks, key=lambda d: d["started_at"])

    columns = get_terminal_size((80, 20)).columns
    name_size = min(name_len_max, 16)
    command_size = columns - 21 - name_size
    template = r"{0:4} {1:NAME_SIZE} {2:6} {3:7} {4:COMMAND_SIZE}"
    template = template.replace("NAME_SIZE", str(name_size))
    template = template.replace("COMMAND_SIZE", str(command_size))
    print(template.format("ID", "NAME", "UPTIME", "PID", "COMMAND"))
    for task in tasks:
        print(
            template.format(
                task["id"],
                task["name"] or "-",
                task["uptime"],
                task["pid"],
                shlex.join(task["command"]),
            )
        )


#######
# MISC


def build_cmd(command: List[str], shell: bool):
    if shell:
        if len(command) == 1:
            return command[0]
        else:
            return shlex.join(command)
    else:
        return command


def format_seconds(seconds, long=False):
    if long:
        raise NotImplementedError()
    else:
        days, remainder = divmod(seconds, 86400)
        hours, remainder = divmod(remainder, 3600)
        minutes, seconds = divmod(remainder, 60)
        s = ""
        if days > 0:
            s += f"{days}d"
        elif hours > 0:
            s += f"{hours}h"
        elif minutes > 0:
            s += f"{minutes}m"
        else:
            s += f"{seconds}s"
        return s


def signal_handler(signum, frame):
    global TERMINATE
    if signum in [SIGINT, SIGTERM]:
        TERMINATE = True


def signals_list():
    return [str(sig.value) for sig in Signals]


def print_error(msg: str, *args, **kwargs):
    if "file" not in kwargs:
        kwargs["file"] = stderr
    print(f"{bcolors.FAIL}{msg}{bcolors.ENDC}", *args, **kwargs)


def print_grey(msg: str, *args, **kwargs):
    print(f"{bcolors.LIGHTGREY}{msg}{bcolors.ENDC}", *args, **kwargs)


def print_warning(msg: str, *args, **kwargs):
    if "file" not in kwargs:
        kwargs["file"] = stderr
    print(f"{bcolors.WARNING}{msg}{bcolors.ENDC}", *args, **kwargs)


def print_success(msg: str, *args, **kwargs):
    print(f"{bcolors.OKGREEN}{msg}{bcolors.ENDC}", *args, **kwargs)


def print_help():
    print()
    print("Usage:  tinytaskmanager [OPTIONS] COMMAND")
    print()
    print("Tiny task manager for Linux, MacOS and Unix-like systems")
    print()
    print("Options:")

    print(
        f"      --cache-dir   Use a custom cache directory instead of the default {CACHE_DIR}"
    )
    print("  -h, --help        Display help")
    print("      --version     Display program version")
    print()
    print("Commands:")
    print("  logs            Display logs of a task")
    print("  ls              List tasks")
    print("  rm              Remove tasks")
    print("  run             Run a new task")
    print("  start           Start tasks")
    print("  stop            Stop tasks")


def print_help_logs():
    print()
    print("Usage:  tinytaskmanager logs TASK")
    print()
    print("Display logs of a task.")
    print("TASK can be a task ID or a task name.")
    print()


def print_help_ls():
    print()
    print("Usage:  tinytaskmanager ls [OPTIONS]")
    print()
    print("List tasks")
    print()
    print("Options:")
    print("  -a, --all        List all tasks, including stopped")
    print()


def print_help_rm():
    print()
    print("Usage:  tinytaskmanager rm [OPTIONS] [TASK]")
    print()
    print("Remove tasks.")
    print("TASK can be a task ID or a task name.")
    print()
    print("Options:")
    print("  -a, --all        Remove all tasks")
    print()


def print_help_run():
    print()
    print("Usage:  tinytaskmanager run [OPTIONS] COMMAND")
    print()
    print("Run a new task")
    print()
    print("Options:")
    print("  -s, --shell        Run COMMAND in a shell")
    print()
    print("Examples:")
    print("  tinytaskmanager run /path/to/my/program arg1 arg2")
    print("  tinytaskmanager run -s 'echo test >> myfile'")
    print()


def print_help_start():
    print()
    print("Usage:  tinytaskmanager start TASK")
    print()
    print("Start a task")
    print()
    print("Examples:")
    print("  tinytaskmanager start 3")
    print("  tinytaskmanager start mytaskname")
    print()


def print_help_stop():
    print()
    print("Usage:  tinytaskmanager stop TASK")
    print()
    print("Stop a task")
    print()
    print("Options:")
    print("  -k, --kill        Send a SIGKILL signal instead of the default SIGTERM.")
    print("                    Equivalent to -9.")
    print("  -SIG              Send a specific signal instead of the default SIGTERM")
    print()
    print("Examples:")
    print("  tinytaskmanager stop 3")
    print("  tinytaskmanager stop mytaskname")
    print("  tinytaskmanager stop --kill my_hanged_task")
    print()
    print("  # Send a SIGINT signal to a task")
    print("  tinytaskmanager stop -2 my_interruptable_task")
    print()
    print("  # Send a SIGTERM signal to a task")
    print("  tinytaskmanager stop -9 my_hanged_task")
    print()


def main():
    for sig in [SIGINT, SIGTERM]:
        signal(sig, signal_handler)

    try:
        if len(argv) == 1:
            print_error("No option provided. Use -h for help.")
            exit(1)
        global_args, option, option_args, command = parse_args(argv)

        init_cache_dir(global_args.get("cache-dir"))

        if option is None:
            if global_args.get("version"):
                print(VERSION)
                return
            if global_args.get("h") or global_args.get("help"):
                print_help()
                return

        elif option == "logs":
            if option_args.get("h") or option_args.get("help"):
                print_help_logs()
                return
            if command is None:
                raise TtmException("Task ID or name must be provided")
            if len(command) > 1:
                raise TtmException(
                    "A single task ID or name must be provided to 'logs'"
                )
            follow = option_args.get("f") or option_args.get("follow") or False
            head = option_args.get("head") or False
            logs(command[0], follow=follow, head=head)

        elif option == "ls":
            if option_args.get("h") or option_args.get("help"):
                print_help_ls()
                return
            ls_all = bool(option_args.get("a") or option_args.get("all"))
            if command:
                if ls_all:
                    raise TtmException(
                        "-a/--all is not allowed when specific tasks are provided"
                    )
            ls(ls_all=ls_all, command=command)

        elif option == "rm":
            if option_args.get("h") or option_args.get("help"):
                print_help_rm()
                return
            rm_all = option_args.get("a") or option_args.get("all")
            if rm_all is True:
                rm(None, rm_all=rm_all)
            else:
                if command is None:
                    raise TtmException("Task ID or name must be provided")
                pool = ThreadPool(len(command))
                results = pool.map(rm, command)
                if not all(results):
                    exit(1)

        elif option == "run":
            if option_args.get("h") or option_args.get("help"):
                print_help_run()
                return
            name = option_args.get("n") or option_args.get("name") or None
            if name is not None:
                if not re.match(r"^[a-zA-Z_]+$", name):
                    raise TtmException(
                        "Only letters and underscore are allowed in task name"
                    )
            shell = bool(option_args.get("s") or option_args.get("shell"))
            if command is None:
                raise TtmException("A command must be provided")
            run(command, name=name, shell=shell)

        elif option == "start":
            if option_args.get("h") or option_args.get("help"):
                print_help_start()
                return
            if command is None:
                raise TtmException("Task ID or name must be provided")
            pool = ThreadPool(len(command))
            results = pool.map(start, command)
            if not all(results):
                exit(1)

        elif option == "stop":
            if option_args.get("h") or option_args.get("help"):
                print_help_stop()
                return
            if command is None:
                raise TtmException("Task ID or name must be provided")
            stop_sig = None
            for sig in signals_list():
                if option_args.get(sig):
                    if stop_sig is not None:
                        raise TtmException("Only one signal can be provided")
                    stop_sig = int(sig)
            if option_args.get("k") or option_args.get("kill"):
                if stop_sig is not None:
                    raise TtmException(
                        "-k/--kill cannot be used when a signal is provided"
                    )
                stop_sig = SIGKILL
            if stop_sig is None:
                stop_sig = SIGTERM
            pool = ThreadPool(len(command))
            arg_list = []
            for c in command:
                arg_list.append((c, stop_sig))
            results = pool.starmap(stop, arg_list)
            if not all(results):
                exit(1)

    except TtmException as e:
        print_error(str(e))
        exit(1)


if __name__ == "__main__":
    main()