# EliteDesk Status Display - Pico MicroPython
# Vixy's sketch - Day 35

from machine import Pin, SPI, PWM
import sys
import select
import time

import st7789

import vga1_16x32 as font_large  # For node name
import vga1_8x16 as font_small   # For labels/values

# ========== Configuration ==========

# Display pins (all GP16+ side for right-angle mounting)
SPI_ID = 0      # SPI0 for GP16+ pins
SCK_PIN = 18    # SPI0 SCK
MOSI_PIN = 19   # SPI0 TX/MOSI
CS_PIN = 17     # Chip select
DC_PIN = 16     # Data/Command
RST_PIN = 20    # Reset
BL_PIN = 21     # Backlight (PWM)

# Display dimensions (landscape)
WIDTH = 320
HEIGHT = 240

# Colors (RGB565)
BLACK = 0x0000
WHITE = 0xFFFF
RED = 0xF800
GREEN = 0x07E0
YELLOW = 0xFFE0
CYAN = 0x07FF
ORANGE = 0xFD20
DARK_GRAY = 0x4208
LIGHT_GRAY = 0x8410

# Thresholds
CPU_WARN = 70
CPU_CRIT = 90
MEM_WARN = 75
MEM_CRIT = 90
TEMP_WARN = 65
TEMP_CRIT = 80

# Layout constants
BAR_X = 70
BAR_W = 180
BAR_H = 18
VAL_X = 260

# ========== Display Setup ==========

def init_display():
    """Initialize ST7789 display"""
    spi = SPI(SPI_ID, baudrate=40000000, polarity=1, phase=1,
              sck=Pin(SCK_PIN), mosi=Pin(MOSI_PIN))
    
    display = st7789.ST7789(
        spi, WIDTH, HEIGHT,
        reset=Pin(RST_PIN, Pin.OUT),
        dc=Pin(DC_PIN, Pin.OUT),
        cs=Pin(CS_PIN, Pin.OUT),
        rotation=1  # Landscape
    )
    
    # Backlight at 70%
    bl = PWM(Pin(BL_PIN))
    bl.freq(1000)
    bl.duty_u16(45875)
    
    return display

# ========== Drawing Functions ==========

def get_status_color(value, warn_thresh, crit_thresh):
    """Return color based on thresholds"""
    if value >= crit_thresh:
        return RED
    elif value >= warn_thresh:
        return ORANGE
    else:
        return GREEN

def clear_rect(display, x, y, w, h):
    """Clear a rectangle to black"""
    display.fill_rect(x, y, w, h, BLACK)

def draw_bar(display, x, y, width, height, value, max_val, color):
    """Draw a progress bar"""
    display.fill_rect(x, y, width, height, DARK_GRAY)
    fill_width = int((value / max_val) * width)
    if fill_width > 0:
        display.fill_rect(x, y, fill_width, height, color)
    display.rect(x, y, width, height, WHITE)

def draw_text(display, font, text, x, y, color=WHITE):
    """Draw text"""
    display.text(font, text, x, y, color)

def draw_static_elements(display):
    """Draw elements that never change"""
    # Divider lines
    display.hline(0, 45, WIDTH, DARK_GRAY)
    display.hline(0, 165, WIDTH, DARK_GRAY)
    
    # Static labels
    draw_text(display, font_small, "CPU", 10, 61, WHITE)
    draw_text(display, font_small, "MEM", 10, 96, WHITE)
    draw_text(display, font_small, "TEMP", 10, 131, WHITE)

def draw_node(display, node):
    """Draw node name"""
    clear_rect(display, 10, 8, 230, 32)
    draw_text(display, font_large, node, 10, 8, CYAN)

def draw_status_badge(display, status):
    """Draw status indicator badge"""
    if status == 'healthy':
        color = GREEN
    elif status == 'warning':
        color = ORANGE
    else:
        color = RED
    
    display.fill_rect(250, 12, 60, 24, color)
    status_short = status[:6].upper()
    draw_text(display, font_small, status_short, 255, 16, BLACK)

def draw_cpu(display, cpu):
    """Draw CPU bar and value"""
    y = 60
    color = get_status_color(cpu, CPU_WARN, CPU_CRIT)
    draw_bar(display, BAR_X, y, BAR_W, BAR_H, cpu, 100, color)
    clear_rect(display, VAL_X, y + 1, 50, 16)
    draw_text(display, font_small, f"{cpu}%", VAL_X, y + 1, color)

