#!/usr/bin/env python3
"""Capture one still from the camera with the Raspberry Pi libcamera fork, as a JPEG.

  capture_still.py OUT.jpg [--frames N] [--width W --height H]

Runs `cam` from /opt/libcamera-rpi for N frames (so auto-exposure settles) and keeps the last.
Run it with the venv's python so OpenCV is the newer one: ~/venv/bin/python capture_still.py ...
"""
import argparse
import glob
import os
import subprocess
import tempfile

import cv2
import numpy as np

PREFIX = '/opt/libcamera-rpi'


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('out')
    p.add_argument('--frames', type=int, default=45)
    p.add_argument('--width', type=int, default=1296)
    p.add_argument('--height', type=int, default=972)
    a = p.parse_args()
    env = dict(os.environ, LD_LIBRARY_PATH=f'{PREFIX}/lib/aarch64-linux-gnu')
    with tempfile.TemporaryDirectory() as tmp:
        subprocess.run([f'{PREFIX}/bin/cam', '-c', '1', '--stream',
                        f'pixelformat=RGB888,width={a.width},height={a.height}',
                        f'--capture={a.frames}', f'--file={tmp}/f-#.bin'],
                       env=env, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=90)
        last = sorted(glob.glob(f'{tmp}/f-*.bin'))[-1]
        raw = np.fromfile(last, np.uint8)
    stride = raw.size // a.height              # rows are padded
    img = raw.reshape(a.height, stride)[:, :a.width * 3].reshape(a.height, a.width, 3)   # B,G,R in memory
    cv2.imwrite(a.out, img)
    print(f'{a.out}: {a.width}x{a.height}, mean brightness {img.mean():.1f}/255')


if __name__ == '__main__':
    main()