#! /usr/bin/env python3
# encoding: utf-8

import os
import socket
import sys

SYM_GIT = "\ue0a0"
SYM_LOCK = "\ue0a2"

def style_color(fg, bg):
	return '\033[38;5;%dm\033[48;5;%dm' %(fg, bg)

def style_bold():
	return '\033[1m'

def style_reset():
	return '\033[0m'

class LoginPart:
	user_fg = 34
	root_fg = 196
	remote_fg = 202
	local_fg = 7
	screen_fg = 27
	bg = 0
	
	def __init__(self):
		self.user = os.getlogin()
		self.remote = bool(os.getenv('SSH_CLIENT'))
		self.sign = ' ⇄ ' if self.remote else '@'
		self.host = socket.gethostname()
		self.screen = os.getenv('WINDOW')
	
	@property
	def length(self):
		return len(self.user) + len(self.host) + len(self.screen) + len(self.sign) + 1
	
	def __str__(self):
		output = ''
		
		if os.geteuid() == 0:
			output += style_color(self.root_fg, self.bg)
		else:
			output += style_color(self.user_fg, self.bg)
		
		output += self.user
		
		if self.remote:
			output += style_color(self.remote_fg, self.bg)
		else:
			output += style_color(self.local_fg, self.bg)
		
		output += self.sign
		output += self.host
		
		if self.screen:
			output += style_reset()
			output += ':'
			output += style_color(self.screen_fg, self.bg)
			output += self.screen
		
		output += style_reset()
		
		return output

def list_homes():
	with open('/etc/passwd') as f:
		for line in f:
			tokens = line.split(':')
			
			if tokens[5] != '/dev/null':
				yield (tokens[5], tokens[0])

class Dir:
	writable_fg = 34
	fg = 196
	bg = 0
	
	def __init__(self, path, name=None, slash='/'):
		self.path = path
		self.slash = slash
		self.lock = not os.access(path, os.W_OK)
		self.truncated = False
		
		if name is None:
			name = os.path.basename(path)
		
		self.raw_str = name + ('' if self.lock else '')
	
	@property
	def length(self):
		return len(self.raw_str) + len(self.slash)
	
	def truncate(self):
		self.truncated = True
		
		if self.raw_str:
			self.raw_str = self.raw_str[0]
	
	def __str__(self):
		output = self.slash
		
		if self.truncated:
			output += style_bold()
		
		if self.raw_str:
			output += style_color(self.fg if self.lock else self.writable_fg, self.bg)
			output += self.raw_str
		
		output += style_reset()
		
		return output

class PathPart:
	def __init__(self, term_width):
		self.term_width = term_width
		self.segments = []
		
		homes = dict(list_homes())
		dirs = os.getcwd().split('/')[1:]
		path = ''
		
		for dir in dirs:
			path += '/' + dir
			
			if path in homes:
				self.segments = []
				
				if homes[path] == os.getlogin():
					self.segments.append(Dir(path, '', '~'))
				else:
					self.segments.append(Dir(path, homes[path], '~'))
			
			else:
				self.segments.append(Dir(path))
	
	@property
	def length(self):
		return sum(s.length for s in self.segments)
	
	def truncate(self, length):
		for s in self.segments:
			if self.length > length:
				s.truncate()
			else:
				break
	
	def truncate_fit(self, line):
		space = self.term_width - sum(part.length for part in line if part is not self)
		self.truncate(space)
	
	def __str__(self):
		return ''.join(str(s) for s in self.segments)

class GitPart:
	def __init__(self):
		...
	
	@property
	def length(self):
		return len(str(self))
	
	def __str__(self):
		return 'GIT'

class Padding:
	def __init__(self, term_width):
		self.term_width = term_width
		self.length = 0
	
	def expand_fit(self, line):
		self.length = self.term_width - sum(part.length for part in line if part is not self) - len(line) + 1
	
	def __str__(self):
		return ' ' * self.length

class StatsPart:
	clock = 27
	space_good = 34
	space_warn = 208
	space_bad = 196
	
	def __init__(self, red_thresh, yellow_thresh, warning_only=True):
		self.red = red_thresh
		self.yellow = yellow_thresh
	
	@property
	def length(self):
		return len(str(self))
	
	def __str__(self):
		return '[ 13:34 | 0.53 | r4.20M'

if __name__ == '__main__':
	red_thresh = float(sys.argv[1])
	yellow_thresh = float(sys.argv[2])
	term_width = int(sys.argv[3])
	exit_status = int(sys.argv[4])
	
	lp = LoginPart()
	pp = PathPart(term_width)
	gp = GitPart()
	pad = Padding(term_width)
	sp = StatsPart(red_thresh, yellow_thresh)
	
	line = [lp, pp, gp, pad, sp]
	pp.truncate_fit(line)
	pad.expand_fit(line)
	
	line_str = ' '.join(str(part) for part in line)
	sys.stderr.write(line_str)
	sys.stderr.flush()