import contextlib
import copy
import dataclasses
import sys
from functools import partialmethod
from pathlib import Path
from typing import Any, Literal, get_args
from collections.abc import Callable, Generator
import shutil
import subprocess
import logging

import pytest

from testlib.fixtures.command import CommandResult, Command
from testlib.fixtures.env import ManagedEnv
from testlib.utils import is_value_of_type

from textwrap import dedent


type _NixSettingValue = str | int | list[str] | bool | None
type _NixValue = str | int | float | list[_NixValue] | dict[str, _NixValue] | bool | None


def _serialise_config(value: _NixSettingValue) -> str:
    if is_value_of_type(value, list[str]):
        return " ".join(_serialise_config(e) for e in value)
    if is_value_of_type(value, bool):
        return "true" if value else "false"
    if is_value_of_type(value, str | int):
        return str(value)

    msg = f"Value is unsupported in nix config: {value!r}, must be {_NixSettingValue.__value__}"
    raise ValueError(msg)


def serialise_nix(value: _NixValue) -> str:
    """
    Serialises the given python object into a nix represenatation.
    NOTE: this interface does not allow accessing variables unless using interpolation, as all strings will be surrounded by quotes.
    """

    def escape(v: str) -> str:
        if "\\r" in v:
            raise ValueError("\\r is not supported for conversion")
        escaped = v.replace("\\", "\\\\").replace("$", "\\$").replace('"', '\\"')
        return f'"{escaped}"'

    if is_value_of_type(value, list[_NixValue.__value__]):
        return f"[{' '.join([serialise_nix(v) for v in value])}]"
    if is_value_of_type(value, dict[str, _NixValue.__value__]):
        return dedent(f"""
            {{
                {"\n                ".join([f"{escape(k)} = {serialise_nix(v)};" for k, v in value.items()])}
            }}
        """)
    if is_value_of_type(value, bool):
        return "true" if value else "false"
    if is_value_of_type(value, int | float):
        return str(value)
    if is_value_of_type(value, str):
        return escape(value)
    if is_value_of_type(value, None):
        return "null"

    msg = f"Value is unsupported in nix code: {value!r}"
    raise ValueError(msg)


class NixSettings:
    """Settings for invoking Nix"""

    def __init__(self):
        self._settings: dict[str, _NixSettingValue] = {
            # Running the test suite creates a lot of stores in the test root (somewhere under TMPDIR).
            # Obviously, they are not critical for system operation, so there is no need to reserve space.
            # The cleanup will only happen a couple of runs later, wasting space in the meantime.
            # Effectively disable this space reserve to reduce the waste considerably (by about 98%).
            "gc-reserved-space": 0,
            "show-trace": True,
            "sandbox": True,
            # explicitly disable substitution by default, otherwise we may attempt to contact
            # substituters and slow down many tests with pointless connection retry timeouts.
            "substituters": [],
            "extra-sandbox-paths": [],
            "extra-experimental-features": [],
            "extra-deprecated-features": [],
        }

    def __getattr__(self, attr: str) -> _NixSettingValue:
        if attr.startswith("__"):
            return super().__getattr__(attr)
        return self._settings[attr.replace("_", "-")]

    def __setattr__(self, attr: str, value: _NixSettingValue):
        if attr == "_settings":
            super().__setattr__(attr, value)
        else:
            self._settings[attr.replace("_", "-")] = value

    def __getitem__(self, attr: str) -> _NixSettingValue:
        return self._settings[attr]

    def __setitem__(self, attr: str, value: str):
        self._settings[attr] = value

    def add_xp_feature(self, *names: str):
        self["extra-experimental-features"] += names

    def add_dp_feature(self, *names: str):
        self["extra-deprecated-features"] += names

    def update(self, args: dict[str, _NixSettingValue] | None = None, **kwargs):
        """
        Overrides the settings with the given dict or kwargs.
        """
        self._settings.update(kwargs | (args or {}))

    def clone(self) -> "NixSettings":
        """
        shortcut to clone the settings to a new object
        """
        return self.with_settings()

    def with_settings(
        self, args: dict[str, _NixSettingValue] | None = None, **kwargs
    ) -> "NixSettings":
        """
        Copies the current settings into a new object, overriding the provided ones.

        :returns: A new Settings object with overridden settings
        """
        new_settings = NixSettings()
        new_settings._settings = copy.deepcopy(self._settings)
        new_settings.update(args, **kwargs)
        return new_settings

    def to_config(self, env: ManagedEnv) -> str:
        config = ""

        self["extra-sandbox-paths"] += env.path.to_sandbox_paths()

        def field_may(name: str, value: Any, serializer: Callable[[Any], str] = _serialise_config):
            nonlocal config
            if value is not None:
                config += f"{name} = {serializer(value)}\n"

        for name, value in self._settings.items():
            field_may(name, value)

        return config

    def to_env_overlay(self, env: ManagedEnv) -> None:
        cfg = self.to_config(env)
        (env.dirs.nix_conf_dir / "nix.conf").write_text(cfg)
        env.set_env("NIX_CONFIG", cfg)


