#!/usr/bin/env python3

import sys

import click
import requests


@click.command()
@click.argument("job")
@click.option("--dc", default="bom1", help="Name of the data center")
@click.option("--token", help="Token for Nomad")
def cli(job, dc, token):
    if token is None:
        print("you need to pass a valid token")
        sys.exit(1)

    headers = {"Authorization": f"Bearer {token}"}
    url = f"https://{dc}-nomad.simulprod.com/v1/job/{job}/allocations"
    try:
        resp = requests.get(url, headers=headers)
        resp.raise_for_status()
    except Exception as e:
        print("return {}".format(str(e)))

    running_tasks = []
    terminated_tasks = []
    for task in resp.json():
        task_name = list(task["TaskStates"].keys())[0]
        if task["TaskStates"][task_name]["State"] == "running":
            running_tasks.append(f"https://{dc}-nomad.simulprod.com/ui/allocations/{task['ID']}/{task_name}/logs")
        else:
            terminated_tasks.append(f"https://{dc}-nomad.simulprod.com/ui/allocations/{task['ID']}/{task_name}/logs")

    if len(running_tasks) > 0:
        print("running tasks")
        for t in running_tasks:
            print(f"→ {t}")

    if len(terminated_tasks) > 0:
        print("terminated tasks")
        for t in terminated_tasks:
            print(f"→ {t}")


if __name__ == "__main__":
    cli()