pi-frame-server/src/api.rs

use crate::display::{create_display_thread, EInkPanel};
use crate::dither::{DitherMethod, DitheredImage};
use crate::eink::Palette;
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::{imageops::resize, imageops::FilterType, ImageReader, RgbImage};
use std::io::Cursor;
use std::str;
use std::str::FromStr;
use std::sync::mpsc;
use std::sync::Arc;
use std::thread::JoinHandle;
use tracing::{error, info, instrument};

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

#[derive(Clone)]
pub struct AppState {
    display_channel: mpsc::Sender<Box<DitheredImage>>,
    display_task: Arc<JoinHandle<()>>,
}

impl AppState {
    #[must_use]
    pub fn new(disp: Box<dyn EInkPanel + Send>) -> Self {
        let (handle, tx) = create_display_thread(disp);
        Self {
            display_channel: tx,
            display_task: Arc::new(handle),
        }
    }
}

// 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())
    }
}

/// 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<AppState> {
    Router::new()
        .route("/setimage", post(set_image))
        .route("/preview", post(preview_image))
}

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

#[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(Palette::from_str(val)?);
                }
                Some("dither_method") => {
                    let data = field.bytes().await?;
                    let val = str::from_utf8(&data)?;
                    dither_method = Some(DitherMethod::from_str(val)?);
                }
                _ => {}
            }
        }
        img.map_or_else(
            || Err(ApiError::MissingImage.into()),
            |i| {
                Ok(Self {
                    image: i,
                    dither_method: dither_method.unwrap_or(DitherMethod::NearestNeighbor),
                    palette: palette.unwrap_or(Palette::Default),
                })
            },
        )
    }
}

#[instrument(skip_all)]
async fn set_image(
    State(ctx): State<AppState>,
    img_req: ImageRequest,
) -> Result<impl IntoResponse, AppError> {
    // FIXME: resize image to 800x480 to match the eink panel.
    info!("Got image");
    let mut buf = DitheredImage::new(800, 480, img_req.palette.value().to_vec());
    let resized = resize(&*img_req.image, 800, 480, FilterType::Lanczos3);
    {
        let mut dither = img_req.dither_method.get_ditherer();
        dither.dither(&resized, &mut buf);
    }
    info!("image resized, pushing to channel");
    ctx.display_channel.send(Box::new(buf))?;
    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 = DitheredImage::new(800, 480, img_req.palette.value().to_vec());
    let resized = resize(&*img_req.image, 800, 480, FilterType::Lanczos3);
    {
        let mut dither = img_req.dither_method.get_ditherer();
        dither.dither(&resized, &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()))
}