@dataclasses.dataclass
class Nix:
    env: ManagedEnv
    logger: logging.Logger
    _settings: NixSettings | None = dataclasses.field(init=False, default=None)

    @property
    def _nix_executable(self) -> Path:
        if nix_bin_dir := self.env.dirs.nix_bin_dir:
            return Path(nix_bin_dir) / "nix"

        if from_path := shutil.which("nix"):
            return Path(from_path)

        raise ValueError(
            "Couldn't find a Nix command to execute! Set NIX_BIN_DIR or fix your environment"
        )

    @property
    def settings(self) -> NixSettings:
        """
        :return: the settings for the nix instance
        """
        if self._settings is None:
            self._settings = NixSettings()
            self._settings.store = f"local?root={self.env.dirs.test_root}"
            if sys.platform == "linux":
                # sandbox build dir cannot be withing store dir. choose a short non-overlapping path.
                self._settings.sandbox_build_dir = (
                    "/build-f2" if self.env.dirs.test_root.parts[1] != "build-f2" else "/build.f2"
                )

        return self._settings

    def nix_cmd(self, argv: list[str], flake: bool = False, cwd: Path | None = None) -> Command:
        """
        Constructs a NixCommand with the appropriate settings.
        """
        # Create a copy of settings to not have a writing side effect
        settings = self.settings.clone()
        if flake:
            settings.add_xp_feature("nix-command", "flakes")

        settings.to_env_overlay(self.env)
        return Command(argv=argv, exe=self._nix_executable, _env=self.env, cwd=cwd)

    def nix(
        self, cmd: list[str], nix_exe: str = "nix", flake: bool = False, cwd: Path | None = None
    ) -> Command:
        return self.nix_cmd([nix_exe, *cmd], flake=flake, cwd=cwd)

    # Mark each of these as correct as they are not ClassVars, but we also don't want to turn off RUF045
    nix_build = partialmethod(nix, nix_exe="nix-build")  # noqa: RUF045
    nix_shell = partialmethod(nix, nix_exe="nix-shell")  # noqa: RUF045
    nix_store = partialmethod(nix, nix_exe="nix-store")  # noqa: RUF045
    nix_env = partialmethod(nix, nix_exe="nix-env")  # noqa: RUF045
    nix_instantiate = partialmethod(nix, nix_exe="nix-instantiate")  # noqa: RUF045
    nix_channel = partialmethod(nix, nix_exe="nix-channel")  # noqa: RUF045
    nix_prefetch_url = partialmethod(nix, nix_exe="nix-prefetch-url")  # noqa: RUF045

    def eval(
        self, expr: str, settings: NixSettings | None = None, flags: list[str] | None = None
    ) -> CommandResult:
        """
        calls `nix eval --json --expr {expr}` using the given expression
        :param expr: what to evaluate
        :param settings: if none, the global settings will be used, otherwise the given one
        :param flags: if none, empty list, otherwise pass flags to the CLI invocation
        :return: result of the evaluation
        """
        if flags is None:
            flags = []
        orig = self.settings.clone()
        self._settings = settings or self.settings
        self.settings.add_xp_feature("nix-command")

        cmd = self.nix(["eval", "--json", *flags, "--expr", expr])
        # restore previous settings
        self._settings = orig
        return cmd.run()

    def eval_builtin(self, name: str, *args: _NixValue) -> CommandResult:
        """
        This is a high-level wrapper, to easily evaluate and obtain the result of a builtin.

        In the background, it calls `nix eval --json --expr builtins.{name} {args}`
        It is recommended to call `.json()` on the return value.

        :param name: name of the builtin
        :param args: list of arguments to provide to the builtin
        :return: result of the evaluation
        """
        args = [serialise_nix(arg) for arg in args]
        return self.eval(f"builtins.{name} {' '.join(args)}")

    @property
    def store_dir(self) -> Path:
        """
        The actual NIX_STORE_DIR this Nix command uses.
        """
        assert self.env.dirs.real_store_dir is not None, "bug in ManagedEnv"
        return self.env.dirs.real_store_dir

    def physical_store_path_for(self, path: str | Path) -> Path:
        """
        Takes a /nix/store/… path and rewrites it to be relative to this Nix's NIX_STORE_DIR.

        Nix accepts and returns store paths as `/nix/store` even when that's not where `NIX_STORE_DIR`
        physically is on the filesystem. Since we move the conceptual root for Nix to `test_root`,
        these "virtual" paths differ from the physical ones. So this function will convert `/nix/store`
        "virtual" paths to their real, physical location on the system.

        Basically, if you're passing it to `nix build` or `nix-store` or whatever, you want the
        `/nix/store` version. If you're passing it to a Python API (like pathlib.Path.exists()) or a
        command that operates on arbitrary files instead of store paths, you want the output of this
        function.


        :param path: a string or Path to convert
        :return: a Path object holding the rewritten, physical system path to the store entry
        """
        return (
            Path(str(path).replace("/nix/store", self.store_dir.as_posix()))
            if str(path).startswith("/nix/store")
            else Path(path)
        )

    def hash_path(self, store_path: str | Path, *args: str) -> str:
        """
        Shortcut to use `nix hash path {store_path}`, converting "virtual" store paths returned
        from Nix to their physical system paths including the test root.

        :param store_path: store path of the derivation or entry to hash
        """
        actual_path = self.physical_store_path_for(store_path).as_posix()
        res = self.nix(["hash", "path", actual_path, *args], flake=True).run().ok()
        return res.stdout_plain

    def clear_store(self):
        """
        Clears the test-owned store (and state) and resets them to an empty state
        """
        nix_store_dir = self.env.dirs.real_store_dir
        state_dir = self.env.dirs.nix_state_dir

        # Make store writable
        Command(["chmod", "-R", "+w", nix_store_dir], self.env).run().ok()
        shutil.rmtree(nix_store_dir)
        shutil.rmtree(state_dir)

        # Re-create the directories
        nix_store_dir.mkdir()
        state_dir.mkdir()


