# Copyright 2023 Oliver Smith
# SPDX-License-Identifier: GPL-3.0-or-later
import os
from pathlib import Path, PurePath
import shutil
import subprocess
from collections.abc import Sequence
from typing import overload, Literal

import pmb.config
import pmb.chroot
import pmb.chroot.binfmt
import pmb.helpers.run
import pmb.helpers.run_core
from pmb.core import Chroot
from pmb.types import (
    Env,
    PathString,
    RunOutputType,
    RunOutputTypeDefault,
    RunOutputTypePopen,
    RunReturnType,
)


def executables_absolute_path():
    """
    Get the absolute paths to the sh and chroot executables.
    """
    ret = {}
    for binary in ["sh", "chroot"]:
        path = shutil.which(binary, path=pmb.config.chroot_host_path)
        if not path:
            raise RuntimeError(
                f"Could not find the '{binary}'"
                " executable. Make sure that it is in"
                " your current user's PATH."
            )
        ret[binary] = path
    return ret


def rootm(
    cmds: Sequence[Sequence[PathString]],
    chroot: Chroot = Chroot.native(),
    working_dir: PurePath = PurePath("/"),
    output: RunOutputType = "log",
    output_return: bool = False,
    check: bool | None = None,
    env: Env = {},
    disable_timeout: bool = False,
    add_proxy_env_vars: bool = True,
) -> RunReturnType:
    """
    Run a list of commands inside a chroot as root.

    :param env: dict of environment variables to be passed to the command, e.g.
                {"JOBS": "5"}
    :param working_dir: chroot-relative working directory
    :param add_proxy_env_vars: if True, preserve HTTP_PROXY etc. vars from host
                               environment. pmb.chroot.user sets this to False
                               when calling pmb.chroot.root, because it already
                               makes the variables part of the cmd argument.

    See pmb.helpers.run_core.core() for a detailed description of all other
    arguments and the return value.
    """

    # Convert any Path objects to their string representation
    cmd_strs = [[os.fspath(x) for x in cmd] for cmd in cmds]

    # Readable log message (without all the escaping)
    msg = f"({chroot}) % "
    for key, value in env.items():
        msg += f"{key}={value} "
    if working_dir != PurePath("/"):
        msg += f"cd {working_dir}; "
    msg += "; ".join([" ".join(cmd_str) for cmd_str in cmd_strs])

    # Merge env with defaults into env_all
    env_all: Env = {
        "CHARSET": "UTF-8",
        "HISTFILE": "~/.ash_history",
        "HOME": "/root",
        "LANG": "UTF-8",
        "PATH": pmb.config.chroot_path,
        "PYTHONUNBUFFERED": "1",
        "SHELL": "/bin/ash",
        "TERM": "xterm",
    }
    for key, value in env.items():
        env_all[key] = value
    if add_proxy_env_vars:
        pmb.helpers.run_core.add_proxy_env_vars(env_all)

    # Build the command in steps and run it, e.g.:
    # cmd: ["echo", "test"]
    # cmd_chroot: ["/sbin/chroot", "/..._native", "/bin/sh", "-c", "echo test"]
    # cmd_sudo: ["sudo", "env", "-i", "sh", "-c", "PATH=... /sbin/chroot ..."]
    executables = executables_absolute_path()
    cmd_chroot = [
        executables["chroot"],
        chroot.path,
        "/bin/sh",
        "-c",
        pmb.helpers.run_core.flat_cmd(cmd_strs, Path(working_dir)),
    ]
    cmd_sudo = pmb.config.sudo(
        [
            "env",
            "-i",
            executables["sh"],
            "-c",
            pmb.helpers.run_core.flat_cmd([cmd_chroot], env=env_all),
        ]
    )
    return pmb.helpers.run_core.core(
        msg, cmd_sudo, None, output, output_return, check, True, disable_timeout
    )


@overload
def root(
    cmds: Sequence[PathString],
    chroot: Chroot = ...,
    working_dir: PurePath = ...,
    output: RunOutputTypePopen = ...,
    output_return: Literal[False] = ...,
    check: bool | None = ...,
    env: Env = ...,
    disable_timeout: bool = ...,
    add_proxy_env_vars: bool = ...,
) -> subprocess.Popen: ...


