#!/usr/bin/env python3

import os
import json
from time import sleep

import requests

import config

DOCKER_COMPOSE_PATH = config.portainer_compose_dir + '/{}/docker-compose.yml'
SCRIPT_DIR = os.path.dirname(__file__)
BASH_COMPANION_SCRIPT = f'{SCRIPT_DIR}/docker_updater.sh'
REPO_UPDATER_SCRIPT = f'{SCRIPT_DIR}/update_repos.sh'


class PortainerClient:

    def __init__(self, portainer_url, access_token):
        self.portainer_url = portainer_url
        self.access_token = access_token

        self.api_url = f"{self.portainer_url}/api"

    def _get_headers(self):
        return {
            "accept": "application/json",
            "X-API-Key": self.access_token
        }

    def _get(self, uri):
        response = requests.get(f"{self.api_url}/{uri}", headers=self._get_headers())
        return response.json()

    def list_endpoints(self):
        return self._get("endpoints")

    def list_stacks(self, endpoint_id=None):
        filters = f'?filters={{"EndpointID":{endpoint_id}}}' if endpoint_id else ""
        return self._get(f'stacks{filters}')

    def get_stack(self, stack_id):
        return self._get(f'stacks/{stack_id}')

    def get_stack_file(self, stack_id):
        return self._get(f'stacks/{stack_id}/file')

    def _put(self, uri, payload):
        response = requests.put(
            f"{self.api_url}/uri",
            headers=self._get_headers(),
            json=payload
        )

        return response.json()

    def upgrade_stack(self, stack):
        file_response = self.get_stack_file(stack['Id'])
        payload = {
            "pullImage": True,
            "prune": True,
            "StackFileContent": file_response["StackFileContent"],
            "Env": stack["Env"]
        }

        return self._put(f"stacks/{stack['Id']}?endpointId={stack['EndpointId']}", payload)

    def list_stack_containers(self, endpoint, stack_name):
        filters = '{"label":["com.docker.compose.project=' + stack_name + '"]}'
        return self._get(f'endpoints/{endpoint}/docker/containers/json?all=1&filters={filters}')

    def update_image(self, endpoint, image):
        response = requests.post(
            f"{self.api_url}/endpoints/{endpoint}/docker/images/create?fromImage={image}",
            headers=self._get_headers()
        )

        try:
            status = json.loads(response.text.strip().split('\r\n')[-1])['status']
        except Exception:
            return False

        return f"Downloaded newer image for {image}" in status

    def delete_image(self, endpoint, image):
        response = requests.delete(
            f"{self.api_url}/endpoints/{endpoint}/docker/images/{image}?force=false",
            headers=self._get_headers()
        )

    def ping_server(self):
        try:
            r = requests.get(f"{self.api_url}/", headers=self._get_headers())
        except requests.exceptions.ConnectionError:
            return False

        return r.status_code == 200


def update_repositories():
    for repo_path in config.git_repositories:
        ret = os.system(f"{REPO_UPDATER_SCRIPT} {repo_path}")
        if ret > 0:
            exit(ret)


def update_stack_images(endpoint, stack_name):
    containers = client.list_stack_containers(endpoint, stack_name)

    updated_images = []
    for c in containers:
        if client.update_image(endpoint, c['Image']):
            print(f"{c['Image']} was updated.")
            updated_images.append(c['ImageID'])

    return updated_images


def upgrade_local_stack(stack):
    project_name = stack['Name']

    print(f"Updating stack: {project_name}")

    compose_file = DOCKER_COMPOSE_PATH.format(stack['Id'])
    ret = os.system(f"{BASH_COMPANION_SCRIPT} {project_name} {compose_file}")
    if ret > 0:
        exit(ret)

    print(f"Stack {project_name} updated!\n")


def upgrade_endpoint_stack(client, endpoint):
    endpoint_id = endpoint['Id']
    endpoint_updated_images = []
    for stack in client.list_stacks(endpoint_id):
        if stack['Status'] == 1:
            stack_name = stack['Name']

            print(f"Checking {stack_name} for updates.")

            updated_images = update_stack_images(endpoint_id, stack_name)
            if updated_images:
                endpoint_updated_images += updated_images
                print(f"{stack_name} was updated. Restarting...")
                if endpoint_id == 2 and stack_name in ('portainer', 'traefik'):
                    upgrade_local_stack(stack)
                    while not client.ping_server():
                        print("Waiting for server to be back online...")
                        sleep(1)
                else:
                    client.upgrade_stack(stack)
                print(f"{stack_name} restarted.")
            else:
                print(f"{stack_name} is up to date. Skipping.")

    if endpoint_updated_images:
        print("Deleting obsolete images")
        for image_id in endpoint_updated_images:
            client.delete_image(endpoint_id, image_id)
    else:
        print("No obsolete image to delete")


if __name__ == '__main__':
    update_repositories()

    client = PortainerClient(portainer_url=config.portainer_url, access_token=config.access_token)

    for e in client.list_endpoints():
        upgrade_endpoint_stack(client, e)