pi-frame-server/src/api.rs

use crate::display::EInkPanel;
use crate::imageproc::{DitherMethod, DitherPalette, EInkImage};
use axum::async_trait;
use axum::extract::{FromRequest, Multipart, State};
use axum::http::{header, StatusCode};
use axum::response::IntoResponse;
use axum::{response::Response, routing::post, Router};
use image::{ImageReader, RgbImage};
use std::io::Cursor;
use std::str;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc::{self, Receiver, Sender};
use tokio::task::JoinHandle;
use tracing::{debug, error, info, instrument};

#[derive(thiserror::Error, Debug)]
pub enum ApiError {
    #[error("missing image field")]
    MissingImage,
}

#[derive(Clone)]
pub struct Context {
    display_channel: Sender<DisplaySetCommand>,
    display_task: Arc<JoinHandle<()>>,
}

impl Context {
    #[must_use]
    pub fn new(disp: Box<dyn EInkPanel + Send>) -> Self {
        let (tx, rx) = mpsc::channel(2);
        let task = tokio::spawn(display_task(rx, disp));
        Self {
            display_channel: tx,
            display_task: Arc::new(task),
        }
    }
}

// Make our own error that wraps `anyhow::Error`.
struct AppError(anyhow::Error);

// Tell axum how to convert `AppError` into a response.
impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Something went wrong: {}", self.0),
        )
            .into_response()
    }
}

// This enables using `?` on functions that return `Result<_, anyhow::Error>` to turn them into
// `Result<_, AppError>`. That way you don't need to do that manually.
impl<E> From<E> for AppError
where
    E: Into<anyhow::Error>,
{
    fn from(err: E) -> Self {
        Self(err.into())
    }
}

#[derive(Debug)]
pub struct DisplaySetCommand {
    img: Box<EInkImage>,
}

#[instrument(skip_all)]
pub async fn display_task(
    mut rx: Receiver<DisplaySetCommand>,
    mut display: Box<dyn EInkPanel + Send>,
) {
    while let Some(cmd) = rx.recv().await {
        info!("Got a display set command");
        if let Err(e) = display.display(&cmd.img) {
            error!("Error displaying command {e}");
        }
        info!("Done setting display");
    }
}

/// API routes for axum
/// Start with the basics: Send an image, crop it, dither, and upload.
/// we defer the upload to a separate task.
pub fn router() -> Router<Context> {
    Router::new()
        .route("/setimage", post(set_image))
        .route("/preview", post(preview_image))
}

#[derive(Debug)]
struct ImageRequest {
    image: Box<RgbImage>,
    dither_method: DitherMethod,
    palette: DitherPalette,
}

#[async_trait]
impl<S> FromRequest<S> for ImageRequest
where
    S: Send + Sync,
{
    type Rejection = AppError;

    async fn from_request(req: axum::extract::Request, state: &S) -> Result<Self, Self::Rejection> {
        let mut parts = Multipart::from_request(req, state).await?;
        let mut img = None;
        let mut palette = None;
        let mut dither_method = None;
        while let Some(field) = parts.next_field().await? {
            match field.name() {
                Some("image") => {
                    let data = field.bytes().await?;
                    let reader = ImageReader::new(Cursor::new(data))
                        .with_guessed_format()
                        .expect("cursor never fails");
                    let image = reader.decode()?;
                    img = Some(Box::new(image.into()));
                }
                Some("palette") => {
                    let data = field.bytes().await?;
                    let val = str::from_utf8(&data)?;
                    palette = Some(DitherPalette::from_str(val)?);
                }
                Some("dither_method") => {
                    let data = field.bytes().await?;
                    let val = str::from_utf8(&data)?;
                    dither_method = Some(DitherMethod::from_str(val)?);
                }
                _ => {}
            }
        }
        if let Some(i) = img {
            Ok(Self {
                image: i,
                dither_method: dither_method.unwrap_or(DitherMethod::NearestNeighbor),
                palette: palette.unwrap_or(DitherPalette::Default),
            })
        } else {
            Err(ApiError::MissingImage.into())
        }
    }
}

#[instrument(skip(ctx))]
async fn set_image(
    State(ctx): State<Context>,
    img_req: ImageRequest,
) -> Result<impl IntoResponse, AppError> {
    let mut buf = EInkImage::new(img_req.palette.value().to_vec());
    {
        let mut dither = img_req.dither_method.get_ditherer();
        dither.dither(&img_req.image, &mut buf);
    }
    let cmd = DisplaySetCommand { img: Box::new(buf) };
    ctx.display_channel
        .send_timeout(cmd, Duration::from_secs(10))
        .await?;
    Ok(StatusCode::OK)
}

/// generates a dithered image based on the given image and the dithering parameters.
/// Can be used to see how the dithering and palette choices affect the result.
async fn preview_image(img_req: ImageRequest) -> Result<impl IntoResponse, AppError> {
    let mut buf = EInkImage::new(img_req.palette.value().to_vec());
    {
        let mut dither = img_req.dither_method.get_ditherer();
        dither.dither(&img_req.image, &mut buf);
    }
    // Convert buf into a png image.
    let img = buf.into_rgbimage();

    let mut buffer = Cursor::new(Vec::new());
    img.write_to(&mut buffer, image::ImageFormat::Png)?;

    let headers = [(header::CONTENT_TYPE, mime::IMAGE_PNG.to_string())];
    Ok((StatusCode::OK, headers, buffer.into_inner()))
}