#! /usr/bin/env python3
from argparse import ArgumentParser
from os import environ
from typing import Any
from typing import cast

import requests_unixsocket
from requests.exceptions import HTTPError
from requests.exceptions import JSONDecodeError


session = requests_unixsocket.Session()


NOMAD_ADDR = environ.get("NOMAD_ADDR", "http://127.0.0.1:4646")
NOMAD_TOKEN = environ.get("NOMAD_TOKEN")


def nomad_req(
    *path: str, params: dict[str, Any] | None = None, method="GET"
) -> list[dict[str, Any]] | dict[str, Any] | str:
    headers = {}
    if NOMAD_TOKEN:
        headers["X-Nomad-Token"] = NOMAD_TOKEN

    response = session.request(
        method,
        f"{NOMAD_ADDR}/v1/{'/'.join(path)}",
        params=params,
        headers=headers,
    )
    response.raise_for_status()

    try:
        return response.json()
    except JSONDecodeError:
        return response.text


exit_code = 0
parser = ArgumentParser(
    description="Checks for orphaned services and optionally deletes them.",
)
parser.add_argument(
    "-d", "--delete", action="store_true", help="Delete orphan services"
)
args = parser.parse_args()


for namespace in nomad_req("services"):
    namespace = cast(dict[str, Any], namespace)
    for service in namespace["Services"]:
        service_name = service["ServiceName"]
        for service_instance in nomad_req("service", service_name):
            service_instance = cast(dict[str, Any], service_instance)
            service_id = service_instance["ID"]
            alloc_id = service_instance["AllocID"]

            alloc_found = True

            try:
                alloc = nomad_req("allocation", alloc_id)
                continue
            except HTTPError as e:
                if e.response.status_code == 404:
                    alloc_found = False
                    message = f"alloc {alloc_id} not found for {service_name}."
                    if args.delete:
                        message += f" Deleting {service_id}"

                    print(message)
                else:
                    raise e

            if not alloc_found and args.delete:
                nomad_req("service", service_name, service_id, method="DELETE")


exit(exit_code)