#!/usr/bin/env python3
#
# archey3 [version 0.2-2]
#
# Copyright 2010 Melik Manukyan <melik@archlinux.us>
# Copyright 2010 Laurie Clark-Michalek <bluepeppers@archlinux.us>
# Distributed under the terms of the GNU General Public License v3.
# See http://www.gnu.org/licenses/gpl.txt for the full license text.
#
# Simple python script to display an Archlinux logo in ASCII art 
# Along with basic system information.

# Display [Comment/Uncomment to Enable/Disable information.]
# Protocols:
#    uname:x = return output of `uname -x` (see UNAME_FLAG_MEANINGS for more info)
#    sensors:x = return output of `sensors x`
#    env:x = return value of env variable x
#    fs:x = return space of partition at x
#    mpd:stat,hostname,port = returns value of an mpd stat (options: artists|albums|songs)

DISPLAY = [
    'distro', # Display Operating System
    'uname:n', # Display Machine Hostname
    'uname:r',  # Display Kernel Version
    'uptime',  # Display System Uptime
    'sensors:coretemp-*', # Display System Core tempature
    'wm',  # Display Window Manager
    'de', # Display Desktop Environment
    'packages', # Display Number of Packages Installed
    'ram', # Display RAM Usage
    'uname:p', # Display CPU Model
    'env:shell', # Display Current Shell
    'env:editor', # Display Editor, using EDITOR env 
    'mpd:songs',
#    'fs:/boot', # Display /boot Partition Usage
#    'fs:/home', # Display /home Partition Usage
#    'fs:/MOUNT/POINT', # Display * Partition, Edit To Your Needs
    'fs:/', # Display / Partition Usage
]


# Import libraries
import subprocess, optparse, re, sys
from subprocess import Popen, PIPE
from optparse import OptionParser
from getpass import getuser
from time import ctime, sleep
from os import getenv

UNAME_FLAG_MEANINGS = {
    'a': 'System Infomation',
    's': 'Kernel Name',
    'n': 'Hostname',
    'r': 'Kernel Release',
    'v': 'Kernel Version',
    'm': 'Machine Hardware name',
    'p': 'Processor Type',
    'i': 'Hardware Platform',
}

LOGOS = {'Arch Linux': '''{c1}
{c1}               +                {results[0]}
{c1}               #                {results[1]}
{c1}              ###               {results[2]}
{c1}             #####              {results[3]}
{c1}             ######             {results[4]}
{c1}            ; #####;            {results[5]}
{c1}           +##.#####            {results[6]}
{c1}          +##########           {results[7]}
{c1}         ######{c2}#####{c1}##;         {results[8]}
{c1}        ###{c2}############{c1}+        {results[9]}
{c1}       #{c2}######   #######        {results[10]}
{c2}     .######;     ;###;`\".      {results[11]}
{c2}    .#######;     ;#####.       {results[12]}
{c2}    #########.   .########`     {results[13]}
{c2}   ######'           '######    {results[14]}
{c2}  ;####                 ####;   {results[15]}
{c2}  ##'                     '##  {results[16]} 
{c2} #'                         `#  {results[17]}
\x1b[0m'''
}

FUNC_MAPPINGS = {
    'distro': 'distroCheck',
    'uname': 'unameDisplay',
    'uptime': 'uptimeDisplay',
    'sensors': 'sensorDisplay',
    'wm': 'wmDisplay',
    'de': 'deDisplay',
    'packages': 'packageDisplay',
    'ram': 'ramDisplay',
    'env': 'envDisplay',
    'fs': 'fsDisplay',
    'mpd': 'mpdDisplay',
}


DE_DICT = {'gnome-session': 'GNOME',
    'ksmserver': 'KDE',
    'xfce4-session': 'Xfce 4.6',
    'lxsession': 'LXDE',
    '': 'None',
}
        
WM_DICT = {
    'awesome': 'Awesome',
    'beryl': 'Beryl',
    'blackbox': 'Blackbox',
    'compiz': 'Compiz',
    'dwm': 'DWM',
    'enlightenment': 'Enlightenment',
    'fluxbox': 'Fluxbox',
    'fvwm': 'FVWM',
    'i3': 'i3',
    'icewm': 'IceWM',
    'kwin': 'KWin',
    'metacity': 'Metacity',
    'musca': 'Musca',
    'openbox': 'Openbox',
    'pekwm': 'PekWM',
    'ratpoison': 'ratpoison',
    'scrotwm': 'ScrotWM',
    'wmaker': 'Window Maker',
    'wmfs': 'Wmfs',
    'wmii': 'wmii',
    'xfwm4': 'Xfwm',
    'xmonad': 'xmonad',
    '': 'None',
}

