# SPDX-FileCopyrightText: 2026 Qyriad <qyriad@qyriad.me>
#
# SPDX-License-Identifier: EUPL-1.1

import functools
from pathlib import Path
import json
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"
testing_client = "/root/.nix-profile/libexec/dynix-testing-client.py"

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)

    try:
        config_file_path.unlink()
    except Exception as e:
        machine.log(f"Couldn't unlike path {config_file_path}: {e}")
        raise
    return config_data

@beartype
def dynix_append_daemon(option: str, value: Any):
    payload = json.dumps(dict(
        action="append",
        args=dict(
            name=option,
            value=value,
        ),
    ))

    machine.succeed(f"echo '{payload}' | {testing_client} /run/user/0/dynix.sock")

@beartype
def dynix_append_cli(option: str, value: Any):
    machine.succeed(f'''
        dynix append {shlex.quote(option)} {shlex.quote(str(value))}
    '''.strip())

    expr = textwrap.dedent("""
        (import <nixpkgs/nixos> { }).config.dynamicism.applyDynamicConfiguration { }
    """).strip()
    machine.succeed(rf"""
        nix run --show-trace --log-format raw-with-logs --impure -E {shlex.quote(expr)}
    """.strip())

@beartype
def run_all_tests(machine: Machine, *, use_daemon: bool):
    dynix_append = dynix_append_daemon if use_daemon else dynix_append_cli

    #
    # Setup.
    #
    dynix_out = machine.succeed("dynix --version")
    assert "dynix" in dynix_out, f"dynix not in {dynix_out=}"

    machine.succeed("systemctl start user@0.service")
    machine.wait_for_unit("user@0.service")

    run_log(machine, "systemctl start dynix-daemon.service")
    machine.wait_for_unit("dynix-daemon.service")

    machine.log("Checking initial harmonia.service conditions")

    # Config should have our initial values.
    config_toml = get_config_file()
    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("env PAGER= nixos-rebuild switch --log-format raw-with-logs --no-reexec -v --fallback")

    # Config should not have changed.
    config_toml = get_config_file()
    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"

    machine.log("Testing dynamic workers=20")
    new_workers = 20
    dynix_append("services.harmonia.settings.workers", new_workers)

    machine.log("Testing that workers, but not max_connection_rate, changed")
    # Workers, but not max connection rate, should have changed.
    config_toml = get_config_file()
    from pprint import pformat
    machine.log(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"

    machine.log("Testing dynamic max_connection_rate=100")
    new_max_connection_rate = 100
    dynix_append("services.harmonia.settings.max_connection_rate", new_max_connection_rate)

    # Max connection rate should have changed, and workers should be the same as before.
    config_toml = get_config_file()
    print(f"checking connection rate, {use_daemon=}")
    assert int(config_toml['max_connection_rate']) == new_max_connection_rate, f"{config_toml['max_connection_rate']=} != {new_max_connection_rate}"
    print(f"checking workers, {use_daemon=}")
    assert int(config_toml['workers']) == new_workers, f"{config_toml['workers']=} != {new_workers}"

    machine.log("Done with tests; stopping dynix-daemon")
    machine.succeed("systemctl stop dynix-daemon.service")

    # And this should set everything back.
    machine.succeed("env PAGER= nixos-rebuild switch --log-format raw-with-logs --no-reexec -v --fallback")
    machine.wait_for_unit("harmonia.service")
    config_toml = get_config_file()
    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'

machine.start(allow_reboot=True)
machine.wait_for_unit("install-dynix.service")
try:
    run_all_tests(machine, use_daemon=False)
except Exception as e:
    machine.log(f"ERROR running CLI tests: {e}")

machine.reboot()

machine.wait_for_unit("install-dynix.service")
try:
    run_all_tests(machine, use_daemon=True)
except Exception as e:
    machine.log(f"ERROR running DAEMON tests: {e}")
    raise