#!/usr/bin/env python3
#
# archey3 [version 0.2-1]
#
# 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.

# 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

# 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

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 
#	'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
        ]

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]}
{c2} '''
}

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


DE_DICT = {'gnome-session': 'GNOME',
        'ksmserver': 'KDE',
        'xfconfd': 'Xfce 4.6',
        'lxsession': 'LXDE',
        '': 'None',
}
        
WM_DICT = {'awesome': 'Awesome',
        'beryl': 'Beryl',
        'blackbox': 'Blackbox',
        'compiz': 'Compiz',
        'dwm': 'DWM',
        'enlightenment': 'Enlightenment',
        'fluxbox': 'Fluxbox',
        'fvwm': 'FVWM',
        'icewm': 'IceWM',
        'kwin': 'KWin',
        'metacity': 'Metacity',
        'musca': 'Musca',
        'openbox': 'Openbox',
        'pekwm': 'PekWM',
        'wmaker': 'Window Maker',
        'wmii': 'wmii',
        'xfwm4': 'Xfwm',
        'xmonad': 'xmonad',
        'scrotwm': 'ScrotWM',
        '': '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
            self.proccess = Popen(cmd.split(), stdin=PIPE, stdout=PIPE, stderr=PIPE)
    
    def render(self):
        (stdoutdata, stderrdata) = self.proccess.communicate(self.stdindata or None)
        return self.format_output(stdoutdata.decode())
        

class fsDisplay(display):
    command_line = "df -TPh {arg1}"
    
    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"
        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'
        part = '{used} MB / {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}'.format(
                daystring='{days} day{s}, '.format(days=day, s=int(day) > 1) 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]

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

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

#create process list

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)
            else:
                outputs.append(out)
        
        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):
        return "{color}{title}:{clear} {data}".format(
                color=self.color(1),
                title=item[0].rstrip(':'),
                data=str(item[1]).rstrip(),
                clear=self.color("clear")
        )
        
    def color(self, code):
        if code == 1:
            return '\x1b[1;3%sm' % self.colorcode
        elif code == 2:
            return '\x1b[0;3%sm' % self.colorcode
        if code == "clear":
            return '\x1b[0m'
        return '\x1b[0;3%sm' % COLORS[code]


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()