picox/src/main.rs

mod api;
mod config;
mod model;
mod state;
mod storage_manager;

use crate::config::PicOxConfig;
use crate::model::album::Album;
use crate::model::image::{Image, ImageData};
use crate::state::Context;
use crate::storage_manager::{StorageManager, StoreError};
use axum::extract::{Multipart, Path, Query, State};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{Json, Router};
use axum_macros::FromRequest;
use j_db::database::Database;
use j_db::model::JdbModel;
use log::info;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use structopt::StructOpt;
use tokio::sync::RwLock;

type PicContext = Arc<Context>;

#[derive(Debug, Clone, StructOpt)]
struct Args {
    pub config: PathBuf,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CreateAlbum {
    pub album_name: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AlbumQuery {
    pub album_name: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AddImage {
    pub album: AlbumQuery,
    pub tags: Vec<String>,
}

#[derive(FromRequest)]
#[from_request(via(axum::Json), rejection(PicOxError))]
struct Response<T>(T);

impl<T> IntoResponse for Response<T>
where
    Json<T>: IntoResponse,
{
    fn into_response(self) -> axum::response::Response {
        Json(self.0).into_response()
    }
}

enum PicOxError {
    StoreError(StoreError),
    DbError(j_db::error::JDbError),
}

impl From<StoreError> for PicOxError {
    fn from(value: StoreError) -> Self {
        Self::StoreError(value)
    }
}

impl From<j_db::error::JDbError> for PicOxError {
    fn from(value: j_db::error::JDbError) -> Self {
        Self::DbError(value)
    }
}

#[derive(Serialize)]
struct ErrorResponse {
    pub message: String,
}

impl IntoResponse for PicOxError {
    fn into_response(self) -> axum::response::Response {
        let (status, message) = match self {
            PicOxError::StoreError(err) => match err {
                StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()),
                StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()),
                StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()),
                StoreError::IOError(_) => (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "IO Error Has Occurred!".to_string(),
                ),
            },
            PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
        };

        (status, Response(ErrorResponse { message })).into_response()
    }
}

async fn create_album(
    State(context): State<PicContext>,
    Json(album): Json<CreateAlbum>,
) -> Result<Response<Album>, PicOxError> {
    info!("Creating new album '{}'", album.album_name);

    let new_album = Album::new(&album.album_name, Vec::new(), 0);

    let new_album = context.db.insert(new_album)?;

    Ok(Response(new_album))
}

async fn get_album(
    album_id: Path<u64>,
    State(context): State<PicContext>,
) -> Result<Response<Album>, PicOxError> {
    let album = context.db.get::<Album>(*album_id)?;

    Ok(Response(album))
}

async fn add_image(
    State(context): State<PicContext>,
    mut img_data: Multipart,
) -> Result<Response<Image>, PicOxError> {
    let mut data: Vec<u8> = Vec::new();
    let mut metadata: Option<AddImage> = None;
    let mut file_name = None;
    while let Some(field) = img_data.next_field().await.unwrap() {
        let field_name = field.name();
        if let Some(field_name) = field_name {
            if field_name == "metadata" {
                let metadata_json = field.text().await.unwrap();
                metadata = Some(serde_json::from_str(&metadata_json).unwrap())
            } else if field_name == "img_data" {
                file_name = Some(field.file_name().unwrap().to_string());
                data.extend_from_slice(field.bytes().await.unwrap().as_ref());
            }
        }
    }

    let mut album =
        Album::find_album_by_query(&context.db, metadata.clone().unwrap().album).unwrap();

    let mut store_manager = context.store_manager.write().await;
    let img = store_manager
        .store_img(
            &context.db,
            None,
            ImageData::Bytes(data),
            &file_name.unwrap(),
            0,
        )
        .await?;

    album.images.push(img.id().unwrap());

    context.db.insert::<Album>(album)?;

    Ok(Response(img))
}

async fn query_album(
    album_query: Query<AlbumQuery>,
    State(context): State<PicContext>,
) -> Result<Response<Option<Album>>, PicOxError> {
    let resp = Album::find_album_by_query(&context.db, album_query.0);
    Ok(Response(resp))
}

#[tokio::main]
async fn main() {
    let args = Args::from_args();

    let config = PicOxConfig::new(args.config);

    let db = Database::new(&config.db_path).unwrap();

    let store_manager = StorageManager::new(config.storage_config.clone());

    let context = Context {
        db,
        config: config.clone(),
        store_manager: RwLock::new(store_manager),
    };

    let context = Arc::new(context);

    // initialize tracing
    tracing_subscriber::fmt::init();
    // build our application with a route
    let app = Router::new()
        // `GET /` goes to `root`
        .route("/api/album/create", post(create_album))
        .route("/api/album/:id", get(get_album))
        .route("/api/album/", get(query_album))
        .route("/api/image/", post(add_image))
        .with_state(context);

    // run our app with hyper, listening globally on port 3000
    let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap();
    info!("Serving at {}", config.host);
    axum::serve(listener, app).await.unwrap();
}