_fully_sandboxed = (
    sys.platform == "linux"
    and Path("/proc/self/ns/user").is_symlink()
    and subprocess.run(["unshare", "--user", "--mount", "--pid", "true"]).returncode == 0
)


def pytest_runtest_setup(item: Any):
    for mark in item.iter_markers(name="full_sandbox"):
        if not _fully_sandboxed:
            pytest.skip(f"{sys.platform} does not support full sandboxing")


@pytest.fixture
def nix(tmp_path: Path, env: ManagedEnv, logger: logging.Logger) -> Generator[Nix, Any, None]:
    """
    Provides a rich way of calling `nix`.
    For pre-applied commands use `nix.nix_instantiate`, `nix.nix_build` etc.
    After configuring the command, use `.run()` to run it
    """
    yield Nix(env, logger)
    # when things are done using the nix store, the permissions for the store are read only
    # after the test was executed, we set the permissions to rwx (write being the important part)
    # for pytest to be able to delete the files during cleanup
    cmd = Command(argv=["chmod", "-R", "+w", str(tmp_path.absolute())], _env=env)
    cmd.run().ok()


type NixDaemon = Callable[..., contextlib.AbstractAsyncContextManager[Nix]]

# NOTE: the order of items here is important. the daemon fixture requires
# the first item in this list to be the last socket opened by the daemon.
type NixDaemonProtocol = Literal["legacy-combined", "legacy", "lix-xp-1"]

