import json
import os
import shlex

API = "http://localhost"
SOCKET = os.environ["FULL_SETUP_SOCKET"]
FORGE = "https://" + os.environ["FULL_SETUP_FORGE_DOMAIN"]
HIVE = os.environ["FULL_SETUP_HIVE"]
TOPLEVELS = json.loads(os.environ["FULL_SETUP_TOPLEVELS"])
REPO = os.environ["FULL_SETUP_HYPERHIVE_REPO"]
CONTROL_NIX = os.environ["FULL_SETUP_CONTROL_NIX"]
INIT = "/root/bao-init.json"
FORGE_AUTH = "/root/forge-auth.header"
NIX = "nix --extra-experimental-features nix-command"
ARCHETYPE = {"source": "hyperhive", "name": "reviewer"}
POLICY_UNITS_ACTIVE = (
    "for u in $(systemctl list-unit-files 'swarm-bao-*-policy.service' --no-legend | cut -d' ' -f1 | grep -v operator-viewer) "
    "swarm-bao-agent-pki.service; do "
    '[ "$(systemctl is-active $u)" = active ] || { echo "$u is not active"; exit 1; }; done'
)


def controller(method, path, body=None, check="--fail-with-body"):
    data = (
        f"-H 'Content-Type: application/json' -d {shlex.quote(json.dumps(body))}"
        if body
        else ""
    )
    return machine.succeed(
        f"curl {check} --silent --unix-socket {SOCKET} -X {method} {API}{path} {data}"
    )


def forge(method, path, body=None):
    data = (
        f"-H 'Content-Type: application/json' -d {shlex.quote(json.dumps(body))}"
        if body
        else ""
    )
    return machine.succeed(
        f"curl --fail-with-body --silent --insecure -H @{FORGE_AUTH} "
        f"-X {method} {FORGE}/api/v1{path} {data}"
    )


def diagnose(unit=None, container=None):
    # What a timed-out wait leaves behind: the failed units and the end of the
    # journal of the units that were waited on.
    where = f"--machine={container} " if container else ""
    units = [unit] if isinstance(unit, str) else list(unit or [])
    flags = " ".join(f"-u '{u}'" for u in units)
    names = " ".join(f"'{u}'" for u in units)
    for command in [
        f"systemctl {where}--failed --no-legend --no-pager",
        *(
            [
                f"systemctl {where}status {names} --no-pager -l",
                f"journalctl {where}{flags} --no-pager -n 50 | tail -c 6000",
            ]
            if units
            else []
        ),
    ]:
        print(f"--- {command}")
        print(machine.execute(command)[1])


def settle(command, timeout, unit=None, container=None):
    try:
        machine.wait_until_succeeds(command, timeout=timeout)
    except Exception:
        diagnose(unit, container)
        raise


def create_agent(body):
    # The controller refuses a create until its dependencies answer (the
    # archetype list comes from the builder's evaluation), so ask again.
    for _ in range(60):
        status, out = machine.execute(
            f"curl --silent --unix-socket {SOCKET} -o /root/create.out -w '%{{http_code}}' "
            f"-X POST {API}/api/agents -H 'Content-Type: application/json' -d {shlex.quote(json.dumps(body))}"
        )
        if out.strip().startswith("2"):
            return
        machine.sleep(10)
    diagnose("swarm-controller")
    # The queue refuses every client when its responder cannot answer. nats and
    # the responder run in the container; the units that fetch its keys and
    # certificate run on the host.
    diagnose(["nats.service", "swarm-nats-auth.service"], "swarm-nats")
    diagnose(["swarm-bao-nats-tls.service", "swarm-bao-nats-callout-xkey.service"])
    raise Exception(
        f"the controller refused to create {body['name']}: HTTP {out}: {machine.succeed('cat /root/create.out')}"
    )


machine.start()
machine.wait_for_unit("multi-user.target")

with subtest("the agents' toplevels are in the store, and the daemon builds nothing"):
    for toplevel in TOPLEVELS:
        machine.succeed(f"{NIX} path-info {toplevel}")
    # The control for the lookup: a path that is not there is not found.
    machine.fail(f"{NIX} path-info /nix/store/00000000000000000000000000000000-absent")
    # The control for the missing build: a derivation no store holds is refused.
    machine.fail(
        f"NIX_CONFIG='experimental-features = nix-command' nix build --no-link --impure -f {CONTROL_NIX}"
    )

# setup.md step 1: the secret store.
with subtest("step 1: bao operator init"):
    machine.wait_until_succeeds(
        '[ "$(bao status >/dev/null 2>&1; echo $?)" = 2 ]', timeout=600
    )
    machine.succeed(
        f"umask 077 && bao operator init -key-shares=1 -key-threshold=1 -format=json > {INIT}"
    )

