#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pygame
import sys
import time
from subprocess import check_output

threshold = 1000


def render_stat(stat, font, screen):
    for i, (machine, status) in enumerate(stat.items()):
        textsurface = None
        if status == 'busy':
            text = u"{} is busy".format(machine)
            color = (255, 0, 0)
        elif status == 'free':
            text = u"{} is free".format(machine)
            color = (0, 255, 0)
        else:
            text = u"{} is unavailable".format(machine)
            color = (128, 0, 128)
        textsurface = font.render(text, False, color)
        screen.blit(textsurface, (20, 200 + 250 * i))


def check_local():
    try:
        free_mem = int(
            check_output(
                [
                    'nvidia-smi', '--query-gpu=memory.used',
                    '--format=csv,nounits,noheader'
                ],
                timeout=10))
        if free_mem > threshold:
            return 'busy'
        else:
            return 'free'
    except:
        return 'err'


def check_remote(address, port=22):
    try:
        free_mem = int(
            check_output(
                [
                    'ssh', address, '-p',
                    str(port), '-t', 'nvidia-smi', '--query-gpu=memory.used',
                    '--format=csv,nounits,noheader'
                ],
                timeout=10))
        if free_mem > threshold:
            return 'busy'
        else:
            return 'free'
    except:
        return 'err'


pygame.init()
size = (pygame.display.Info().current_w, pygame.display.Info().current_h)
screen = pygame.display.set_mode(size, pygame.FULLSCREEN)
pygame.font.init()
myfont = pygame.font.SysFont('DejaVuSans', 100, bold=True)
stat = {}

while True:
    stat[u'Name1'] = check_local()
    stat[u'Name2'] = check_remote('user@address')

    screen.fill((0, 0, 0))
    render_stat(stat, myfont, screen)
    pygame.display.flip()
    time.sleep(10)