@overload
def root(
    cmds: Sequence[PathString],
    chroot: Chroot = ...,
    working_dir: PurePath = ...,
    output: RunOutputTypeDefault = ...,
    output_return: Literal[False] = ...,
    check: bool | None = ...,
    env: Env = ...,
    disable_timeout: bool = ...,
    add_proxy_env_vars: bool = ...,
) -> int: ...


@overload
def root(
    cmds: Sequence[PathString],
    chroot: Chroot = ...,
    working_dir: PurePath = ...,
    output: RunOutputType = ...,
    output_return: Literal[True] = ...,
    check: bool | None = ...,
    env: Env = ...,
    disable_timeout: bool = ...,
    add_proxy_env_vars: bool = ...,
) -> str: ...


def root(
    cmds: Sequence[PathString],
    chroot: Chroot = Chroot.native(),
    working_dir: PurePath = PurePath("/"),
    output: RunOutputType = "log",
    output_return: bool = False,
    check: bool | None = None,
    env: Env = {},
    disable_timeout: bool = False,
    add_proxy_env_vars: bool = True,
) -> RunReturnType:
    return rootm(
        [cmds],
        chroot,
        working_dir,
        output,
        output_return,
        check,
        env,
        disable_timeout,
        add_proxy_env_vars,
    )


def userm(
    cmds: Sequence[Sequence[PathString]],
    chroot: Chroot = Chroot.native(),
    working_dir: Path = Path("/"),
    output: RunOutputType = "log",
    output_return: bool = False,
    check: bool | None = None,
    env: Env = {},
) -> RunReturnType:
    """
    Run a command inside a chroot as "user". We always use the BusyBox
    implementation of 'su', because other implementations may override the PATH
    environment variable (#1071).

    :param env: dict of environment variables to be passed to the command, e.g.
                {"JOBS": "5"}

    See pmb.helpers.run_core.core() for a detailed description of all other
    arguments and the return value.
    """
    env = env.copy()
    pmb.helpers.run_core.add_proxy_env_vars(env)

    if "HOME" not in env:
        env["HOME"] = "/home/pmos"

    flat_cmd = pmb.helpers.run_core.flat_cmd(cmds, env=env)
    cmd = ["busybox", "su", "pmos", "-c", flat_cmd]
    # Can't figure out why this one fails :(
    return pmb.chroot.root(  # type: ignore[call-overload]
        cmd, chroot, working_dir, output, output_return, check, {}, add_proxy_env_vars=False
    )


@overload
def user(
    cmd: Sequence[PathString],
    chroot: Chroot = ...,
    working_dir: Path = ...,
    output: RunOutputTypePopen = ...,
    output_return: Literal[False] = ...,
    check: bool | None = ...,
    env: Env = ...,
) -> subprocess.Popen: ...


@overload
def user(
    cmd: Sequence[PathString],
    chroot: Chroot = ...,
    working_dir: Path = ...,
    output: RunOutputTypeDefault = ...,
    output_return: Literal[False] = ...,
    check: bool | None = ...,
    env: Env = ...,
) -> int: ...


@overload
def user(
    cmd: Sequence[PathString],
    chroot: Chroot = ...,
    working_dir: Path = ...,
    output: RunOutputType = ...,
    output_return: Literal[True] = ...,
    check: bool | None = ...,
    env: Env = ...,
) -> str: ...


def user(
    cmd: Sequence[PathString],
    chroot: Chroot = Chroot.native(),
    working_dir: Path = Path("/"),
    output: RunOutputType = "log",
    output_return: bool = False,
    check: bool | None = None,
    env: Env = {},
) -> RunReturnType:
    return userm([cmd], chroot, working_dir, output, output_return, check, env)


def exists(username: str, chroot: Chroot = Chroot.native()) -> bool:
    """
    Checks if username exists in the system

    :param username: User name
    :returns: bool
    """
    output = pmb.chroot.root(
        ["getent", "passwd", username], chroot, output_return=True, check=False
    )
    return len(output) > 0