#!/usr/bin/env python3

## -=Notebook=-
#
# - Exception hierarchy: https://docs.python.org/3/library/exceptions.html#exception-hierarchy
# Non-default Prerequisites:
# - Font: Cascadia Code
# - APT packages: python3-dev
# - PIP packages: psutil, pillow

import socket, io, psutil, datetime
from os.path import exists as file_exists
from PIL import Image, ImageDraw, ImageFont

# Socket constants
KUBESOCK="/var/run/pibox/framebuffer.sock"
HEADER1=b"POST /image HTTP/1.1\x0d\x0aHost: localhost\x0d\x0aUser-Agent: python-socket\x0d\x0aAccept: */*\x0d\x0a"
HEADER2=b"\x0d\x0aContent-Type: application/x-www-form-urlencoded\x0d\x0a\x0d\x0a"

# Drawing constants
SCREENWIDTH=240
SCREENHEIGHT=240
CASCADIAFONTPATH="/usr/share/fonts/truetype/CascadiaCode/static/CascadiaMono-Regular.ttf"
FONTSIZE={"CHONKY": 35,"THICC": 25, "BIG":15, "MED":12,"SMOL":10}
#DISPLAYWIDTHCHARS={"CHONKY": 11, "THICC": 16, "BIG":26, "MED":34, "SMOL":40}
FONT={"CHONKY":ImageFont.truetype(CASCADIAFONTPATH,FONTSIZE["CHONKY"]),
      "THICC":ImageFont.truetype(CASCADIAFONTPATH,FONTSIZE["THICC"]),
      "BIG":ImageFont.truetype(CASCADIAFONTPATH,FONTSIZE["BIG"]),
      "MED":ImageFont.truetype(CASCADIAFONTPATH,FONTSIZE["MED"]),
      "SMOL":ImageFont.truetype(CASCADIAFONTPATH,FONTSIZE["SMOL"])}
BLACK,GRAY,WHITE=(0,0,0,255),(127,127,127,255),(255,255,255,255)
RED,ORANGE,YELLOW=(200,0,0,255),(200,125,0,255),(200,200,0,255)
DARK_GREEN, DARK_YELLOW, DARK_RED=(0,80,0,255),(150,150,0,255),(100,0,0,255)
GREEN,BLUE,PURPLE=(0,230,0,255),(80,190,255,255),(180,100,255,255)
AQUA,PINK=(55, 255, 255, 255),(255,150,230,255)


# Metrics constants
NET_ADAPTER="eth0"
CPULOADLABELS=["1min", "5min", "15min"]
# only added to as-anticipated
SUBNET_CIDR_MAP={"255.0.0.0": "8", "255.255.0.0": "16", "255.255.255.0": "24", "255.255.255.128":"25"}

mystats=Image.new("RGBA", (240, 240), BLACK)
drawer=ImageDraw.Draw(mystats)

def precheck():
 if not file_exists(KUBESOCK):
  raise FileNotFoundError(f"Cannot find socket {KUBESOCK} - ensure socket server is running")
 if not file_exists(CASCADIAFONTPATH):
  raise FileNotFoundError(f"Cannot find font file {CASCADIAFONTPATH} - please ensure it is installed on your system")
 return

def sendimage(imgdata):
 if type(imgdata) is Image.Image:
  # https://stackoverflow.com/a/33117447
  newbytes=io.BytesIO()
  imgdata.save(newbytes, format="PNG")
  imgbytes=newbytes.getvalue()
 elif type(imgdata) is io.BytesIO:
  imgbytes=imgdata.getvalue()
 imglength=f"{len(imgbytes)}"
 request_data=HEADER1+b"Content-Length: "+bytes(imglength,"UTF-8")+HEADER2+imgbytes
 with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
  client.connect(KUBESOCK)
  client.send(request_data)
  #print(request_data)
  print(client.recv(4096))
  client.close()
 return