COLORS = {
    'black': '0',
    'red': '1',
    'green': '2',
    'yellow': '3',
    'blue': '4',
    'magenta': '5',
    'cyan': '6',
    'white': '7'
}

class display(object):
    __slots__ = ['command_line', 'arg1', 'arg2', 'arg3', 'stdindata', 'proccess', '_parent']
    
    command_line = ''
    arg1 = ''
    arg2 = ''
    arg3 = ''
    stdindata = ''
    
    def __init__(self, parent=None):
        self._parent = parent
    
    def run_command(self):
        if self.command_line:
            if '{arg3}' in self.command_line:
                cmd = self.command_line.format(arg1=self.arg1, arg2=self.arg2, arg3=self.arg3)
            elif '{arg2}' in self.command_line:
                cmd = self.command_line.format(arg1=self.arg1, arg2=self.arg2)
            elif '{arg1}' in self.command_line:
                cmd = self.command_line.format(arg1=self.arg1)
            else:
                cmd = self.command_line
            try:
                self.proccess = Popen(cmd.split(), stdin=PIPE, stdout=PIPE, stderr=PIPE)
            except Exception as e:
                pass
    
    def render(self):
        (stdoutdata, stderrdata) = self.proccess.communicate(self.stdindata or None)
        return self.format_output(stdoutdata.decode())
        
    def color_me(self, output, number=None, low=30, low_color='green', medium=60, medium_color='yellow', high_color='red'):
        if number is None and output.isdigit():
            number = int(output)
        elif number is None:
            return output
        
        if number <= low:
            color = low_color
        elif low < number <= medium:
            color = medium_color
        elif medium < number:
            color = high_color
        
        return '{0}{1}{2}'.format(self._parent.color(color), output, self._parent.color('clear'))

class fsDisplay(display):
    command_line = "df -TPh {arg1}"
    
    conversions = {
        'M': 1 * 10 ** 6,
        'K': 1 * 10 ** 3,
        'G': 1 * 10 ** 9,
        'T': 1 * 10 ** 12,
        'P': 1 * 10 ** 15,
    }
    
    def __init__(self, path='/', parent=None):
        self.arg1 = path
        self._parent = parent
    
    def format_output(self, instring):
        values = [line for line in instring.split('\n') if line][1].split()
        used = values[3]
        total = values[2]
        fstype = values[1]
        
        mount = '/root' if self.arg1 == '/' else self.arg1 
        title = mount.split('/')[-1].title() + " FS"
        try:
            used_ = int(used[:-1]) * self.conversions[used[-1]]
            total_ = int(total[:-1]) * self.conversions[total[-1]]
            persentage = used_ / total_ * 100
        except:
            pass
        else:
            used = self.color_me(used, persentage)
        part = '{used} / {total} ({fstype})'.format(used=used, total=total, fstype=fstype)
        return title, part

class ramDisplay(display):
    command_line = "free -m"
    
    def format_output(self, instring):
        ram = ''.join(line for line in str(instring).split('\n') if line.startswith('Mem:')).split()
        used = int(ram[2]) - int(ram[5]) - int(ram[6])
        total = int(ram[1])
        title = 'RAM'
        try:
            persentage = (used / total * 100)
        except:
            used += ' MB'
        else:
            used = self.color_me(number=persentage, output=str(used) + ' MB')
        part = '{used} / {total} MB'.format(used=used, total=total)
        return title, part
        
class sensorDisplay(display):
    command_line = "sensors {arg1}"
    
    def __init__(self, sensors="coretemp-*", parent=None):
        self.arg1 = sensors
        self._parent = parent
    
    def format_output(self, instring):
        tempinfo = instring.split('\n')[2::4]
        
        out = []
        for line in tempinfo:
            info = [re.sub("\s\s+", "", line) for line in line.split('  ') if line]
            value = info[1]
            intvalue = int(value[:3])
            if intvalue > 45:
                temp = self._parent.color("red") + info[1] + self._parent.color("clear")
            elif intvalue in range(30, 45):
                temp = self._parent.color("magenta") + info[1] + self._parent.color("clear")
            else:
                temp = self._parent.color("green") + info[1] + self._parent.color("clear")
            out.append((info[0], temp))
        return out

