#!/usr/bin/env python3

import yaml
import subprocess
import time
import sys

groupyaml = '/etc/hiera/data/group.yaml'

def get_frontends(data):
  try:
    return list(data['sunet_frontend']['load_balancer']['websites'].keys())
  except KeyError:
    return []

def check_docker_instance_status(instance):
  cmd = f"docker inspect -f {r'{{.State.Status}}'} {instance}"
  result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
  return result.stdout.strip() == 'running'

def is_exabgp_running():
  cmd = ["systemctl", "is-active", "exabgp.service"]
  result = subprocess.run(cmd, capture_output=True, text=True)
  return result.stdout.strip() == 'active'

def check_docker_instances(instances, max_retries=3, initial_wait=10):
  for instance in instances:
    retries = 0
    while retries < max_retries:
      if check_docker_instance_status(instance):
        print(f"Instance: {instance} is running!")
        break
      else:
        print(f"Instance: {instance} is not running! Will try again in {initial_wait * (2**retries)} seconds.")
        time.sleep(initial_wait * (2**retries))
        retries += 1
    if retries == max_retries:
      print(f"Max retries reached for instance: {instance}, exiting!")
      sys.exit(1)

def check_exabgp_running(max_retries=3, initial_wait=10):
  retries = 0
  while retries < max_retries:
    if is_exabgp_running():
      print("ExaBGP service is running!")
      break
    else:
      print(f"Exabgp is not running! Will try again in {initial_wait * (2**retries)} seconds.")
      time.sleep(initial_wait * (2**retries))
      retries += 1
  if retries == max_retries:
    print(f"Max retries reached for checking if exabgp is running, exiting!")
    sys.exit(1)

with open(groupyaml, 'r') as f:
  data = yaml.safe_load(f)

frontends = get_frontends(data)
instances = []

for frontend in frontends:
  instances.append(frontend + '-haproxy-1')
  instances.append(frontend + '-monitor-1')
  instances.append(frontend + '-config-1')

instances.append('frontend-api-1')
instances.append('frontend-telegraf-1')

check_exabgp_running()
check_docker_instances(instances)
sys.exit(0)