def writewords(text, line=5, spacing=3, indent=0, right=SCREENWIDTH, font="BIG", color=WHITE, alignment="LEFT", outline=None):
 alignment=alignment.upper()
 font=font.upper()
 pixelwidth=round(drawer.textlength(text, font=FONT[font]),0)
 #print(f"text \"{text}\" has width {pixelwidth} pixels.")
 strokewidth=0 if outline == None else 1
 if alignment == "LEFT":
  drawer.text((indent, line), text, font=FONT[font], fill=color, stroke_fill=outline, stroke_width=strokewidth)
 elif alignment == "RIGHT":
  drawer.text((right-pixelwidth-indent, line), text, font=FONT[font], fill=color, stroke_fill=outline, stroke_width=strokewidth)
 elif alignment == "CENTER":
  drawer.text(((((right-indent)/2)-(pixelwidth/2))+indent, line), text, font=FONT[font], fill=color, stroke_fill=outline, stroke_width=strokewidth)
 else:
  raise ValueError(f"Unsupported alignment provided to writewords function: {alignment}")
 return line+FONTSIZE[font]+spacing

def drawbar(top=50, indent=20, right=220, height=30, color=WHITE, corners=(True, True, True, True), radius=10, percent=.5):
 # Minimum 11% for easy draw (radius*2+1 pixels)
 # Maximum 95% for easy draw
 radius = radius if radius*2+2 < height else height/2-2
 outlinecorners=corners
 fillcorners=corners
 #print(f"Top left: ({indent}, {top}) / Bottom right: ({right}, {top+height})")
 #             0          1           2             3
 # corners=(top_left, top_right, bottom_right, bottom_left)
 if percent < .9: fillcorners=(corners[0], False, False, corners[3])
 if percent >= .15:
  drawer.rounded_rectangle(((indent, top), (((right-indent)*percent)+indent, top+height)), radius=radius, outline=None, fill=color, corners=fillcorners)
 else:
  # Due to how the rounded_rectangle function draws the shape we have to draw at least 11% (20% is just a nice number)
  # and then "subtract" the leftover percents from the bar to show only what we want
  # (also the top 5% isn't neatly rounded but that's a problem for a later time)
  drawer.rounded_rectangle(((indent, top), (((right-indent)*.2)+indent, top+height)), radius=radius, outline=None, fill=color, corners=fillcorners)
  drawer.rectangle(((((right-indent)*percent)+indent, top), (((right-indent)*.2)+indent, top+height)), fill=BLACK, outline=None)
 # Then we draw the outline afterwards so it goes over our blacked-out region (if present)
 drawer.rounded_rectangle(((indent, top), (right, top+height)), radius=radius, outline=color, fill=None, corners=outlinecorners)
 return top+height+1

def substring_indicies(string, character):
 # https://stackoverflow.com/a/13009866
 return [i for i, letter in enumerate(string) if letter == character]



