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

# This script is a simple wrapper which prefixes each i3status line with custom
# information. It is a python reimplementation of:
# http://code.stapelberg.de/git/i3status/tree/contrib/wrapper.pl
#
# To use it, ensure your ~/.i3status.conf contains this line:
#     output_format = "i3bar"
# in the 'general' section.
# Then, in your ~/.i3/config, use:
#     status_command i3status | ~/i3status/contrib/wrapper.py
# In the 'bar' section.
#
# In its current version it will display the cpu frequency governor, but you
# are free to change it to display whatever you like, see the comment in the
# source code below.
#
# © 2012 Valentin Haenel <valentin.haenel@gmx.de>
#
# This program is free software. It comes without any warranty, to the extent
# permitted by applicable law. You can redistribute it and/or modify it under
# the terms of the Do What The Fuck You Want To Public License (WTFPL), Version
# 2, as published by Sam Hocevar. See http://sam.zoy.org/wtfpl/COPYING for more
# details.

import sys
import json
#from docker import Client

#cli = Client(base_url='unix://var/run/docker.sock')

def is_container_running(name):
    """ Check whether the given container is running. """
    return False #len(cli.containers(filters={'name':name})) == 1

def print_line(message):
    """ Non-buffered printing to stdout. """
    sys.stdout.write(message + '\n')
    sys.stdout.flush()

def read_line():
    """ Interrupted respecting reader for stdin. """
    # try reading a line, removing any extra whitespace
    try:
        line = sys.stdin.readline().strip()
        # i3status sends EOF, or an empty line
        if not line:
            sys.exit(3)
        return line
    # exit on ctrl-c
    except KeyboardInterrupt:
        sys.exit()

def get_psi_metric(type, group='some', metric='avg60'):
    input = open('/proc/pressure/%s' % type)
    for line in input.readlines():
        entries = line.strip().split(' ')
        if entries[0] == group:
            for metric_pair in entries[1:]:
                [metric_key, value] = metric_pair.split('=')
                if metric_key == metric:
                    return float(value)
    return -1

if __name__ == '__main__':
    # Skip the first line which contains the version header.
    print_line(read_line())

    # The second line contains the start of the infinite array.
    print_line(read_line())

    while True:
        line, prefix = read_line(), ''
        # ignore comma at start of lines
        if line.startswith(','):
            line, prefix = line[1:], ','

        j = json.loads(line)
        # insert information into the start of the json, but could be anywhere
        #for container in [ 'mongodb', 'elasticsearch', 'app' ]:
        #    j.insert(0, {'instance' : '%s' % container, 'name' : 'docker', 'full_text' : '%s' % container, 'color' : '%s' % ('#00FFFF' if is_container_running('collaborne_%s_1' % container) else '#7F7F7F')})
        
        for psi_type in ['cpu', 'io', 'memory']:
            value = get_psi_metric(psi_type, 'some')
            j.insert(0, {
                'instance': psi_type,
                'name': 'psi',
                'full_text': '%2.2f%% %s' % (value * 100, psi_type),
                'color': '%s' % ('#FF0000' if value > 0.8 else '#7F7F7F')
            })
        print_line(prefix+json.dumps(j))
