import argparse
import functools
#from pathlib import Path
#from pprint import pformat
import shlex
import textwrap
from typing import 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

parser = argparse.ArgumentParser()
#parser.add_argument("--no-detach", action="store_true")
#parser.add_argument("--daemon", action="store_true")
#parser.add_argument("--enable-tcp-insecure", action="store_true")
#parser.add_argument("--port", type=int)
parser.add_argument("--jobs", type=int)
parser.add_argument("--job-lifetime", type=int)
parser.add_argument("--log-level", type=str)
#parser.add_argument("--nice", type=int)
#parser.add_argument("--stats", action="store_true")
#parser.add_argument("--stats-port", type=int)

@beartype
def get_cli_args() -> argparse.Namespace:
    machine.wait_for_unit("distccd.service")
    mainpid = int(machine.get_unit_property("distccd.service", "MainPID"))
    machine.log(f"{mainpid=}")
    pidtext = machine.succeed(f"pgrep -P {mainpid}")
    machine.log(f"{pidtext=}")
    pid = int(pidtext.splitlines()[0])
    machine.log(f"{pid=}")
    execstart = machine.get_unit_property("distccd.service", "ExecStart")
    print(f"{execstart=}")

    cmdline = machine.succeed(f"cat /proc/{pid}/cmdline")
    cmdline_args = cmdline.split("\0")
    machine.log(f"{cmdline_args=}")
    print(f"{cmdline_args=}")

    #return shlex.join(cmdline_args[1:])
    args, rest = parser.parse_known_args(cmdline_args)
    return args

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

# Config should have our initial values.
args = get_cli_args()
#assert '--jobs 12' in args, f'--jobs 12 not in {args=}'
assert args.jobs == 12, f'{args.jobs=} != 12'
assert args.job_lifetime == 900, f'{args.job_lifetime} != 900'
assert args.log_level == 'warning', f'{args.log_level=} != warning'

with machine.nested("must succeed: initial nixos-rebuild switch"):
    machine.succeed("env PAGER= nixos-rebuild switch --log-format raw-with-logs --fallback")

# Config should not have changed.
args = get_cli_args()
#machine.log(f"config.toml after first rebuild: {indent(pformat(args))}")
#assert int(args['workers']) == 4, f"{args['workers']=} != 4"
#assert int(args['max_connection_rate']) == 256, f"{args['max_connection_rate']=} != 256"
#
new_jobs = 4
expr = textwrap.dedent(f"""
    let
        nixos = import <nixpkgs/nixos> {{ }};
    in nixos.config.dynamicism.doChange {{
        option = "services.distccd.maxJobs";
        value = {new_jobs};
    }}
""").strip()
machine.succeed(rf"""
    nix run --show-trace --log-format raw-with-logs --impure -E {shlex.quote(expr)}
""".strip())

args = get_cli_args()

# Only jobs should have changed. The others should still be default.
assert args.jobs == new_jobs, f'{args.jobs=} != {new_jobs=}'
assert args.job_lifetime == 900, f'{args.job_lifetime} != 900'
assert args.log_level == 'warning', f'{args.log_level=} != warning'

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

args = get_cli_args()
#assert args.jobs == new_jobs, f'{args.jobs=} != {new_jobs=}'
#assert args.job_lifetime == 900, f'{args.job_lifetime} != 900'
assert args.log_level == new_log_level, f'{args.log_level=} != {new_log_level=}'