daemon_protocols: list[NixDaemonProtocol] = get_args(NixDaemonProtocol.__value__)

_daemon_protocol_xp_features: dict[NixDaemon, list[str]] = {"lix-xp-1": ["rpc-sockets"]}


# paramterize every daemon tests to run using all supported nix protocols
@pytest.fixture(params=daemon_protocols)
def daemon(request: pytest.FixtureRequest) -> NixDaemon:
    default_protocol = request.param

    @contextlib.contextmanager
    def wrapper(
        nix: Nix,
        args: list[str] | None = None,
        settings: dict[str, _NixSettingValue] | None = None,
        protocol: NixDaemonProtocol | None = None,
        **kwargs,
    ) -> contextlib.AbstractAsyncContextManager[Nix]:
        protocol = protocol or default_protocol

        daemon = copy.deepcopy(nix)
        daemon.logger = nix.logger.getChild("daemon")
        daemon.settings["allowed-users"] = ["*"]
        daemon.settings["trusted-users"] = []
        daemon.settings.store = f"local?root={nix.env.dirs.test_root}"
        daemon.settings.update(settings)
        if requires_features := _daemon_protocol_xp_features.get(protocol):
            daemon.settings.add_xp_feature(*requires_features)

        sockets_dir = Path(daemon.env.dirs.nix_state_dir) / "daemon-socket"
        sockets = [sockets_dir / "socket", sockets_dir / "lix-xp-1/socket"]
        for p in sockets:
            p.unlink(missing_ok=True)

        proc = daemon.nix(args or [], nix_exe="nix-daemon", **kwargs).start()

        def log_daemon_result(result: CommandResult | None, level: int):
            if result:
                daemon.logger.log(level, "daemon exited with code %i", result.rc)
                daemon.logger.log(level, "stdout: %s", result.stdout_s)
                daemon.logger.log(level, "stderr: %s", result.stderr_s)
            else:
                daemon.logger.error("daemon exited unexpectedly")

        # wait for daemon to come up. this may take a while under load.
        # we wait only for the first socket in the list, expecting that
        # it'll be the last one opened by the daemon. this is to ensure
        # that we always return correctly regardless of rpc xp settings
        while not sockets[0].exists():
            if status := proc.wait(0.01):
                log_daemon_result(status, logging.ERROR)
                raise RuntimeError("daemon exited during startup")

        inner = copy.deepcopy(nix)
        socket_path = sockets_dir / "socket" if protocol == "legacy-combined" else sockets_dir
        inner.settings.store = f"unix://{socket_path}?protocol={protocol}"

        try:
            timeout, level = 1, logging.ERROR
            yield inner
            # 5 seconds should be enough to wait for a *graceful* exit.
            timeout, level = 5, logging.DEBUG
        finally:
            result = proc.terminate(timeout)
            if not result:
                result = proc.kill()
            log_daemon_result(result, level)

    return wrapper


@pytest.fixture
def enable_diverted_store(nix: Nix):
    """
    clear NIX_STORE_DIR, resetting it to the default (ie /nix/store).
    this makes builds impossible on platforms that cannot bind-mount,
    (e.g. macos) but it is important for eval result reproducibility.
    while builds may not work, substitution should still be possible.
    """
    nix.env.dirs.nix_store_dir = None
    nix.settings.sandbox_build_dir = None


def with_diverted_store(func: Callable[[Any], None]) -> Callable[[Any], None]:
    return pytest.mark.usefixtures("enable_diverted_store")(func)