import os
import json
from time import sleep

import requests

import config


DOCKER_COMPOSE_PATH = config.portainer_compose_dir + '/{}/docker-compose.yml'
BASH_COMPANION_SCRIPT = os.getcwd() + '/docker_updater.sh'
REPO_UPDATER_SCRIPT = os.getcwd() + '/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_stacks(self):
        return self._get('stacks')

    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 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_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")


if __name__ == '__main__':
    update_repositories()

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

    for stack in client.list_stacks():
        endpoint = stack['EndpointId']
        stack_name = stack['Name']
        if stack['Status'] == 1:
            containers = client.list_stack_containers(endpoint, stack_name)

            an_image_was_updated = False
            for c in containers:
                this_image_was_updated = client.update_image(endpoint, c['Image'])
                if this_image_was_updated:
                    print(f"{c['Image']} was updated")
                    an_image_was_updated = True

            if an_image_was_updated:
                print(f"Restarting {stack_name}")
                if endpoint == 2 and stack_name in ('portainer', 'traefik'):
                    update_local_stack(stack)
                    while not client.ping_server():
                        print("Waiting for server to be back online...")
                        sleep(1)
                else:
                    client.upgrade_stack(stack)

    os.system('docker image prune -fa')