JARVIS/display.py

#!/usr/bin/env python3
"""
JARVIS Display
──────────────
Renders the JARVIS UI overlay.  Works in two modes:

  X11 (development laptop):
      python display.py

  Linux framebuffer (device – no windowing system):
      python display.py --framebuffer
      # or
      JARVIS_FB=1 python display.py

Framebuffer prerequisites on Debian (device):
  sudo apt install python3-pygame
  sudo usermod -aG video $USER   # then log out/in
  # If fbdev fails (common on newer kernels) try KMS/DRM:
  SDL_VIDEODRIVER=kmsdrm python display.py --framebuffer
"""

import os
import math
import argparse
from datetime import datetime

# ── CLI / env-var flags  ─────────────────────────────────────────────────────
# These must be parsed BEFORE pygame is imported so we can set SDL env vars
# in time for pygame.init() to pick them up.

_parser = argparse.ArgumentParser(description='JARVIS display', add_help=True)
_parser.add_argument(
    '--framebuffer', '-fb', action='store_true',
    help='Render to Linux framebuffer instead of an X11/Wayland window')
_parser.add_argument(
    '--fullscreen', '-fs', action='store_true',
    help='Run fullscreen under X11/Wayland')
_parser.add_argument('--width',  type=int, default=1280,
                     help='Screen / framebuffer width  (default 1280)')
_parser.add_argument('--height', type=int, default=800,
                     help='Screen / framebuffer height (default 800)')
_args, _ = _parser.parse_known_args()

USE_FB: bool = _args.framebuffer or bool(os.environ.get('JARVIS_FB'))

if USE_FB:
    # SDL2 env vars must be set before pygame.display.init()
    os.environ.setdefault('SDL_VIDEODRIVER', 'kmsdrm')
    os.environ.setdefault('SDL_RENDER_DRIVER', 'software')
    os.environ.setdefault('SDL_FBDEV', os.environ.get('JARVIS_FBDEV', '/dev/fb0'))
    # Prevent SDL trying to open a mouse device on the framebuffer
    os.environ.setdefault('SDL_NOMOUSE', '1')

import pygame  # noqa: E402  — import after env setup

# ── Colour palette ───────────────────────────────────────────────────────────
BLACK    = (  0,   0,   0)
WHITE    = (255, 255, 255)
GRAY     = (160, 160, 160)
DIM_GRAY = ( 80,  80,  80)
WIRE_COL = ( 58,  58,  58)   # subdued wireframe line colour

# ── Wireframe face ────────────────────────────────────────────────────────────
# Pure 2-D face drawing.  Face shape is defined by an explicit width profile;
# eye sockets, nose bridge and chin are drawn as distinct features so the
# brain immediately reads "face" rather than a geometric solid.

WIRE_EYE = ( 95,  95,  95)   # slightly brighter for eye/feature lines

def _hw(t: float, rx: float) -> float:
    """
    Face half-width in pixels at normalised height t.
      t = -1  → crown (top)
      t =  0  → cheek level (widest, ~= rx)
      t = +1  → chin tip
    Upper half: rounded ellipse.
    Lower half: faster taper toward chin.
    """
    if t <= 0.0:
        return rx * math.sqrt(max(0.0, 1.0 - t * t))
    else:
        return rx * math.sqrt(max(0.0, 1.0 - t * t)) * (1.0 - 0.30 * t)