with subtest("step 4: bao operator unseal, with the shamir seal"):
    machine.succeed(f"bao operator unseal \"$(jq -r '.unseal_keys_b64[0]' {INIT})\"")
    machine.succeed("bao status")

with subtest("step 1: the granter bootstrap"):
    machine.succeed(
        "set -o pipefail; "
        f'export BAO_TOKEN="$(jq -r .root_token {INIT})"; '
        "bao policy write bao-bootstrap /etc/hyperhive/bao-bootstrap-policy.hcl; "
        "bao token create -policy=bao-bootstrap -ttl=24h -orphan -display-name=bao-bootstrap -field=token "
        "| install -D -m 0600 /dev/stdin /var/lib/swarm-bao-bootstrap/grant.token; "
        "unset BAO_TOKEN; "
        "systemctl restart swarm-bao-granter-role"
    )
    machine.wait_until_succeeds(
        "journalctl -u swarm-bao-granter-role --no-pager | grep -F 'Uploaded policy: bao-granter'",
        timeout=300,
    )

with subtest("step 1: the policy units that waited for the granter are restarted"):
    # Some of these units retry every 30 seconds until something outside
    # them exists: the bao UI's viewer policy waits for authelia's secrets
    # unit, in its container, to publish the UI's OIDC client secret, and
    # that unit waits for the secret publisher's policy written here. So
    # the units are started without waiting, and polled until all are active.
    machine.succeed(
        "systemctl reset-failed 'swarm-bao-*-policy.service' swarm-bao-agent-pki.service || true"
    )
    machine.succeed(
        "systemctl restart --no-block 'swarm-bao-*-policy.service' swarm-bao-agent-pki.service"
    )
    machine.wait_until_succeeds(
        '[ "$(systemctl is-active swarm-bao-agent-pki.service)" = active ]', timeout=600
    )
    # Under the shamir seal the PKI unit leaves a restart of the store to the
    # operator when its listener's client CA changed, and says which command.
    if (
        machine.execute(
            "journalctl -u swarm-bao-agent-pki --no-pager | grep -F 'restart openbao.service'"
        )[0]
        == 0
    ):
        machine.succeed("systemctl --machine=swarm-bao restart openbao.service")
        machine.wait_until_succeeds(
            '[ "$(bao status >/dev/null 2>&1; echo $?)" = 2 ]', timeout=600
        )
        machine.succeed(
            f"bao operator unseal \"$(jq -r '.unseal_keys_b64[0]' {INIT})\""
        )
    settle(POLICY_UNITS_ACTIVE, 1200, "swarm-bao-*-policy")
    machine.succeed(
        "journalctl -u swarm-bao-controller-policy --no-pager "
        "| grep -F 'Data written to: auth/cert/certs/swarm-controller'"
    )

with subtest("authelia's mint has published the clients' secrets"):
    # The operator viewer policy is left out of the poll above: it waits for
    # the UI client's secret, which authelia's secrets unit publishes from
    # its container once the store and the trust bundle it needs exist. The
    # controller's client secret is published by the same unit, and the
    # controller, the hive's store logins and the UI policy all wait for it.
    settle(
        f'BAO_TOKEN="$(jq -r .root_token {INIT})" bao kv get -field=value '
        "secret/swarm/controller/swarm-controller/oidc/client >/dev/null",
        1200,
        "authelia-swarm-secrets",
        "swarm-authelia",
    )
    settle(
        '[ "$(systemctl is-active swarm-bao-operator-viewer-policy.service)" = active ]',
        300,
        "swarm-bao-operator-viewer-policy",
    )

with subtest("step 2: swarmctl user add"):
    machine.wait_until_succeeds(
        "swarmctl user add mara --display-name Mara --email mara@example.com --group admins "
        ">/root/user-add.out 2>&1",
        timeout=900,
    )
    machine.succeed("[ -s /root/user-add.out ]")
    machine.succeed("swarmctl user list | grep -w mara")

