import argparse
import time

from kubernetes.client.rest import ApiException

from polyaxon_client.client import PolyaxonClient
from polyaxon_k8s.manager import K8SManager
from sidecar import settings
from sidecar.monitor import is_pod_running

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--app_label',
        type=str
    )
    parser.add_argument(
        '--container_id',
        type=str
    )
    parser.add_argument(
        '--sleep_interval',
        default=2,
        type=int
    )
    parser.add_argument(
        '--max_restarts',
        default=0,
        type=int
    )
    args = parser.parse_args()
    arguments = args.__dict__

    container_id = arguments.pop('container_id')
    app_label = arguments.pop('app_label')
    sleep_interval = arguments.pop('sleep_interval')
    max_restarts = arguments.pop('max_restarts')

    k8s_manager = K8SManager(namespace=settings.K8S_NAMESPACE, in_cluster=True)
    client = PolyaxonClient()
    client.set_internal_health_check()
    retry = 0
    is_running = True
    status = None
    while is_running and retry < 3:
        time.sleep(sleep_interval)
        try:
            is_running, status = is_pod_running(k8s_manager,
                                                settings.POD_ID,
                                                container_id,
                                                max_restarts)
        except ApiException:
            retry += 1
            time.sleep(sleep_interval)  # We wait a bit more before try

    if status:
        client.reconcile(status=status)

[((12, 13, 12, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((40, 18, 40, 79), 'polyaxon_k8s.manager.K8SManager', 'K8SManager', (), '', False, 'from polyaxon_k8s.manager import K8SManager\n'), ((41, 13, 41, 29), 'polyaxon_client.client.PolyaxonClient', 'PolyaxonClient', ({}, {}), '()', False, 'from polyaxon_client.client import PolyaxonClient\n'), ((47, 8, 47, 34), 'time.sleep', 'time.sleep', ({(47, 19, 47, 33): 'sleep_interval'}, {}), '(sleep_interval)', False, 'import time\n'), ((49, 33, 52, 61), 'sidecar.monitor.is_pod_running', 'is_pod_running', ({(49, 48, 49, 59): 'k8s_manager', (50, 48, 50, 63): 'settings.POD_ID', (51, 48, 51, 60): 'container_id', (52, 48, 52, 60): 'max_restarts'}, {}), '(k8s_manager, settings.POD_ID, container_id, max_restarts)', False, 'from sidecar.monitor import is_pod_running\n'), ((55, 12, 55, 38), 'time.sleep', 'time.sleep', ({(55, 23, 55, 37): 'sleep_interval'}, {}), '(sleep_interval)', False, 'import time\n')]