picox/src/api/mod.rs

use crate::api::models::{
    AddImage, AlbumQuery, CreateAlbum, ImageQuery, ImageSort, PicContext, PicOxError, Response,
};
use crate::config::PicOxConfig;
use crate::model::album::Album;
use crate::model::api_key::ApiKey;
use crate::model::image::{Image, ImageData};
use crate::state::Context;
use crate::storage_manager::StorageManager;
use axum::body::Bytes;
use axum::extract::{DefaultBodyLimit, Multipart, Path, Query, Request, State};
use axum::http::HeaderMap;
use axum::middleware::Next;
use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{middleware, Json, Router};
use j_db::database::Database;
use j_db::model::JdbModel;
use log::info;
use rand::prelude::SliceRandom;
use rand::thread_rng;
use std::sync::Arc;
use tokio::sync::RwLock;

pub mod models;

async fn create_album(
    State(context): State<PicContext>,
    headers: HeaderMap,
    Json(album): Json<CreateAlbum>,
) -> Result<Response<Album>, PicOxError> {
    let user_id = get_user_id_from_headers(&headers)?;
    let new_album = Album::new(&album.album_name, Vec::new(), user_id);

    info!(
        "Creating new album '{}pub pub pub ' for user {}",
        album.album_name, user_id
    );
    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 get_image(
    image_id: Path<u64>,
    State(context): State<PicContext>,
) -> Result<Response<Image>, PicOxError> {
    let image = context.db.get::<Image>(*image_id)?;

    Ok(Response(image))
}

async fn query_images(
    State(context): State<PicContext>,
    Json(image_query): Json<ImageQuery>,
) -> Result<Response<Vec<Image>>, PicOxError> {
    let album_id = if let Some(album) = &image_query.album {
        Some(
            Album::find_album_by_query(
                &context.db,
                AlbumQuery {
                    album_name: Some(album.to_string()),
                },
            )
            .first()
            .ok_or(PicOxError::AlbumNotFound)?
            .id()
            .unwrap(),
        )
    } else {
        None
    };

    let mut images: Vec<Image> = context
        .db
        .filter(|_, img: &Image| {
            if let Some(album_id) = album_id {
                if img.album != album_id {
                    return false;
                }
            }

            if !image_query.tags.is_empty() {
                let mut found = false;
                for tag in &image_query.tags {
                    if img.tags.contains(tag) {
                        found = true;
                        break;
                    }
                }

                if !found {
                    return false;
                }
            }

            true
        })?
        .collect();

    match image_query.order {
        ImageSort::Random => {
            images.shuffle(&mut thread_rng());
        }
        ImageSort::DateAscending => {
            images.sort_by(|img_a, img_b| img_a.create_date.cmp(&img_b.create_date))
        }
        ImageSort::DateDescending => {
            images.sort_by(|img_a, img_b| img_b.create_date.cmp(&img_a.create_date))
        }
        ImageSort::None => {}
    }

    if images.len() > image_query.limit {
        images.drain(image_query.limit..);
    }

    Ok(Response(images))
}

fn get_user_id_from_headers(headers: &HeaderMap) -> Result<u64, PicOxError> {
    let user = headers.get("user").ok_or(PicOxError::NoUserInHeader)?;
    let user_str = user.to_str().unwrap();
    let user: u64 = user_str.parse().unwrap();

    Ok(user)
}

async fn add_image(
    State(context): State<PicContext>,
    headers: HeaderMap,
    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());
                let file_segment = field.bytes().await.unwrap_or(Bytes::new());
                data.extend_from_slice(file_segment.as_ref());
            }
        }
    }

    let user = headers.get("user").unwrap().to_str().unwrap();
    let user: u64 = user.parse().unwrap();

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

    let mut album = if let Some(album) = album {
        album
    } else {
        let name = metadata
            .clone()
            .unwrap()
            .album
            .album_name
            .ok_or(PicOxError::AlbumNotFound)?;
        let album = Album::new(&name, Vec::new(), user);
        context.db.insert::<Album>(album)?
    };

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

    info!(
        "Creating new image id={} at {:?}",
        img.id().unwrap(),
        img.storage_location
    );

    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<Vec<Album>>, PicOxError> {
    let resp = Album::find_album_by_query(&context.db, album_query.0);
    Ok(Response(resp))
}

async fn check_token_header(
    State(context): State<PicContext>,
    mut request: Request,
    next: Next,
) -> Result<impl IntoResponse, PicOxError> {
    let headers = request.headers();

    if let Some(token) = headers.get("token") {
        if let Some(api_key) = ApiKey::find_api_key_by_token(&context.db, token.to_str().unwrap())?
        {
            info!(
                "Authenticated user {}: '{}'",
                api_key.id().unwrap(),
                api_key.description
            );
            request
                .headers_mut()
                .insert("user", api_key.id().unwrap().into());
            return Ok(next.run(request).await);
        }
    }

    Err(PicOxError::TokenInvalid)
}

pub async fn run_picox(db: Database, config: PicOxConfig) {
    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);

    let app = Router::new()
        .route("/api/image/", post(add_image))
        .layer(DefaultBodyLimit::max(1024 * 1024 * 1024))
        .route("/api/album/create", post(create_album))
        .layer(middleware::from_fn_with_state(
            context.clone(),
            check_token_header,
        ))
        .route("/api/album/:id", get(get_album))
        .route("/api/album/", get(query_album))
        .route("/api/image/", get(query_images))
        .route("/api/image/:id", get(get_image))
        .with_state(context);

    let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap();
    info!("Serving at {}", config.host);
    axum::serve(listener, app).await.unwrap();
}