class envDisplay(display):
    def __init__(self, env='SHELL', parent=None):
        self.arg1 = env
        self._parent = parent
    
    def render(self):
        argvalue = getenv(self.arg1.upper())
        return ('$' + self.arg1.upper(), argvalue)

class unameDisplay(display):
    command_line = "uname {arg1}"
    
    def __init__(self, flag=False, parent=None):
        self.arg1 = '-' + flag if flag else ''
        self._parent = parent
    
    def format_output(self, instring):
        return (UNAME_FLAG_MEANINGS[self.arg1[1]], instring)

class uptimeDisplay(display):
    def render(self):
        with open("/proc/uptime") as upfile:
            raw = upfile.read()
            fuptime = int(raw.split('.')[0])
            
            day = int(fuptime / 86400)
            fuptime = fuptime % 86400
            hour = int(fuptime / 3600)
            fuptime = fuptime % 3600
            minute = int(fuptime / 60)
            uptime = '{daystring}{hours}:{mins:02d}'.format(
                daystring='{days} day{s}, '.format(days=day, s='s' if day > 1 else '') if day else '',
                hours = hour, mins = minute
            )
            return "Uptime", uptime

class packageDisplay(display):
    command_line = "pacman -Q"
    
    def format_output(self, instring):
        return "Packages", len(instring.split('\n'))

class distroCheck(display):
    def render(self):
        try:
            rc = open("/etc/pacman.conf")
        except IOError:
            distro = render(unameDisplay("o"))[1]
        else:
            distro = "Arch Linux"
        distro = ' '.join([distro, render(unameDisplay('m'))[1]])
        return "OS", distro

class proccessCheck(display):
    __slots__ = display.__slots__ + ['_processes']
    command_line = "ps -u {arg1}"
    
    render = lambda self: self
    
    def __init__(self, parent=None):
        self.arg1 = getuser()
        self._parent = parent
    
    def run_command(self):
        super().run_command()
        out = str(self.proccess.communicate()[0])
        self._processes = set([line.split()[3] for line in out.split('\\n') if len(line.split()) == 4])
    
    def __call__(self, proc):
        if proc in self._processes:
            return True
        return False
        
class wmDisplay(display):
    def render(self):
        wm = ''
        for key in WM_DICT.keys():
            if processes(key):
                wm = key
                break
        return "WM", WM_DICT[wm]

class deDisplay(display):
    def render(self):
        de = ''
        for key in DE_DICT.keys():
            if processes(key):
                de = key
                break
        return "DE", DE_DICT[de]

class mpdDisplay(display):
    """Displays certain stat about MPD database. If mpd not installed, output nothing."""
    command_line = "mpc stats --host {arg1} --port {arg2}"
    
    def __init__(self, stat='songs', hostname='localhost', port='6600', parent=None):
        self.stat = stat
        self.hostname = hostname
        self.port = port
        self.arg1 = hostname
        self.arg2 = port
        self._parent = parent
        try:
            global mpd
            import mpd
            self.method = 'module'
        except ImportError:
            self.method = 'shell'
    
    def run_command(self):
        if self.method == 'module':
            self.proccess = mpd.MPDClient()
            try:
                self.proccess.connect(self.hostname, self.port)
            except Exception:
                self.method == 'shell'
                super().run_command()
        else:
            super().run_command()
    
    def render(self):
        if self.method == 'module':
            try:
                return '{statname} in MPD database'.format(statname=self.stat.title()), self.proccess.stats()[self.stat]
            except Exception:
                return False
        else:
            if hasattr(self, "proccess"):
                return super().render()
            else:
                return None
    
    def format_output(self, instring):
        lines = instring.split('\n')
        stats = {}
        try:
            stats['artists'] = lines[0].split(':')[1].strip()
            stats['albums'] = lines[1].split(':')[1].strip()
            stats['songs'] = lines[2].split(':')[1].strip()
        #if people don't have mpc installed (= no mpd installed) then return None)
        except ValueError:
            return False
        
        return '{statname} in MPD database'.format(statname=self.stat.title()), stats[self.stat]

