import os
import signal
import stat
import subprocess
import tempfile
from contextlib import contextmanager
from functools import partial
from io import TextIOWrapper
from typing import Callable

from soniclue_basekit.cli_formatter import CliFormatter

# https://linux.die.net/man/1/timeout
_TIMEOUT_RETURN_CODE = 124


class ScriptRunnerError(RuntimeError):
    """Raised when a script exits with a non-zero return code or times out."""

    def __init__(
        self,
        cmd: list[str],
        returncode: int,
        stdout: str | None,
        stderr: str | None,
    ) -> None:
        super().__init__()
        self.cmd = cmd
        self.returncode = returncode
        self.stdout = stdout
        self.stderr = stderr

    @property
    def is_timeout(self) -> bool:
        """Returns True if the script exceeded its time limit."""
        return self.returncode == _TIMEOUT_RETURN_CODE

    def _outcome_str(self) -> str:
        return "timed out" if self.is_timeout else f"returned {self.returncode}"

    def __str__(self) -> str:
        parts = ["Running script file", self._outcome_str()]
        if self.stdout:
            parts.append(f'stdout="{self.stdout}"')
        if self.stderr:
            parts.append(f'stderr="{self.stderr}"')
        return " ".join(parts)


def _default_cmd_runner(
    cmd_array: list[str],
    work_dir: str | None,
    env: dict | None,
    redirect: TextIOWrapper | int | None = None,
    preexec_fn: Callable[[], None] | None = None,
) -> tuple[str, str] | None:
    """Runs a command in a subprocess.

    Args:
        cmd_array: Command and its arguments, e.g. ['/path/to/script', '-k', 'val'].
        work_dir: Working directory for the subprocess; None inherits the caller's cwd.
        env: Environment for the subprocess; None inherits the caller's environment.
        redirect: Routing for stdout/stderr.
            - None: subprocess inherits the caller's stdout/stderr.
            - subprocess.PIPE: capture and return the output.
            - A file handle: route output to that file.
        preexec_fn: POSIX-only callable invoked in the child process just before exec.

    Raises:
        ScriptRunnerError: If the process exits with a non-zero return code.

    Returns:
        Stripped (stdout, stderr) strings when redirect is subprocess.PIPE, else None.
    """
    return_output = redirect == subprocess.PIPE
    proc = subprocess.Popen(
        cmd_array,
        universal_newlines=True,
        cwd=work_dir,
        env=env,
        preexec_fn=preexec_fn,
        stdout=redirect,
        stderr=redirect,
    )
    stdout, stderr = proc.communicate()
    if return_output:
        stdout, stderr = stdout.strip(), stderr.strip()
    if proc.returncode != 0:
        raise ScriptRunnerError(
            cmd=cmd_array,
            returncode=proc.returncode,
            stdout=stdout,
            stderr=stderr,
        )
    return (stdout, stderr) if return_output else None


class ScriptRunner:
    """Runs inline script content or an existing script file in a subprocess."""

    def __init__(self, content: str, is_path: bool = False) -> None:
        """Constructor.

        Args:
            content: Script body, or a path to an existing file when is_path is True.
            is_path: If True, treat content as a file path instead of inline script text.
        """
        self._content = content
        self._is_path = is_path

    @contextmanager
    def as_temp_file(self, work_dir: str | None, **kwargs):
        """Context manager that yields a path to the script file.

        When is_path is True, yields the stored path directly with no cleanup.
        Otherwise, writes the script to a temporary file and deletes it on exit.

        Args:
            work_dir: Directory in which to create the temporary file.
            kwargs: Forwarded to tempfile.NamedTemporaryFile (e.g. prefix, suffix).
        """
        if self._is_path:
            yield self._content
            return
        with tempfile.NamedTemporaryFile(mode="w", dir=work_dir, delete=False, **kwargs) as fp:
            fp.write(self._content)
            script_file = fp.name
        try:
            yield script_file
        finally:
            os.unlink(script_file)

    def run(
        self,
        fmt: CliFormatter | None = None,
        work_dir: str | None = None,
        env: dict | None = None,
        cmd_runner: Callable | None = None,
        timeout_sec: float | None = None,
        timeout_signal: signal.Signals = signal.SIGINT,
        with_output: bool = False,
        log_file_handle: TextIOWrapper | None = None,
        preexec_fn: Callable[[], None] | None = None,
        timeout_kill_after_sec: float | None = None,
    ) -> tuple[str, str] | None:
        """Runs the script.

        Args:
            fmt: CliFormatter that holds additional arguments passed to the script.
            work_dir: Working directory for the subprocess.
            env: Environment for the subprocess; None inherits the caller's environment.
            cmd_runner: Optional callable invoked as cmd_runner(cmd_array, work_dir, env).
                If None, _default_cmd_runner is used. Mutually exclusive with
                with_output and log_file_handle.
            timeout_sec: Seconds before the script is killed; None means no timeout.
            timeout_signal: Signal sent to the process when the timeout expires.
            with_output: If True, capture and return stdout/stderr.
            log_file_handle: File handle to receive stdout/stderr output.
                Mutually exclusive with with_output.
            preexec_fn: POSIX-only callable invoked in the child process just before exec.
            timeout_kill_after_sec: Seconds to wait after timeout_signal before SIGKILL
                is sent (forwarded to the `timeout` command's --kill-after flag).

        Raises:
            ScriptRunnerError: If execution fails or the timeout is exceeded.

        Returns:
            (stdout, stderr) if with_output is True, else None.
        """
        if fmt is None:
            fmt = CliFormatter()
        assert isinstance(fmt, CliFormatter)
        assert not (with_output and log_file_handle is not None)
        if with_output or log_file_handle is not None:
            assert cmd_runner is None
        if cmd_runner is None:
            output_redirect = subprocess.PIPE if with_output else log_file_handle
            cmd_runner = partial(
                _default_cmd_runner,
                redirect=output_redirect,
                preexec_fn=preexec_fn,
            )
        assert callable(cmd_runner)

        with self.as_temp_file(work_dir, prefix="exec_script_") as script_file:
            os.chmod(script_file, os.stat(script_file).st_mode | stat.S_IEXEC)
            fmt.set_cmd(os.path.abspath(script_file))
            cmd_array = fmt.as_cmd_list()

            if timeout_sec is not None:
                timeout_fmt = CliFormatter("timeout")
                timeout_fmt.add("--signal", int(timeout_signal))
                if timeout_kill_after_sec is not None:
                    timeout_fmt.add("--kill-after", timeout_kill_after_sec)
                # Duration is a bare positional argument to `timeout`, not a key=value pair.
                cmd_array = timeout_fmt.as_cmd_list() + [str(timeout_sec)] + cmd_array

            return cmd_runner(cmd_array, work_dir, env)

    def run_with_output(self, **kwargs) -> tuple[str, str]:
        """Runs the script and returns (stdout, stderr).

        This is a convenience wrapper around run(with_output=True, ...).

        Args:
            kwargs: Keyword arguments forwarded to run().

        Raises:
            ScriptRunnerError: If execution fails or the timeout is exceeded.

        Returns:
            Tuple of (stdout, stderr).
        """
        return self.run(with_output=True, **kwargs)  # type: ignore
