#!/usr/bin/env python3
"""
Wrap the devcontrol binary used for running a local instance of the
Tailscale control plane.

I use a tiny subset of the available flags and scenarios; this script
wraps them in a friendly interface.
"""

import argparse
import atexit
import os
from pathlib import Path
import subprocess
import sys
from typing import TypedDict

from chives.text import coloured


Args = TypedDict(
    "args",
    {
        "audit_logs": bool,
        "persist": bool,
        "tfacc": bool,
    },
)


def parse_args():
    parser = argparse.ArgumentParser(
        prog="devcontrol",
        description="Run devcontrol with different options",
    )

    parser.add_argument("--audit-logs", action="store_true", help="save audit logs")
    # TODO: Can I make this a flag to store state in a specific dir?
    parser.add_argument(
        "--persist",
        action="store_true",
        help="save the devcontrol state in /tmp/devcontrol",
    )
    parser.add_argument(
        "--tfacc",
        action="store_true",
        help="run the Terraform acceptance test scenario",
    )

    args = parser.parse_args()

    return {
        "audit_logs": args.audit_logs,
        "persist": args.persist,
        "tfacc": args.tfacc,
    }


if __name__ == "__main__":
    args = parse_args()

    corp_dir = Path.home() / "repos/corp"

    tailcontrol_cmd = ["./tool/go", "run", "--tags=tailscale_saas", "./cmd/devcontrol"]
    tailcontrol_env = {
        "HOME": str(Path.home()),
        "PATH": os.environ["PATH"],
    }

    # See https://www.notion.so/tailscale/Running-a-local-CONTROL-clients-and-DERP-d69826480c704c1f9f39950478e7303f#f744c21f49f14c4fa748bff0a944e9bb
    if args["audit_logs"]:
        # Compile logzservice before running it, so the terminate() signal
        # goes to the logzservice process rather than `go run`.
        #
        # Otherwise, the compiled logzservice binary can be left running
        # after this script exits.
        subprocess.check_call(["./tool/go", "build", "./cmd/logzservice"], cwd=corp_dir)
        logz_proc = subprocess.Popen(["./logzservice", "-dev"], cwd=corp_dir)
        atexit.register(lambda: logz_proc.terminate())
        tailcontrol_env.update(
            {
                "TS_LOG_TARGET": "http://localhost:9951",
                "TS_DEBUG_CONTROL_LOGTAIL_UPLOAD": "true",
            }
        )

        def stop_logz_proc() -> None:
            logz_proc.terminate()

        atexit.register(stop_logz_proc)

    if args["persist"]:
        tailcontrol_cmd.append("--dir=/tmp/devcontrol")

    if args["tfacc"]:
        tailcontrol_cmd.append("--generate-test-devices=terraform-acceptance-testing")

    debug_string = ["->"]
    for name, value in tailcontrol_env.items():
        if name in {"HOME", "PATH"}:
            continue
        debug_string.append(f"{name}={value}")
    debug_string.extend(tailcontrol_cmd)

    print(coloured(" ".join(debug_string), "blue"))
    try:
        subprocess.check_call(tailcontrol_cmd, env=tailcontrol_env, cwd=corp_dir)
    except KeyboardInterrupt:
        sys.exit(0)