#!/usr/bin/env python3


# Libervia XMPP client
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# 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 Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""tools to launch process in a async way (using Twisted)"""

import os
import signal
from pathlib import Path
from typing import Any
from twisted.internet import defer, protocol, reactor, task
from twisted.python.failure import Failure
from libervia.backend.core.i18n import _
from libervia.backend.core import exceptions
from libervia.backend.core.log import getLogger
from libervia.backend.tools.common import utils

log = getLogger(__name__)

# Processes registered for auto-cleaning on shutdown (pid -> CommandProtocol)
_running: dict[int, "CommandProtocol"] = {}


class CommandProtocol(protocol.ProcessProtocol):
    """handle an external command"""

    # name of the command
    name: str | None = None
    # full path to the command
    command: bytes | None = None
    # True to activate logging of command outputs
    log: bool = False

    def __init__(self, deferred, stdin=None):
        """
        @param deferred(defer.Deferred): will be called when command is completed
        @param stdin(str, None): if not None, will be push to standard input
        """
        self._stdin = stdin
        self._deferred = deferred
        self.data = []
        self.err_data = []
        self.cmd_args: list[str] | None = None
        self.cmd_kwargs: dict[str, Any] | None = None
        self.pid: int | None = None
        self.pid_file: Path | None = None

    @property
    def command_name(self):
        """returns command name or empty string if it can't be guessed"""
        if self.name is not None:
            return self.name
        elif self.command is not None:
            return os.path.splitext(os.path.basename(self.command))[0].decode(
                "utf-8", "ignore"
            )
        else:
            return ""

    def connectionMade(self):
        if self._stdin is not None:
            self.transport.write(self._stdin)
            self.transport.closeStdin()

    def outReceived(self, data):
        if self.log:
            log.info(data.decode("utf-8", "replace"))
        self.data.append(data)

    def errReceived(self, data):
        if self.log:
            log.warning(data.decode("utf-8", "replace"))
        self.err_data.append(data)

    def processEnded(self, reason):
        if self.pid is not None:
            _running.pop(self.pid, None)
        if self.pid_file is not None:
            self.pid_file.unlink(missing_ok=True)
        data = b"".join(self.data)
        if reason.value.exitCode == 0:
            log.debug(f"{self.command_name!r} command succeed")
            # we don't use "replace" on purpose, we want an exception if decoding
            # is not working properly
            self._deferred.callback(data)
        else:
            err_data = b"".join(self.err_data)

            assert self.cmd_args is not None
            assert self.cmd_kwargs is not None
            msg = _(
                "Can't complete {name} command (error code: {code}):\n"
                "Executed command: {command}\n"
                "Keyword arguments:\n"
                "{command_kw}\n\n"
                "stderr:\n{stderr}\n{stdout}\n"
            ).format(
                name=self.command_name,
                code=reason.value.exitCode,
                command=" ".join(self.cmd_args),
                command_kw="\n".join(
                    f"    - {k} = {v!r}" for k, v in self.cmd_kwargs.items()
                ),
                stderr=err_data.decode(errors="replace"),
                stdout="stdout: " + data.decode(errors="replace") if data else "",
            )
            self._deferred.errback(
                Failure(exceptions.CommandException(msg, data, err_data))
            )

    @classmethod
    def run(
        cls,
        *args: str | Path,
        auto_clean: bool = True,
        pid_file: Path | None = None,
        **kwargs: Any,
    ) -> defer.Deferred:
        """Create a new CommandProtocol and execute the given command.

        @param *args: command arguments
            if cls.command is specified, it will be the path to the command to execute
            otherwise, first argument must be the path
        @param auto_clean: if True, the process will be stopped on shutdown if not
            already finished
        @param pid_file: if not None, the process PID is written there (and removed
            on exit), to allow cleaning it up after a crash
        @param **kwargs: can be:
            - stdin: data to push to standard input
            - verbose: if True stdout and stderr will be logged
            other keyword arguments will be used in reactor.spawnProcess
        @return: Deferred called with the process stdout on success
        @raise CommandException: command returned a non zero status, stdout and
            stderr will be attached to the exception
        """
        stdin = kwargs.pop("stdin", None)
        if stdin is not None:
            if isinstance(stdin, str):
                stdin = stdin.encode()
        verbose = kwargs.pop("verbose", False)
        args = list(args)
        d = defer.Deferred()
        prot = cls(d, stdin=stdin)
        if verbose:
            prot.log = True
        if cls.command is None:
            if not args:
                raise ValueError(
                    "You must either specify cls.command or use a full path to command "
                    "to execute as first argument"
                )
            command = args.pop(0)
            if isinstance(command, Path):
                command = str(command)
            if prot.name is None:
                name = os.path.splitext(os.path.basename(command))[0]
                prot.name = name
        else:
            command = cls.command
        cmd_args = [command] + args
        prot.cmd_args = cmd_args
        prot.cmd_kwargs = kwargs
        if "env" not in kwargs:
            # we pass parent environment by default
            # FIXME: `None` doesn't seem to work, despite what documentation says, to be
            #    checked and reported upstream if confirmed.
            kwargs["env"] = os.environ
        transport = reactor.spawnProcess(prot, command, cmd_args, **kwargs)
        prot.pid = transport.pid
        if pid_file is not None:
            prot.pid_file = pid_file
            prot.pid_file.write_text(f"{prot.pid}\n")
        if auto_clean:
            _running[prot.pid] = prot
        return d


def signal_all(sig: signal.Signals) -> None:
    """Send a signal to all the registered processes"""
    for pid, prot in list(_running.items()):
        try:
            prot.transport.signalProcess(sig)
        except Exception as e:
            log.debug(f"Can't send {sig.name} to process {pid}: {e}")
            _running.pop(pid, None)


async def stop_all(grace: float = 5.0) -> None:
    """Stop all processes registered for auto-cleaning.

    Send SIGTERM, wait for the processes to exit, then send SIGKILL to the remaining ones.

    @param grace: time (in seconds) to wait before sending SIGKILL
    """
    if not _running:
        return
    log.info(f"Stopping {len(_running)} external process(es) on shutdown")
    signal_all(signal.SIGTERM)
    while _running and grace > 0:
        await task.deferLater(reactor, 0.5, lambda: None)
        grace -= 0.5
    signal_all(signal.SIGKILL)


def kill_stale_process(pid_file: Path, expected_cwd: Path) -> None:
    """Kill a process recorded in a PID file if it is verified as a previous instance.

    Verification is done by comparing the process working directory with ``expected_cwd``
    (relying on /proc, Linux only). If the process can't be verified, it is left
    untouched.

    @param pid_file: path to the PID file
    @param expected_cwd: expected working directory of the process
    """
    if not utils.is_process_running(pid_file):
        pid_file.unlink(missing_ok=True)
        return
    pid = int(pid_file.read_text().strip())
    try:
        cwd = Path(os.readlink(f"/proc/{pid}/cwd"))
    except OSError:
        log.warning(
            f"Can't verify process {pid} (no /proc access), it won't be killed; if it "
            "is stale, please kill it manually."
        )
        return
    if cwd.resolve() != expected_cwd.resolve():
        log.error(
            f"Process {pid} recorded in {pid_file} has an unexpected working directory "
            f"({cwd!r}), it won't be killed."
        )
        pid_file.unlink(missing_ok=True)
        return
    log.info(f"Killing stale process {pid} left by a previous run.")
    os.kill(pid, signal.SIGTERM)
    pid_file.unlink(missing_ok=True)


run = CommandProtocol.run