with subtest("the forge holds hyperhive/hyperhive at the fixture's commit"):
    # The controller's forge token, from the store as the controller reads it.
    # Minted and stored by `forgejo-swarm-controller-account` in the forge
    # container, which waits for forgejo and for a login to the store.
    settle(
        f'BAO_TOKEN="$(jq -r .root_token {INIT})" bao kv get -field=value '
        "secret/swarm/controller/swarm-controller/forge/token > /root/forge-token",
        900,
        ["forgejo", "forgejo-swarm-controller-account", "forgejo-sso-source"],
        "hive-forge",
    )
    machine.succeed(
        f"umask 077 && printf 'Authorization: token %s\\n' \"$(cat /root/forge-token)\" > {FORGE_AUTH}"
    )
    # The token being published says nothing about the gateway's vhost for
    # the forge, which the calls below go through.
    machine.wait_until_succeeds(
        f"curl --fail --silent --insecure -H @{FORGE_AUTH} {FORGE}/api/v1/user >/dev/null",
        timeout=600,
    )
    status = machine.succeed(
        f"curl --silent --insecure -o /dev/null -w '%{{http_code}}' -H @{FORGE_AUTH} {FORGE}/api/v1/orgs/hyperhive"
    ).strip()
    if status != "200":
        forge("POST", "/orgs", {"username": "hyperhive", "visibility": "public"})
    forge(
        "POST",
        "/orgs/hyperhive/repos",
        {"name": "hyperhive", "private": False, "default_branch": "main"},
    )
    machine.succeed(
        f"cp -r {REPO}/repo.git /root/hyperhive.git && chmod -R u+w /root/hyperhive.git"
    )
    machine.succeed(
        f'GIT_SSL_NO_VERIFY=1 git -C /root/hyperhive.git -c "http.extraHeader=$(cat {FORGE_AUTH})" '
        f"push {FORGE}/hyperhive/hyperhive.git main"
    )
    pushed = machine.succeed(
        f'GIT_SSL_NO_VERIFY=1 git -c "http.extraHeader=$(cat {FORGE_AUTH})" '
        f"ls-remote {FORGE}/hyperhive/hyperhive.git refs/heads/main"
    ).split()[0]
    assert pushed == machine.succeed(f"cat {REPO}/rev").strip(), pushed

with subtest(
    "step 8: spawn, a reviewer-archetype agent, through the endpoint the swarm UI calls"
):
    machine.wait_until_succeeds(
        f"curl --fail --silent --unix-socket {SOCKET} {API}/api/permissions >/dev/null",
        timeout=900,
    )
    # A custom root agent (ruth, as the UI creates her) is a follow-up, not covered yet.
    create_agent({"name": "spawn", "hive": HIVE, "archetype": ARCHETYPE})

with subtest("step 9: the first agent"):
    create_agent({"name": "iris", "hive": HIVE, "archetype": ARCHETYPE})

with subtest("internal/meta holds the flake the toplevels were built from"):
    settle(
        f"curl --fail --silent --insecure -H @{FORGE_AUTH} {FORGE}/api/v1/repos/internal/meta/raw/flake.nix "
        ">/root/meta-flake.nix",
        1800,
        "swarm-controller",
    )
    rev = machine.succeed(f"cat {REPO}/rev").strip()
    settings = json.loads(os.environ["FULL_SETUP_META_SETTINGS"])
    expected = (
        open(os.environ["FULL_SETUP_META_TEMPLATE"])
        .read()
        .replace(
            "@HYPERHIVE_URL@",
            f"git+{settings['git_base']}/hyperhive/hyperhive.git?rev={rev}",
        )
    )
    for name, value in settings["vars"].items():
        expected = expected.replace(f"@{name}@", value)
    committed = machine.succeed("cat /root/meta-flake.nix")
    assert committed == expected, (
        "the controller's flake differs from the one the toplevels came from:\n"
        + committed
    )

with subtest("swarmctl setup status reports the steps done so far"):
    # Step 9 names ruth, the root agent of the real setup; this test creates none, so the report
    # ends on step 9 not done and the exit status is 1.
    reported = [
        "✅   1  secret store initialised",
        "✅   2  secret store unsealed",
        "✅   3  granter bootstrapped",
        "✅   4  granting units active",
        "✅   5  store unsealed after the granting units",
        "✅   6  SSO account in the admins group",
        "👉   9  ruth created and up",
        "✅  10  another agent [optional]",
    ]
    for _ in range(60):
        lines = machine.execute("swarmctl setup status")[1].splitlines()
        if all(any(line.startswith(prefix) for line in lines) for prefix in reported):
            break
        machine.sleep(10)
    else:
        raise Exception(
            "swarmctl setup status did not report the steps done so far:\n"
            + "\n".join(lines)
        )
    assert machine.execute("swarmctl setup status >/dev/null")[0] == 1

with subtest("the agents are deployed and need a login"):
    settle(
        f"curl --fail --silent --unix-socket {SOCKET} {API}/api/agents/status "
        "| jq -e '[.[] | select(.hive == \""
        + HIVE
        + '" and .freshness == "fresh" and .needs_login == true)'
        ' | .name] | map(select(. == "iris" or . == "spawn")) | unique | length == 2\'',
        2400,
        ["swarm-controller", "hive-c0re"],
    )