#!/usr/bin/env python3
"""Bench-test the Waveshare 2.13" e-Paper HAT.

Alternates every 5 seconds between:
  - a data screen: time, CPU temp, BME280, SD card storage, recording status
  - a camera preview: both stereo eyes side by side, dithered to 1-bit

Prereqs on the Pi (see README / install.sh for the base packages):
  sudo apt install -y python3-gpiozero python3-spidev python3-lgpio python3-pil git ffmpeg
  git clone --depth 1 https://github.com/waveshare/e-Paper ~/e-Paper
  sudo pip3 install --break-system-packages ~/e-Paper/RaspberryPi_JetsonNano/python
  sudo pip3 install --break-system-packages smbus2 RPi.bme280

The camera must not be in use by anything else (stop bubblecam-camera@*
services first) or frame grabs will fail -- the preview then shows a
"no frame" message instead.

Run:
  python3 epaper_test.py
  python3 epaper_test.py --driver epd2in13_V3     # if the PCB sticker says V3
  python3 epaper_test.py --no-camera              # data screen only
"""

from __future__ import annotations

import argparse
import importlib
import shutil
import subprocess
import sys
import tempfile
import time
from datetime import datetime
from pathlib import Path

WIDTH, HEIGHT = 250, 122  # landscape, matches bubblecam/hardware/epaper.py

FONT_PATH = "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"

SCREEN_SECONDS = 5
# Full refresh (with the black/white flash) every N partial updates to clear
# the ghosting that dithered photo content builds up on e-paper.
FULL_REFRESH_EVERY = 20


def cpu_temp_c() -> float:
    with open("/sys/class/thermal/thermal_zone0/temp") as f:
        return int(f.read().strip()) / 1000.0


def sd_line() -> str:
    du = shutil.disk_usage("/")
    used_gb = (du.total - du.free) / 1e9
    total_gb = du.total / 1e9
    pct = 100 * (du.total - du.free) / du.total
    return f"SD: {used_gb:.1f}/{total_gb:.1f}GB ({pct:.0f}%)"


def recording_status() -> bool:
    """Mirror state_lib.is_running("camera_primary"): the camera service's
    /run/bubblecam heartbeat says running and is <30s fresh."""
    import json

    try:
        with open("/run/bubblecam/camera_primary.json") as f:
            state = json.load(f)
    except (FileNotFoundError, ValueError, OSError):
        return False
    last = state.get("last_update")
    if not isinstance(last, (int, float)):
        return False
    return state.get("state") == "running" and (time.time() - last) <= 30


class BME280:
    """Minimal wrapper mirroring bubblecam/hardware/bme280.py (smbus2 + RPi.bme280)."""

    def __init__(self, bus: int, address: int):
        import bme280
        import smbus2

        self._bme280 = bme280
        self._address = address
        self._bus = smbus2.SMBus(bus)
        self._cal = bme280.load_calibration_params(self._bus, address)

    def read(self) -> dict:
        s = self._bme280.sample(self._bus, self._address, self._cal)
        return {"t": s.temperature, "h": s.humidity, "p": s.pressure}


def load_fonts():
    from PIL import ImageFont

    try:
        return ImageFont.truetype(FONT_PATH, 15), ImageFont.truetype(FONT_PATH, 13)
    except OSError:
        f = ImageFont.load_default()
        return f, f


def build_data_image(bme: "BME280 | None"):
    from PIL import Image, ImageDraw

    image = Image.new("1", (WIDTH, HEIGHT), 255)
    draw = ImageDraw.Draw(image)
    font_big, font_small = load_fonts()

    if bme is not None:
        try:
            r = bme.read()
            bme_line = f"BME280: {r['t']:.1f}C {r['h']:.0f}%RH {r['p']:.0f}hPa"
        except Exception as exc:
            bme_line = f"BME280 read err: {type(exc).__name__}"
    else:
        bme_line = "BME280: not found"

    lines = [
        datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        f"CPU temp: {cpu_temp_c():.1f} C",
        bme_line,
        sd_line(),
        "Camera: RECORDING" if recording_status() else "Camera: NOT RECORDING",
    ]
    y = 6
    for line in lines:
        draw.text((6, y), line, font=font_big if y == 6 else font_small, fill=0)
        y += 22

    draw.rectangle((0, 0, WIDTH - 1, HEIGHT - 1), outline=0)
    return image


def grab_camera_frame(device: str, cam_size: str, timeout_s: float = 5.0):
    """Grab one MJPEG frame via ffmpeg; return a PIL Image or None."""
    from PIL import Image

    tmp = Path(tempfile.gettempdir()) / "epaper_cam_frame.jpg"
    cmd = [
        "ffmpeg", "-loglevel", "error", "-y",
        "-f", "v4l2", "-input_format", "mjpeg", "-video_size", cam_size,
        "-i", device, "-frames:v", "1", str(tmp),
    ]
    try:
        subprocess.run(cmd, check=True, timeout=timeout_s, capture_output=True)
        return Image.open(tmp)
    except Exception:
        return None