def draw_wireframe_face(surface: pygame.Surface,
                        cx: int, cy: int,
                        rx: float, ry: float) -> None:
    """
    Draw a 2-D face wireframe centred at (cx, cy).
    Coordinate system:  t ∈ [−1, +1]  (−1=crown, 0=cheeks, +1=chin)
                        x_fraction ∈ [−1, +1]  (fraction of half-width at t)
    """
    N_H  = 13    # horizontal grid lines
    N_V  = 11    # vertical grid lines (including edges)
    PTS  = 40    # interpolation steps per curve

    # Helpers ─────────────────────────────────────────────────────────────────
    def sc(xf: float, t: float) -> tuple[int, int]:
        """xf (fraction of half-width) + height t → screen pixel."""
        w = _hw(t, rx)
        # Slight forward bow: face curves toward viewer at cheek level
        bow = int(0.06 * w * (1.0 - t * t) * (1.0 - xf * xf))
        return (int(cx + xf * w), int(cy + t * ry) - bow)

    def draw(pts, close=False, col=WIRE_COL):
        if len(pts) >= 2:
            pygame.draw.lines(surface, col, close, pts, 1)

    # Face outline (left and right silhouette) ────────────────────────────────
    t_vals = [(-1.0 + 2.0 * i / (PTS - 1)) for i in range(PTS)]
    draw([sc(-1.0, t) for t in t_vals])   # left edge
    draw([sc( 1.0, t) for t in t_vals])   # right edge
    # Crown arc (straight line across the very top)
    draw([sc(-1.0, -1.0), sc(1.0, -1.0)])
    # Chin arc
    draw([sc(-0.6, 1.0), sc(0.0, 1.02), sc(0.6, 1.0)])

    # Horizontal grid lines ───────────────────────────────────────────────────
    for i in range(1, N_H):
        t = -1.0 + 2.0 * i / N_H
        draw([sc(-1.0 + 2.0 * j / (PTS - 1), t) for j in range(PTS)])

    # Vertical grid lines ─────────────────────────────────────────────────────
    for i in range(1, N_V):
        xf = -1.0 + 2.0 * i / N_V
        draw([sc(xf, t) for t in t_vals])

    # ── Eye sockets ───────────────────────────────────────────────────────────
    # Each eye is a small ellipse drawn in a slightly brighter colour so it
    # immediately reads as a facial feature.
    eye_t   = -0.38            # height: upper face (~35% from crown)
    eye_xf  =  0.35            # lateral offset (fraction of half-width at eye_t)
    eye_w   = _hw(eye_t, rx)   # face half-width at eye level
    e_rx    = int(eye_w * 0.24)
    e_ry    = int(ry    * 0.09)

    for sign in (-1, +1):
        ecx = int(cx + sign * eye_xf * eye_w)
        ecy = int(cy + eye_t * ry)
        eye_pts = [
            (int(ecx + e_rx * math.cos(2 * math.pi * k / 28)),
             int(ecy + e_ry * math.sin(2 * math.pi * k / 28)))
            for k in range(28)
        ]
        draw(eye_pts + [eye_pts[0]], col=WIRE_EYE)

    # ── Nose bridge ───────────────────────────────────────────────────────────
    # Two short lines running from inner eye corners down to nose tip
    inner = int(eye_w * 0.12)
    e_btm = int(cy + eye_t * ry) + e_ry
    nose_y = int(cy + 0.08 * ry)
    nose_w = int(eye_w * 0.07)
    draw([(cx - inner, e_btm), (cx - nose_w, nose_y)], col=WIRE_COL)
    draw([(cx + inner, e_btm), (cx + nose_w, nose_y)], col=WIRE_COL)


# ── Clock / date ─────────────────────────────────────────────────────────────

def draw_clock(surface: pygame.Surface, fonts: dict,
               x: int, y: int) -> None:
    now = datetime.now()
    date_str = now.strftime('%A, %B %-d, %Y')
    time_str = now.strftime('%-H:%M')
    secs_str = now.strftime('%S')

    # Date row
    surface.blit(fonts['small'].render(date_str, True, GRAY), (x, y))

    # Large time + superscript seconds
    t_surf = fonts['large'].render(time_str, True, WHITE)
    s_surf = fonts['small'].render(secs_str,  True, DIM_GRAY)
    ty = y + fonts['small'].get_height() + 4
    surface.blit(t_surf, (x, ty))
    surface.blit(s_surf, (x + t_surf.get_width() + 2,
                           ty + t_surf.get_height() - s_surf.get_height() - 8))


# ── Status line ───────────────────────────────────────────────────────────────

