import functools
from pathlib import Path
from pprint import pformat
import shlex
import textwrap
import tomllib
from typing import Any, cast, TYPE_CHECKING

from beartype import beartype

from test_driver.machine import Machine
from test_driver.errors import RequestedAssertionFailed

if TYPE_CHECKING:
    global machine
    machine = cast(Machine, ...)
    assert machine.shell is not None

ls = "eza -lah --color=always --group-directories-first"

indent = functools.partial(textwrap.indent, prefix='  ')

@beartype
def run_log(machine: Machine, *commands: str, timeout: int | None = 60) -> str:
    output = ""
    for command in commands:
        with machine.nested(f"must succeed: {command}"):
            (status, out) = machine.execute(f"{command} | tee /dev/stderr", timeout=timeout)
            if status != 0:
                machine.log(f"output: {out}")
                raise RequestedAssertionFailed(
                    f"command `{command}` failed (exit code {status})",
                )
            output += out

    return output

@beartype
def get_config_file() -> dict[str, Any]:
    machine.wait_for_unit("harmonia.service")
    pid = int(machine.get_unit_property("harmonia.service", "MainPID"))
    env_lines: list[str] = machine.succeed(f"cat /proc/{pid}/environ").replace("\0", "\n").splitlines()
    pairs: list[list[str]] = [line.split("=", maxsplit=1) for line in env_lines]
    env = dict(pairs)

    config_file = Path(env["CONFIG_FILE"])

    machine.log(f"copying from VM: {config_file=}")
    machine.copy_from_vm(config_file.as_posix())

    config_file_path = machine.out_dir / config_file.name
    with open(config_file_path, "rb") as f:
        config_data = tomllib.load(f)

    config_file_path.unlink()
    return config_data

machine.wait_for_unit("default.target")
assert "lix" in machine.succeed("nix --version").lower()
machine.log("INIT")

# Config should have our initial values.
config_toml = get_config_file()
machine.log(f"config.toml BEFORE first rebuild (initial): {indent(pformat(config_toml))}")
assert int(config_toml['workers']) == 4, f"{config_toml['workers']=} != 4"
assert int(config_toml['max_connection_rate']) == 256, f"{config_toml['max_connection_rate']=} != 256"

with machine.nested("must succeed: initial nixos-rebuild switch"):
    machine.succeed("nixos-generate-config")
    machine.succeed("mkdir -vp /etc/nixos")
    # Dereference is required since that configuration.nix is probably a symlink to the store.
    machine.succeed("cp -rv --dereference /run/current-system/sw/share/nixos/configuration.nix /etc/nixos/")
    machine.succeed("env PAGER= nixos-rebuild switch --log-format raw-with-logs --fallback")

# Config should not have changed.
config_toml = get_config_file()
machine.log(f"config.toml after first rebuild: {indent(pformat(config_toml))}")
assert int(config_toml['workers']) == 4, f"{config_toml['workers']=} != 4"
assert int(config_toml['max_connection_rate']) == 256, f"{config_toml['max_connection_rate']=} != 256"

new_workers = 20
expr = textwrap.dedent(f"""
    let
        nixos = import <nixpkgs/nixos> {{ }};
    in nixos.config.dynamicism.doChange {{
        option = "services.harmonia.settings.workers";
        value = {new_workers};
    }}
""").strip()
machine.succeed(rf"""
    nix run --show-trace --log-format raw-with-logs --impure -E {shlex.quote(expr)}
""".strip())

# Workers, but not max connection rate, should have changed.
config_toml = get_config_file()
machine.log(f"config.toml after DYNAMIC ACTIVATION: {indent(pformat(config_toml))}")
assert int(config_toml['workers']) == new_workers, f"{config_toml['workers']=} != {new_workers}"
assert int(config_toml['max_connection_rate']) == 256, f"{config_toml['max_connection_rate']=} != 256"

new_max_connection_rate = 100
expr = textwrap.dedent(f"""
    let
        nixos = import <nixpkgs/nixos> {{ }};
    in nixos.config.dynamicism.doChange {{
        option = [ "services" "harmonia" "settings" "max_connection_rate" ];
        value = {new_max_connection_rate};
    }}
""").strip()
machine.succeed(rf"""
    nix run --show-trace --log-format raw-with-logs --impure -E {shlex.quote(expr)}
""".strip())

# Max connection rate should have changed, but workers should have reverted.
config_toml = get_config_file()
machine.log(f"config_toml after SECOND dynamic activation: {indent(pformat(config_toml))}")
assert int(config_toml['max_connection_rate']) == new_max_connection_rate, f"{config_toml['max_connection_rate']=} != {new_max_connection_rate}"

# And this should set everything back.
machine.succeed("env PAGER= nixos-rebuild switch --log-format raw-with-logs --fallback")
machine.systemctl("restart harmonia.service")
machine.wait_for_unit("harmonia.service")
config_toml = get_config_file()
machine.log(f"config_toml after NORMAL activation: {indent(pformat(config_toml))}")
assert int(config_toml['max_connection_rate']) == 256, f'{config_toml["max_connection_rate"]=} != 256'
assert int(config_toml['workers']) == 4, f'{config_toml["workers"]=} != 4'