def textstats():
 line=FONTSIZE['BIG']
 column=(5,120)
 # CPU PERCENTS
 cpu_use=psutil.cpu_percent(interval=1,percpu=True)
 core=0
 while core < len(cpu_use):
  drawer.text((column[core%2],line), f"CPU {core}: {cpu_use[core]}%", font=FONT['BIG'], fill=WHITE)
  line+= 0 if core%2==0 else FONTSIZE['BIG']+5
  core+=1
 # CPU LOAD
 cpu_loadsss=psutil.getloadavg()
 load=0
 line+=10
 while load < len(cpu_loadsss):
  drawer.text((column[0], line), f"{CPULOADLABELS[load]} load: {round(cpu_loadsss[load],4)}", font=FONT['BIG'], fill=WHITE)
  line+=FONTSIZE['BIG']+5
  load+=1
 line+=10
 # MEMORY USAGE
 memories=psutil.virtual_memory()
 memtotal=round(memories.total/1024/1024/1024, 2)
 memavailable=round(memories.available/1024/1024/1024, 2)
 drawer.text((column[0], line), f"Total: {memtotal}G", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+5
 drawer.text((column[0], line), f"Avail: {memavailable}G", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+15
 temp=round(psutil.sensors_temperatures()['cpu_thermal'][0].current, 1)
 drawer.text((column[0], line),f"Temp: {temp}°C", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+15
 # Hello, world!
 drawer.text((column[0],line), "Hello, world!", font=FONT['BIG'], fill=WHITE)

def texttest():
 line=5
 spacing=3
 smolspacing=3
 indent=0
 linebreak=15
 drawer.text((indent, line), "{:^26}".format("--Letter spacing test--"), font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+linebreak
 drawer.text((indent, line), f"font: {FONT['BIG'].getname()[0]} {FONTSIZE['BIG']}px", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+spacing
 drawer.text((indent, line), "123456789012345678901234567890", font=FONT['MED'], fill=RED)
 line+=FONTSIZE['BIG']+spacing
 drawer.text((indent, line), "abcdefghijklmnopqrstuvwxyz", font=FONT['BIG'], fill=ORANGE)
 line+=FONTSIZE['BIG']+spacing
 drawer.text((indent, line), "ABCDEFGHIJKLMNOPQRSTUVWXYZ", font=FONT['BIG'], fill=YELLOW)
 line+=FONTSIZE['BIG']+spacing
 drawer.text((indent, line), f"line spacing: {spacing}", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+linebreak

 drawer.text((indent, line), f"smol font: {FONT['SMOL'].getname()[0]} {FONTSIZE['SMOL']}px", font=FONT["SMOL"], fill=WHITE)
 line+=FONTSIZE["SMOL"]+smolspacing
 drawer.text((indent, line), "1234567890123456789012345678901234567890", font=FONT["SMOL"], fill=GREEN)
 line+=FONTSIZE["SMOL"]+smolspacing
 drawer.text((indent, line), "abcdefghijklmnopqrstuvwxyz", font=FONT["SMOL"], fill=BLUE)
 line+=FONTSIZE["SMOL"]+smolspacing
 drawer.text((indent, line), "ABCDEFGHIJKLMNOPQRSTUVWXYZ", font=FONT["SMOL"], fill=PURPLE)
 line+=FONTSIZE["SMOL"]+smolspacing
 drawer.text((indent, line), f"smol line spacing: {smolspacing}", font=FONT["SMOL"], fill=WHITE)
 line+=FONTSIZE["SMOL"]+linebreak

 drawer.text((indent, line), f"left indent: {indent}", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+spacing

def testguistats():
 line=5
 footerline=240-FONTSIZE['BIG']-5
 barheight=10
 barwidth=220
 column=(5,120)
 # CPU PERCENTS
 line = writewords("CPU Usage", indent=5, line=line, font="big")
 cpu_use=psutil.cpu_percent(interval=1,percpu=True)
 core=0
 while core < len(cpu_use):
  rect_start=(column[0], line)
  rect_end=(5+round(barwidth/100*cpu_use[core],0), line+barheight)
  rect_outline=(5+barwidth,line+barheight)
  drawer.rectangle([rect_start, rect_outline],fill=None,outline=GREEN)
  drawer.rectangle([rect_start, rect_end],fill=GREEN,outline=None)
  line+=barheight
  core+=1
 line+=FONTSIZE['BIG']
 # CPU LOAD
 # cpu_loadsss=psutil.getloadavg()
 # load=0
 # line+=10
 # while load < len(cpu_loadsss):
 #  drawer.text((column[0], line), f"{CPULOADLABELS[load]} load: {round(cpu_loadsss[load],4)}", font=FONT['BIG'], fill=WHITE)
 #  line+=FONTSIZE['BIG']+5
 #  load+=1
 # line+=10
 # MEMORY USAGE
 line = writewords("Memory Usage", line=line, indent=5)
 memories=psutil.virtual_memory()
 memtotal=round(memories.total/1024/1024/1024, 2)
 memavailable=round(memories.available/1024/1024/1024, 2)
 rect_start=(column[0], line)
 rect_end=(5+round(barwidth/100*memories.percent,0), line+barheight)
 rect_outline=(5+barwidth,line+barheight)
 drawer.rectangle([rect_start, rect_outline],fill=None,outline=GREEN)
 drawer.rectangle([rect_start, rect_end],fill=GREEN,outline=None)
 line+=barheight+FONTSIZE['BIG']
 # From https://github.com/kubesail/pibox-os/blob/main/pwm-fan/pi_fan_hwpwm.c
 # Fan off = 60 / on = 65 / high = 80
 temp=round(psutil.sensors_temperatures()['cpu_thermal'][0].current, 1)
 tempcolor=GREEN if temp <= 66 else ORANGE if temp <= 80 else RED
 writewords(f"Temp: {temp}°C", line=3, color=tempcolor, alignment='RIGHT')

 # FOOTER
 hostname=socket.gethostname()
 # Trusts that IPv4 is in the address list before IPv6
 eth=psutil.net_if_addrs()[NET_ADAPTER][0]
 ip=eth.address
 cidr=SUBNET_CIDR_MAP[eth.netmask]
 writewords(f"{ip}/{cidr}", line=footerline, indent=5)
 writewords(f"{hostname}", line=footerline, alignment="RIGHT")
 return
 line+=barheight
 drawer.text((column[0], line), f"Total: {memtotal}G", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+5
 drawer.text((column[0], line), f"Avail: {memavailable}G", font=FONT['BIG'], fill=WHITE)
 line+=FONTSIZE['BIG']+5
 # Hello, world!
 drawer.text((column[0],line+24), "Hello, world!", font=FONT['BIG'], fill=WHITE)
 return

def guistats():
 headerline=round(FONTSIZE['BIG']*1.7,0)
 line=headerline+20
 padding=20
 barindent=75
 linepadding=10
 diskheight=15
 #nowstr=datetime.datetime.now().strftime("%Y-%m-%d %I:%M%p").lower()
 nowstr=datetime.datetime.now().strftime("%b%d %I:%M%p")
 timestr=datetime.datetime.now().strftime("%H:%M")
 datestr=datetime.datetime.now().strftime("%b %d ")

 ## Header
 # From https://github.com/kubesail/pibox-os/blob/main/pwm-fan/pi_fan_hwpwm.c
 # Fan off = 60 / on = 65 / high = 80
 temp=round(psutil.sensors_temperatures()['cpu_thermal'][0].current, 1)
 tempcolor=GREEN if temp <= 66 else ORANGE if temp <= 80 else RED
 writewords(f"{temp}°C", line=4, color=tempcolor, alignment='LEFT', indent=10)
 writewords(f"{timestr}", line=4, indent=5, alignment='CENTER')
 writewords(f"{datestr}", line=4, indent=5, alignment='RIGHT')
 drawer.line(((linepadding, headerline), (SCREENWIDTH-linepadding, headerline)), fill=WHITE)

 ## Body
 # CPU USAGE
 writewords("CPU", line=line,font="THICC", spacing=5, indent=padding)
 cpupercent=psutil.cpu_percent(interval=1)/100
 cpucolor=GREEN if cpupercent <= .70 else ORANGE if cpupercent <= .85 else RED
 corecount=psutil.cpu_count()
 fiveminload=psutil.getloadavg()[1]
 fiveminperc=round(fiveminload/corecount,2)
 fiveminperc=.5
 fivemincolor=DARK_GREEN if fiveminperc <= .70 else DARK_YELLOW if fiveminperc <= .85 else DARK_RED
 drawbar(top=line, indent=barindent, right=SCREENWIDTH-padding, corners=(True, True, False, False), color=fivemincolor, percent=fiveminperc)
 line=drawbar(top=line, indent=barindent, right=SCREENWIDTH-padding, corners=(True, True, False, False), color=cpucolor, percent=cpupercent)
 writewords(f"{int(cpupercent*100)}%", indent=barindent, right=SCREENWIDTH-padding,line=line-FONTSIZE['BIG']-8,alignment='CENTER',outline=BLACK)

 # MEMORY USAGE
 memories=psutil.virtual_memory()
 mempercent=round(memories.used/memories.total,2)
 memcolor=AQUA if mempercent <= .70 else ORANGE if mempercent <= .85 else RED
 memtotal=round(memories.total/1024/1024/1024, 2)
 memused=round(memories.used/1024/1024/1024, 2)
 memavailable=round(memories.available/1024/1024/1024, 2)
 writewords("RAM", line=line,font="THICC", spacing=5, indent=padding)
 line=drawbar(top=line, indent=barindent, right=SCREENWIDTH-padding, corners=(False, False, True, True), color=memcolor, percent=mempercent)
 writewords(f"{memused}G/{memtotal}G", indent=barindent, right=SCREENWIDTH-padding, line=line-23,alignment='CENTER',outline=BLACK, font='big')
#  line=drawbar(top=line, right=240-padding, corners=(False, False, True, True), color=memcolor, percent=mempercent)
#  writewords(f"{memused}G/{memtotal}G", line=line-27,alignment='CENTER',outline=BLACK, font='big')
#  line=writewords("Memory", line=line,font="THICC", spacing=5, indent=padding)
 line+=25

 # DISK USAGE
 rawdisklist=psutil.disk_partitions(all=False)
 disklist=[]
 diskcolors=[BLUE,PINK,WHITE]
 for disk in rawdisklist:
  #if 'boot' in disk.mountpoint:
  # continue
  disklist.append(disk)
  #print(f'found disk {disk.device}')
 if len(disklist) == 1:
  #diskname=disklist[0].device[substring_indicies(disklist[0].device, '/')[-1]+1:]
  diskname=disklist[0].mountpoint
  writewords(f"{diskname}", line, indent=padding, font='BIG')
  drawbar(top=line,indent=SCREENWIDTH/2,height=diskheight, right=SCREENWIDTH-padding, color=diskcolors[0],percent=psutil.disk_usage(disklist[0].mountpoint).percent/100)
 else:
  diskindex=0
  for disk in disklist:
   # e.g. /dev/mmcblk0p2 or /dev/mapper/pibox--group-k3s
   #diskname=disk.device[substring_indicies(disk.device, '/')[-1]+1:]
   # e.g. / or /mnt/pibox
   diskname=disk.mountpoint
   diskusage=psutil.disk_usage(disk.mountpoint).percent/100
   writewords(f"{diskname}", line=line, indent=padding, font='MED')
   if disk.device == disklist[0].device:
    corners=(True, True, False, False)
   elif disk.device == disklist[-1].device:
    corners=(False, False, True, True)
   else:
    corners=(False, False, False, False)
   line=drawbar(top=line,indent=SCREENWIDTH/2,height=diskheight, right=SCREENWIDTH-padding, color=diskcolors[diskindex%len(diskcolors)],percent=diskusage, corners=corners)
   diskindex+=1
 hostname=socket.gethostname()
 # Trusts that IPv4 is in the address list before IPv6
 eth=psutil.net_if_addrs()[NET_ADAPTER][0]
 ip=eth.address
 cidr=SUBNET_CIDR_MAP[eth.netmask]
 footerfont="BIG" if drawer.textlength(f"{ip}/{cidr}  {hostname}", FONT['BIG']) <= SCREENWIDTH else "MED"
 footerline=240-FONTSIZE[footerfont]-5
 writewords(f"{ip}/{cidr}", line=footerline, indent=5, font=footerfont)
 writewords(f"{hostname}", line=footerline, alignment="RIGHT", font=footerfont, indent=5)
 footerlineline=SCREENHEIGHT-round(FONTSIZE[footerfont]*1.8,0)
 drawer.line(((linepadding, footerlineline), (SCREENWIDTH-linepadding, footerlineline)), fill=WHITE)
 return

precheck()
#textstats()
#testguistats()
#texttest()
guistats()
#sendimage(mystats)
mystats.save("stats.png")