def draw_status(surface: pygame.Surface, fonts: dict,
                text: str, cx: int, y: int) -> None:
    surf = fonts['small'].render(text, True, GRAY)
    surface.blit(surf, (cx - surf.get_width() // 2, y))


def draw_dot(surface: pygame.Surface, cx: int, y: int,
             radius: int = 9) -> None:
    pygame.draw.circle(surface, WHITE, (cx, y), radius)


# ── Font helper ───────────────────────────────────────────────────────────────

def _load_font(size: int, bold: bool = False) -> pygame.font.Font:
    """Try a list of clean system fonts, fall back to pygame built-in."""
    for name in ('DejaVuSans', 'FreeSans', 'LiberationSans',
                 'Helvetica', 'Arial', None):
        try:
            f = pygame.font.SysFont(name, size, bold=bold)
            if f:
                return f
        except Exception:
            pass
    return pygame.font.Font(None, size)


# ── Main loop ─────────────────────────────────────────────────────────────────

def main() -> None:
    pygame.init()

    W, H = _args.width, _args.height

    if USE_FB or _args.fullscreen:
        screen = pygame.display.set_mode((W, H), pygame.FULLSCREEN | pygame.NOFRAME)
    else:
        screen = pygame.display.set_mode((W, H))
        pygame.display.set_caption('JARVIS')

    pygame.mouse.set_visible(False)

    fonts = {
        'large': _load_font(56),
        'small': _load_font(18),
    }

    # Layout anchors — all relative to screen size so they adapt to resolution
    face_rx  = int(min(W, H) * 0.115)   # half-width at cheek level
    face_ry  = int(min(W, H) * 0.135)   # half-height crown→chin
    face_cx  = W // 2
    face_cy  = int(face_ry * 1.2) + 8   # nearly touches the top edge
    status_y = int(H * 0.82)
    dot_y    = int(H * 0.895)

    # ── Load face wireframe PNG if present ────────────────────────────────────
    # Place any face wireframe PNG at  assets/face_wire.png  and it will be
    # used automatically.  Black background is composited away via BLEND_ADD
    # so only the bright lines show.  Falls back to parametric drawing.
    _here     = os.path.dirname(os.path.abspath(__file__))
    _png_path = os.path.join(_here, 'assets', 'face_wire.png')
    face_img  = None
    if os.path.exists(_png_path):
        raw      = pygame.image.load(_png_path).convert()
        img_src_w, img_src_h = raw.get_size()
        # Scale to fit face slot while preserving aspect ratio
        scale    = min((face_rx * 2.2) / img_src_w,
                       (face_ry * 2.4) / img_src_h)
        img_w    = max(1, int(img_src_w * scale))
        img_h    = max(1, int(img_src_h * scale))
        face_img = pygame.transform.smoothscale(raw, (img_w, img_h))

    status_text = 'Initializing...'

    clock = pygame.time.Clock()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            if event.type == pygame.KEYDOWN:
                if event.key in (pygame.K_ESCAPE, pygame.K_q):
                    pygame.quit()
                    return
            # Touch / mouse tap anywhere also quits in kiosk mode
            if USE_FB and event.type == pygame.MOUSEBUTTONDOWN:
                pygame.quit()
                return

        screen.fill(BLACK)
        if face_img:
            screen.blit(face_img,
                        face_img.get_rect(center=(face_cx, face_cy)),
                        special_flags=pygame.BLEND_ADD)
        else:
            draw_wireframe_face(screen, face_cx, face_cy, face_rx, face_ry)
        clock_block_h = fonts['small'].get_height() + 4 + fonts['large'].get_height()
        draw_clock(screen, fonts, 20, face_cy - clock_block_h // 2)
        draw_status(screen, fonts, status_text, W // 2, status_y)
        draw_dot(screen, W // 2, dot_y)

        pygame.display.flip()
        clock.tick(10)   # 10 fps is plenty for a status screen


if __name__ == '__main__':
    main()