#!/usr/bin/env python

import sys
import time
import curses
import argparse

from foundation.utils import Workers
from foundation.workers import list_workers

parser = argparse.ArgumentParser(description="Start an HCI worker.")
parser.add_argument('-a', '--advertise_addr', default=None, help="Advertise address.")
args = parser.parse_args()


workers = Workers(swarm_advertise_addr=args.advertise_addr)

foundation_services = set([
    'kafka-service',
    'zookeeper-service',
    'kafka-logs-service',
    'zookeeper-logs-service',
    'timescaledb-service',
    'jupyterlab-service',
])

system_workers = set([f'{worker}-service-worker' for worker in list_workers()])


########################################################################
class Stats:
    """"""
    COL = [0, 40 + 5, 53 + 5, 62 + 10, 84 + 10, 101 + 10, 139 + 10]
    COL_NAMES = ["SERVICE", "RUNNING", "NODE", "SERVICE ID", "CONTAINER ID", "IMAGE", "URL"]

    # ----------------------------------------------------------------------
    def __init__(self, stdscr):
        """Constructor"""
        curses.resizeterm(100, 500)
        curses.curs_set(0)
        self.stdscr = stdscr

        curses.init_pair(1, curses.COLOR_RED, curses.COLOR_BLACK)
        curses.init_pair(2, curses.COLOR_GREEN, curses.COLOR_BLACK)

        while True:
            self.stdscr.clear()

            self.stdscr.addstr(0, 0, '-' * 80)
            start = self.update_services_stats(foundation_services, name='FOUNDATION SERVICES', start=2)
            start = self.update_services_stats(system_workers, name='SYSTEM WORKERS', start=start)
            start = self.update_services_stats(set(workers.swarm.services) - foundation_services - system_workers, name='USER WORKERS', start=start)

            self.stdscr.addstr(start, 0, '-' * 80)
            start = self.update_volumes_stats(name='VOLUMES', start=start + 2)

            self.stdscr.addstr(start, 0, '-' * 80)

            self.stdscr.addstr(start + 2, 0, workers.swarm.get_join_command())

            self.stdscr.refresh()
            time.sleep(1)

    # ----------------------------------------------------------------------
    def write_row(self, items, row=0):
        """"""
        for item, col in zip(items, self.COL):
            if ('----' in items[:-1] or 'N/A' in items[:-1]) and (col == 0):
                color = curses.color_pair(1)
            elif (items[1] == 'True') and (col == 0):
                color = curses.color_pair(2)
            else:
                color = curses.color_pair(0)
            self.stdscr.addstr(row, col, item, color)

    # ----------------------------------------------------------------------
    def service_get(self, service, attr='id', default='----'):
        """"""
        service = workers.swarm.client.services.list(filters={'name': service, })
        if service:
            try:
                match attr:
                    case 'id':
                        return service[0].id[:12]
                    case 'container_id':
                        return service[0].tasks()[0]['Status']['ContainerStatus']['ContainerID'][:12]
                    case 'image':
                        return service[0].tasks()[0]['Spec']['ContainerSpec']['Image']
                    case 'ip':
                        return list(filter(lambda s: s.startswith('PORT='), service[0].tasks()[0]['Spec']['ContainerSpec']['Env']))[0].replace('PORT=', '')
                    case 'service_name':
                        try:
                            return "@" + list(filter(lambda s: s.startswith('SERVICE_NAME='), service[0].tasks()[0]['Spec']['ContainerSpec']['Env']))[0].replace('SERVICE_NAME=', '')
                        except:
                            return ''
                    case 'endpoint':
                        try:
                            return list(filter(lambda s: s.startswith('ENDPOINT='), service[0].tasks()[0]['Spec']['ContainerSpec']['Env']))[0].replace('ENDPOINT=', '')
                        except:
                            return ''
                    case 'node':

                        node_id = service[0].tasks()[0]['NodeID']

                        node_info = workers.swarm.client.nodes.get(node_id)

                        return node_info.attrs['Description']['Hostname']

            except:
                return "N/A"
        else:
            return default

    # ----------------------------------------------------------------------
    def update_services_stats(self, services, name, start=0):
        """"""
        if not services:
            return start
        self.write_row([name] + self.COL_NAMES[1:], row=start)
        for row, service in enumerate(services, start=start + 1):

            url = '----'
            if self.service_get(service, attr='ip') not in ['----', 'N/A']:
                url = f"http://127.0.0.1:{self.service_get(service, attr='ip')}{self.service_get(service, attr='endpoint')}"

            self.write_row([
                f"{service}{self.service_get(service, attr='service_name',  default='')}",
                f"{service in workers.swarm.services}",
                f"{self.service_get(service, attr='node')}",
                f"{self.service_get(service, attr='id')}",
                f"{self.service_get(service, attr='container_id')}",
                f"{self.service_get(service, attr='image')}",
                url,
            ],
                row=row)
        return row + 2

    # ----------------------------------------------------------------------
    def update_volumes_stats(self, name, start=0):
        """"""
        self.stdscr.addstr(start, 0, name)
        networks = sorted([v for v in workers.swarm.client.volumes.list() if
                           v.name.endswith('-volume')], key=lambda s: s.name)
        for row, volume in enumerate(networks, start=start + 1):
            self.stdscr.addstr(row, 0, volume.name)
        return row + 2


try:
    curses.wrapper(Stats)
except:
    sys.exit()