def build_camera_image(device: str, cam_size: str):
    """Both stereo eyes side by side (the camera's native combined frame),
    scaled + dithered to 1-bit, with a divider and L/R labels."""
    from PIL import Image, ImageDraw

    frame = grab_camera_frame(device, cam_size)

    canvas = Image.new("1", (WIDTH, HEIGHT), 255)
    draw = ImageDraw.Draw(canvas)
    _, font_small = load_fonts()

    if frame is None:
        msg = ("camera busy (recording?)" if recording_status()
               else "camera: no frame")
        draw.text((6, 50), msg, font=font_small, fill=0)
        draw.rectangle((0, 0, WIDTH - 1, HEIGHT - 1), outline=0)
        return canvas

    # The combined frame is already left eye | right eye. Fit the whole
    # thing, preserving aspect; grayscale then Floyd-Steinberg dither to
    # 1-bit (PIL's convert("1") default), which reads far better on
    # e-paper than a hard threshold.
    frame.thumbnail((WIDTH, HEIGHT))
    mono = frame.convert("L").convert("1")
    x0 = (WIDTH - mono.width) // 2
    y0 = (HEIGHT - mono.height) // 2
    canvas.paste(mono, (x0, y0))

    # Divider between the eyes + labels.
    mid = x0 + mono.width // 2
    draw.line((mid, y0, mid, y0 + mono.height), fill=0)
    draw.text((x0 + 2, y0), "L", font=font_small, fill=0)
    draw.text((mid + 3, y0), "R", font=font_small, fill=0)
    return canvas


def main() -> int:
    parser = argparse.ArgumentParser(description="e-Paper HAT bench test: data + camera preview")
    parser.add_argument(
        "--driver", default="epd2in13_V4",
        help="waveshare_epd driver module matching the HAT's PCB revision "
             "(epd2in13_V4 / epd2in13_V3 / epd2in13_V2; default V4)",
    )
    parser.add_argument("--bme-bus", type=int, default=1, help="I2C bus (default 1)")
    parser.add_argument("--bme-addr", type=lambda s: int(s, 0), default=0x76,
                        help="BME280 I2C address (default 0x76; some boards are 0x77)")
    parser.add_argument("--cam-dev", default="/dev/video0", help="camera device (default /dev/video0)")
    parser.add_argument("--cam-size", default="1280x480",
                        help="combined side-by-side capture size (default 1280x480)")
    parser.add_argument("--no-camera", action="store_true", help="data screen only, no preview")
    args = parser.parse_args()

    try:
        bme = BME280(args.bme_bus, args.bme_addr)
        print(f"BME280 found on bus {args.bme_bus} addr {args.bme_addr:#x}")
    except Exception as exc:
        bme = None
        print(f"BME280 unavailable ({exc}) -- will show 'not found' on screen")

    print(f"Loading driver waveshare_epd.{args.driver} ...")
    driver = importlib.import_module(f"waveshare_epd.{args.driver}")
    epd = driver.EPD()

    print("Initializing panel (a few seconds, expect black/white flashing)...")
    epd.init()
    epd.Clear(0xFF)

    partial = hasattr(epd, "displayPartial") and hasattr(epd, "displayPartBaseImage")
    show_camera = False  # start on the data screen
    updates = 0

    try:
        image = build_data_image(bme)
        if partial:
            epd.displayPartBaseImage(epd.getbuffer(image))
        else:
            epd.display(epd.getbuffer(image))
        print(f"First frame up. Alternating data/camera every {SCREEN_SECONDS}s "
              f"({'partial' if partial else 'FULL'} refresh) -- Ctrl+C to stop.")

        while True:
            time.sleep(SCREEN_SECONDS)
            show_camera = (not show_camera) and not args.no_camera

            if show_camera:
                image = build_camera_image(args.cam_dev, args.cam_size)
            else:
                image = build_data_image(bme)

            updates += 1
            if partial and updates % FULL_REFRESH_EVERY != 0:
                epd.displayPartial(epd.getbuffer(image))
            else:
                # periodic full refresh clears accumulated ghosting
                epd.init()
                epd.display(epd.getbuffer(image))
                if partial:
                    epd.displayPartBaseImage(epd.getbuffer(image))

            print(f"{datetime.now().strftime('%H:%M:%S')} "
                  f"{'camera' if show_camera else 'data'} screen")
    except KeyboardInterrupt:
        print("\nStopping.")
    finally:
        print("Putting panel to deep sleep (image stays visible).")
        epd.sleep()

    return 0


if __name__ == "__main__":
    sys.exit(main())