def draw_mem(display, mem):
    """Draw memory bar and value"""
    y = 95
    color = get_status_color(mem, MEM_WARN, MEM_CRIT)
    draw_bar(display, BAR_X, y, BAR_W, BAR_H, mem, 100, color)
    clear_rect(display, VAL_X, y + 1, 50, 16)
    draw_text(display, font_small, f"{mem}%", VAL_X, y + 1, color)

def draw_temp(display, temp):
    """Draw temperature bar and value"""
    y = 130
    color = get_status_color(temp, TEMP_WARN, TEMP_CRIT)
    draw_bar(display, BAR_X, y, BAR_W, BAR_H, temp, 100, color)
    clear_rect(display, VAL_X, y + 1, 50, 16)
    draw_text(display, font_small, f"{temp}C", VAL_X, y + 1, color)

def draw_pods(display, pods):
    """Draw pods count"""
    clear_rect(display, 10, 180, 200, 32)
    draw_text(display, font_large, f"PODS: {pods}", 10, 180, YELLOW)

def draw_full_screen(display, stats):
    """Draw everything - used on startup"""
    display.fill(BLACK)
    draw_static_elements(display)
    draw_node(display, stats.get('NODE', '???'))
    draw_status_badge(display, stats.get('STATUS', 'unknown'))
    draw_cpu(display, stats.get('CPU', 0))
    draw_mem(display, stats.get('MEM', 0))
    draw_temp(display, stats.get('TEMP', 0))
    draw_pods(display, stats.get('PODS', 0))

def update_display(display, old_stats, new_stats):
    """Update only changed elements"""
    if old_stats.get('NODE') != new_stats.get('NODE'):
        draw_node(display, new_stats.get('NODE', '???'))
    
    if old_stats.get('STATUS') != new_stats.get('STATUS'):
        draw_status_badge(display, new_stats.get('STATUS', 'unknown'))
    
    if old_stats.get('CPU') != new_stats.get('CPU'):
        draw_cpu(display, new_stats.get('CPU', 0))
    
    if old_stats.get('MEM') != new_stats.get('MEM'):
        draw_mem(display, new_stats.get('MEM', 0))
    
    if old_stats.get('TEMP') != new_stats.get('TEMP'):
        draw_temp(display, new_stats.get('TEMP', 0))
    
    if old_stats.get('PODS') != new_stats.get('PODS'):
        draw_pods(display, new_stats.get('PODS', 0))

# ========== Serial Communication ==========

def parse_stats(lines):
    """Parse stats from serial lines"""
    stats = {}
    for line in lines:
        line = line.strip()
        if ':' in line:
            key, value = line.split(':', 1)
            key = key.strip().upper()
            value = value.strip()
            
            if key in ('CPU', 'MEM', 'TEMP', 'PODS'):
                try:
                    value = int(value)
                except ValueError:
                    try:
                        value = float(value)
                    except ValueError:
                        pass
            
            stats[key] = value
    
    return stats

def read_serial():
    """Read available lines from USB serial"""
    lines = []
    poll = select.poll()
    poll.register(sys.stdin, select.POLLIN)
    
    while poll.poll(0):
        line = sys.stdin.readline()
        if line:
            lines.append(line)
        else:
            break
    
    return lines

# ========== Main Loop ==========

def main():
    print("EliteDesk Status Display - Starting...")
    
    display = init_display()
    
    stats = {
        'NODE': 'ed?',
        'CPU': 0,
        'MEM': 0,
        'TEMP': 0,
        'PODS': 0,
        'STATUS': 'waiting'
    }
    prev_stats = {}
    
    # Initial full draw
    draw_full_screen(display, stats)
    prev_stats = stats.copy()
    print("Display initialized")
    
    buffer = []
    last_update = time.ticks_ms()
    
    while True:
        lines = read_serial()
        buffer.extend(lines)
        
        for i, line in enumerate(buffer):
            if line.strip().upper().startswith('STATUS:'):
                new_stats = parse_stats(buffer[:i+1])
                if new_stats:
                    stats.update(new_stats)
                buffer = buffer[i+1:]
                break
        
        # Update display every second, only changed elements
        if time.ticks_diff(time.ticks_ms(), last_update) > 1000:
            update_display(display, prev_stats, stats)
            prev_stats = stats.copy()
            last_update = time.ticks_ms()
        
        time.sleep_ms(10)

if __name__ == '__main__':
    main()