#------------ Functions -----------

def screenshot():
        print('Screenshotting in')
        for x in sorted(range(1,6), reverse=True):
                print('%s' % x, end='')
                sys.stdout.flush()
                sleep(1.0/3)
                for x in range(3):
                    print('.', end='')
                    sys.stdout.flush()
                    sleep(1.0/3)
                    
        print('Say Cheese!')
        sys.stdout.flush()
        subprocess.check_call(['import', '-window', 'root', ctime().replace(' ','_')+'.jpg'])

def render(instance):
    instance.run_command()
    return instance.render()

#------------ Display object ---------

class Archey(object):
    def __init__(self, display=None, colorscheme='blue'):
        self.display = display or []
        for key in COLORS.keys():
            if key == colorscheme:
                self.colorcode = COLORS[key]
        
        self.distro_name = ' '.join(render(distroCheck())[1].split()[:-1])
    
    def render(self):
        results = self.prepare_results()
        
        print(LOGOS[self.distro_name].format(c1=self.color(1),
                                             c2=self.color(2),
                                             results = results
                                            ))
    
    def prepare_results(self):
        outputs = []
        # Run functions found in 'display' array.
        for func in self.display:
            if len(func.split(':')) > 1:
                cls = eval(FUNC_MAPPINGS[func.split(':')[0]])
                args = tuple(func.split(':')[1].split(','))
            else:
                cls = eval(FUNC_MAPPINGS[func])
                args = ()
            
            out = render(cls(*args, parent=self))
            if isinstance(out, list):
                for o in out:
                    outputs.append(o)
            elif out:
                outputs.append(out)
            else:
                continue
        
        for i in range(18):
            if i < len(outputs):
                outputs[i] = self.format_item(outputs[i])
            else:
                outputs.append('')
        
        return outputs

    def format_item(self, item):
        title = item[0].rstrip(':')
        data = str(item[1]).rstrip()
        
        #if we're dealing with a fraction
        if len(data.split('/')) == 2:
            numerator = data.split('/')[0]
            numerator = self.color(1, bold=True) + numerator + self.color('clear')
            denominator = data.split('/')[1]
            data = '/'.join((numerator, denominator))
        
        return "{color}{title}:{clear} {data}".format(
                color=self.color(1),
                title=title,
                data=data,
                clear=self.color("clear")
        )
        
    def color(self, code, bold=False):
        if code == 2:
            bold = True
        first_bitty_bit = '\x1b[{0};'.format(int(not bold))
        if code in range(3):
            second_bitty_bit = '3{0}m'.format(self.colorcode)
        elif code == "clear":
            return '\x1b[0m'
        else:
            second_bitty_bit = '3{0}m'.format(COLORS[code])
       
        return first_bitty_bit + second_bitty_bit


def main():
    global processes
    processes = render(proccessCheck())
    
    parser = OptionParser(usage='%prog [-c COLOR] [-s, --screenshot] [-d, --display]', 
            description="""Archey is a utility to take """,
        version="%prog 0.2-1")
    parser.add_option('-c',
            action='store', default='blue', type='choice', dest='color',
            choices=('black',
                     'red',
                     'green',
                     'yellow',
                     'blue',
                     'magenta',
                     'cyan',
                     'white'),
            help='choose a color: black, red, green, yellow, blue, magenta, cyan, white [Default: blue]')
    parser.add_option('-s', '--screenshot',
            action='store_true', dest='screenshot', help='Take a screenshot')
    parser.add_option('-d', '--display',
            action='store', dest='display', help='Define info to be displayed by archey. DISPLAY must be in the \
format "archey_display_module:arg1,arg2,arg3.second_display_module". Availible display modules are uname, fs, env, \
uptime, sensors, ram, de and wm. See the DISPLAY list in the archey source file for more info.')
    (options, args) = parser.parse_args()
    
    if options.display:
        global DISPLAY
        DISPLAY = options.display.split('.')
    
    archey = Archey(display=DISPLAY, colorscheme=options.color)
    archey.render()
    
    if options.screenshot:
        screenshot()

if __name__ == "__main__":
    main()