fileserver/backend/src/routes/fs/routes.rs

use std::{
    collections::{BTreeSet, HashMap},
    fs::File,
    io::{Read, Write},
    num::NonZeroU32,
    ops::DerefMut,
    sync::atomic::Ordering
};

use fast_image_resize as fir;
use image::DynamicImage;
use rustracing_jaeger::Span;
use tiny_http::{Request, Response, ResponseBox, StatusCode};

use crate::{
    db,
    db::DBConnection,
    dto,
    header,
    metrics,
    routes::{filters::UserInfo, fs::resize_dimensions, get_reply, AppError, ChannelReader}
};

pub fn root(_: &Span, _: &mut Request, _: &mut DBConnection, info: UserInfo) -> Result<ResponseBox, AppError> {
    get_reply(&dto::responses::Root {
        statusCode: 200,
        rootId: info.0.root_id
    })
}

pub fn node(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    node: i32
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.read(span);
    let node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?;

    get_reply(&dto::responses::GetNode {
        statusCode: 200,
        id: node.id,
        name: node.name,
        isFile: node.is_file,
        preview: node.has_preview,
        parent: node.parent_id,
        size: node.size,
        children: (!node.is_file).then(|| {
            db.get_children(span, node.id)
                .iter()
                .map(|child| dto::responses::GetNodeEntry {
                    id: child.id,
                    name: child.name.clone(),
                    isFile: child.is_file,
                    preview: child.has_preview,
                    parent: child.parent_id,
                    size: child.size
                })
                .collect()
        })
    })
}

pub fn path(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    node: i32
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.read(span);
    let node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?;

    get_reply(&super::generate_path_dto(span, &node, db))
}

pub fn create_node(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    data: dto::requests::CreateNode,
    file: bool
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.read(span);
    let node = super::create_node(span, data.name, &info.0, file, Some(data.parent), false, db);

    match node {
        Ok(v) => get_reply(&dto::responses::NewNode {
            statusCode: 200,
            id: v.id
        }),
        Err(v) => match v {
            super::CreateNodeResult::InvalidName => AppError::BadRequest("Invalid name").err(),
            super::CreateNodeResult::InvalidParent => AppError::BadRequest("Invalid parent").err(),
            super::CreateNodeResult::Exists(file, id) => get_reply(&dto::responses::NodeExists {
                statusCode: 200,
                id,
                exists: true,
                isFile: file
            })
        }
    }
}

pub fn delete_node(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    node: i32,
    pool: &db::DBPool
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let inner_guard_lock = guard_lock.clone();
    let _guard = guard_lock.read(span);

    let node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?;

    if node.parent_id.is_none() {
        return AppError::BadRequest("Can't delete root").err();
    }

    let (tx, rx) = std::sync::mpsc::channel::<String>();
    let inner_pool = pool.clone();

    tx.send("Waiting in queue\n".to_owned()).unwrap();

    let inner_span = metrics::span("DELETE_POOL - queueing", span);
    super::DELETE_POOL.spawn(move || {
        let del_span = metrics::span("DELETE_POOL - deleting", &inner_span);
        drop(inner_span);
        let mut db = &mut DBConnection::from(inner_pool.get().unwrap());
        let guard_lock = inner_guard_lock.clone();
        let _guard = guard_lock.write(&del_span);
        super::delete_node(&del_span, &node, &tx, db.deref_mut());
    });

    Ok(Response::new(
        StatusCode::from(200),
        vec![header("content-type", "text/plain; charset=utf-8")],
        ChannelReader(rx),
        None,
        None
    )
    .boxed())
}

pub fn upload(
    span: &Span,
    req: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    node: i32
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.read(span);
    let mut node = super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?;

    if !node.is_file {
        return AppError::BadRequest("Can't upload to a directory node").err();
    }

    let mut file_size = 0_i64;
    let file_name = format!("./files/{}", node.id);
    let mut file_buf = Vec::<u8>::new();
    {
        let _span = metrics::span("receive_file", span);
        let mut buf = vec![0_u8; 8 * 1024 * 1024];
        let mut file = File::create(file_name.clone()).unwrap();

        let reader = req.as_reader();

        loop {
            let r = reader.read(&mut buf).unwrap();
            if r == 0 {
                break;
            }
            if file_size < 20 * 1024 * 1024 {
                file_buf.write_all(&buf[..r]).unwrap();
            }
            file.write_all(&buf[..r]).unwrap();
            file_size += r as i64;
        }
    }

    metrics::DISK_USAGE
        .with_label_values(&[node.owner_id.to_string().as_str()])
        .add(file_size - node.size.unwrap_or(0));
    {
        let prev_span = metrics::span("generate_preview", span);
        node.has_preview = (|| {
            if file_size > file_buf.len() as i64 {
                return None;
            }
            let mime = mime_guess::from_path(std::path::Path::new(&node.name)).first()?.to_string();
            let img = {
                let _span = metrics::span("generate_preview_load", &prev_span);
                image::load_from_memory_with_format(
                    file_buf.as_slice(),
                    image::ImageFormat::from_mime_type(mime)?
                )
                .ok()?
            };
            let img = {
                let _span = metrics::span("generate_preview_convert", &prev_span);
                let width = NonZeroU32::try_from(img.width()).unwrap();
                let height = NonZeroU32::try_from(img.height()).unwrap();
                match img {
                    DynamicImage::ImageLuma8(v) =>
                        fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8),
                    DynamicImage::ImageLumaA8(v) =>
                        fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8x2),
                    DynamicImage::ImageRgb8(v) =>
                        fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8x3),
                    DynamicImage::ImageRgba8(v) =>
                        fir::Image::from_vec_u8(width, height, v.into_raw(), fir::PixelType::U8x4),
                    DynamicImage::ImageLuma16(_) =>
                        fir::Image::from_vec_u8(width, height, img.to_luma8().into_raw(), fir::PixelType::U8),
                    DynamicImage::ImageLumaA16(_) => fir::Image::from_vec_u8(
                        width,
                        height,
                        img.to_luma_alpha8().into_raw(),
                        fir::PixelType::U8x2
                    ),
                    DynamicImage::ImageRgb16(_) => fir::Image::from_vec_u8(
                        width,
                        height,
                        img.to_rgb8().into_raw(),
                        fir::PixelType::U8x3
                    ),
                    DynamicImage::ImageRgba16(_) => fir::Image::from_vec_u8(
                        width,
                        height,
                        img.to_rgba8().into_raw(),
                        fir::PixelType::U8x4
                    ),
                    DynamicImage::ImageRgb32F(_) => fir::Image::from_vec_u8(
                        width,
                        height,
                        img.to_rgb8().into_raw(),
                        fir::PixelType::U8x3
                    ),
                    DynamicImage::ImageRgba32F(_) => fir::Image::from_vec_u8(
                        width,
                        height,
                        img.to_rgba8().into_raw(),
                        fir::PixelType::U8x4
                    ),
                    _ => fir::Image::from_vec_u8(
                        width,
                        height,
                        img.to_rgba8().into_raw(),
                        fir::PixelType::U8x4
                    )
                }
                .expect("Failed to convert preview image")
            };
            let img = {
                let _span = metrics::span("generate_preview_resize", &prev_span);
                let new_dim = resize_dimensions(img.width().get(), img.height().get());
                let mut dst = fir::Image::new(
                    NonZeroU32::try_from(new_dim.0).unwrap(),
                    NonZeroU32::try_from(new_dim.1).unwrap(),
                    img.pixel_type()
                );
                fir::Resizer::new(fir::ResizeAlg::SuperSampling(fir::FilterType::Hamming, 2))
                    .resize(&img.view(), &mut dst.view_mut())
                    .expect("Failed to resize preview image");
                dst
            };
            let _span = metrics::span("generate_preview_save", &prev_span);
            let mut file = std::io::BufWriter::new(
                File::create(std::path::Path::new(&(file_name + "_preview.jpg")))
                    .expect("Failed to open preview image file")
            );
            image::codecs::jpeg::JpegEncoder::new(&mut file)
                .encode(
                    img.buffer(),
                    img.width().get(),
                    img.height().get(),
                    match img.pixel_type() {
                        fir::PixelType::U8 => image::ColorType::L8,
                        fir::PixelType::U8x2 => image::ColorType::La8,
                        fir::PixelType::U8x3 => image::ColorType::Rgb8,
                        fir::PixelType::U8x4 => image::ColorType::Rgba8,
                        _ => unreachable!()
                    }
                )
                .expect("Failed to save preview image");
            Some(())
        })()
        .is_some();
    }

    node.size = Some(file_size);
    db.save_node(span, &node);

    get_reply(&dto::responses::Success { statusCode: 200 })
}

pub fn create_zip(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    data: dto::requests::CreateZip,
    pool: &db::DBPool
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let inner_guard_lock = guard_lock.clone();
    let _guard = guard_lock.read(span);
    let mut nodes: Vec<db::Inode> = Vec::new();
    for node in data.nodes.clone() {
        nodes.push(super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?);
    }
    let zip_nodes = BTreeSet::from_iter(data.nodes.iter().copied());

    {
        let guard = super::ZIP_TO_PROGRESS.read();
        if let Some(entry) = guard.get(&zip_nodes) {
            return get_reply(&dto::responses::CreateZipDone {
                statusCode: 200,
                done: entry.done.load(Ordering::Relaxed),
                progress: Some(entry.progress.load(Ordering::Relaxed)),
                total: Some(entry.total.load(Ordering::Relaxed))
            });
        }
    }
    let entry = {
        let mut guard = super::ZIP_TO_PROGRESS.write();
        guard.insert(
            zip_nodes.clone(),
            std::sync::Arc::from(super::ZipProgressEntry {
                temp_id: super::NEXT_TEMP_ID.fetch_add(1, Ordering::Relaxed),
                done: std::sync::atomic::AtomicBool::new(false),
                progress: std::sync::atomic::AtomicU64::new(0),
                total: std::sync::atomic::AtomicU64::new(1),
                delete_after: std::sync::atomic::AtomicI64::new(0)
            })
        );
        guard.get(&zip_nodes).unwrap().clone()
    };

    let inner_pool = pool.clone();

    let inner_span = metrics::span("ZIP_POOL - queueing", span);
    super::ZIP_POOL.spawn(move || {
        let mut zip_span = metrics::span("ZIP_POOL - zipping", &inner_span);
        drop(inner_span);
        let db = &mut DBConnection::from(inner_pool.get().unwrap());
        type NodeMap = HashMap<i32, db::Inode>;

        super::cleanup_temp_zips(&zip_span);

        let _guard = inner_guard_lock.read(&zip_span);

        fn get_path(node: &db::Inode, dirs: &NodeMap) -> String {
            let mut path = node.name.clone();
            let mut _node = dirs.get(&node.parent_id.unwrap_or(-1));
            while let Some(node) = _node {
                path.insert_str(0, &(node.name.clone() + "/"));
                _node = dirs.get(&node.parent_id.unwrap_or(-1));
            }
            path
        }
        {
            let i_span = metrics::span("ZIP_POOL - calc total", &zip_span);
            nodes.iter().for_each(|node| {
                entry.total.fetch_add(
                    super::get_total_size(&i_span, node.clone(), db),
                    Ordering::Relaxed
                );
            });
            entry.total.fetch_sub(1, Ordering::Relaxed);
        }
        {
            let comp_span = metrics::span("ZIP_POOL - compressing total", &zip_span);
            let mut buf = vec![0_u8; 16 * 1024 * 1024];
            let file = File::create(format!("./temp/{}", entry.temp_id)).expect("Failed to create temp file");
            let mut zip = zip::ZipWriter::new(file);
            let zip_options = zip::write::FileOptions::default().large_file(true);
            let (files, dirs): (NodeMap, NodeMap) = {
                let _span = metrics::span("ZIP_POOL - gathering nodes", &comp_span);
                nodes
                    .iter()
                    .flat_map(|node| super::get_nodes_recursive(&comp_span, node.clone(), db))
                    .map(|node| (node.id, node))
                    .partition(|v| v.1.is_file)
            };
            zip_span.set_tags(|| {
                vec![
                    rustracing::tag::Tag::new("files", files.len() as i64),
                    rustracing::tag::Tag::new("dirs", dirs.len() as i64),
                ]
            });
            {
                let _span = metrics::span("ZIP_POOL - dirs", &comp_span);
                dirs.values().for_each(|dir| {
                    zip.add_directory(get_path(dir, &dirs), zip_options).expect("Failed to add dir to zip");
                });
            }
            {
                let _span = metrics::span("ZIP_POOL - files", &comp_span);
                files.values().for_each(|node| {
                    zip.start_file(get_path(node, &dirs), zip_options).expect("Failed to start zip file");
                    let mut file = File::open(format!("./files/{}", node.id)).expect("Failed to open file for zip");

                    loop {
                        let count = file.read(&mut buf).expect("Failed to read file for zip");
                        if count == 0 {
                            break;
                        }
                        zip.write_all(&buf[..count]).expect("Failed to write zip");
                        entry.progress.fetch_add(count as u64, Ordering::Relaxed);
                    }
                });
            }

            zip.finish().expect("Failed to finish zip");
        }
        entry.done.store(true, Ordering::Relaxed);
        entry.delete_after.store(chrono::Utc::now().timestamp() + 10 * 60, Ordering::Relaxed);
    });
    get_reply(&dto::responses::CreateZipDone {
        statusCode: 200,
        done: false,
        progress: Some(0),
        total: Some(1)
    })
}

pub fn download(span: &Span, req: &mut Request, db: &mut DBConnection) -> Result<ResponseBox, AppError> {
    let data: dto::requests::Download =
        serde_urlencoded::from_reader(req.as_reader()).map_err(|_| AppError::BadRequest("Invalid form data"))?;
    let info = crate::routes::filters::authorize_jwt(span, &data.jwtToken, db)?;
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.read(span);
    let node: db::Inode =
        super::get_node_and_validate(span, &info.0, data.id, db).ok_or(AppError::BadRequest("Unknown node"))?;

    if node.is_file {
        let file_name = format!("./files/{}", node.id);
        let resp = Response::from_file(File::open(std::path::Path::new(&file_name)).unwrap())
            .with_header(header("content-type", "application/octet-stream"))
            .with_header(header(
                "content-disposition",
                &("attachment; filename=".to_owned() + &node.name)
            ))
            .boxed();
        Ok(resp)
    } else {
        let nodes_key = BTreeSet::from([node.id]);
        let guard = super::ZIP_TO_PROGRESS.read();
        let entry = guard.get(&nodes_key).ok_or(AppError::BadRequest("Unknown node"))?;
        if !entry.done.load(Ordering::Relaxed) {
            AppError::BadRequest("Unknown node").err()
        } else {
            let file_name = format!("./temp/{}", entry.temp_id);
            let resp = Response::from_file(File::open(std::path::Path::new(&file_name)).unwrap())
                .with_header(header("content-type", "application/zip"))
                .with_header(header(
                    "content-disposition",
                    &("attachment; filename=".to_owned() + &node.name + ".zip")
                ))
                .boxed();
            Ok(resp)
        }
    }
}

pub fn download_multi(span: &Span, req: &mut Request, db: &mut DBConnection) -> Result<ResponseBox, AppError> {
    let data: dto::requests::DownloadMulti =
        serde_urlencoded::from_reader(req.as_reader()).map_err(|_| AppError::BadRequest("Invalid form data"))?;
    let info = crate::routes::filters::authorize_jwt(span, &data.jwtToken, db)?;
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.read(span);
    let nodes: Vec<db::Inode> = {
        let _span = metrics::span("parsing_nodes", span);
        data.id
            .split(',')
            .map(|v| {
                v.parse::<i32>().map_err(|_| AppError::BadRequest("Failed to parse")).map(|n| {
                    super::get_node_and_validate(span, &info.0, n, db).ok_or(AppError::BadRequest("Unknown node"))
                })
            })
            .into_iter()
            .collect::<Result<Result<Vec<db::Inode>, AppError>, AppError>>()??
    };

    let nodes_key = BTreeSet::from_iter(nodes.iter().map(|node| node.id));
    let guard = super::ZIP_TO_PROGRESS.read();
    let entry = guard.get(&nodes_key).ok_or(AppError::BadRequest("Unknown zip"))?;
    if !entry.done.load(Ordering::Relaxed) {
        AppError::BadRequest("Unfinished zip").err()
    } else {
        let file_name = format!("./temp/{}", entry.temp_id);
        let resp = Response::from_file(File::open(std::path::Path::new(&file_name)).unwrap())
            .with_header(header("content-type", "application/zip"))
            .with_header(header(
                "content-disposition",
                "attachment; filename=files.zip"
            ))
            .boxed();
        Ok(resp)
    }
}

pub fn download_preview(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    node: i32
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.read(span);
    let node: db::Inode =
        super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?;

    if node.has_preview {
        let file = format!("./files/{}_preview.jpg", node.id);
        get_reply(&dto::responses::DownloadBase64 {
            statusCode: 200,
            data: "data:image/png;base64,".to_owned()
                + &base64::encode(std::fs::read(std::path::Path::new(&file)).unwrap())
        })
    } else {
        AppError::BadRequest("No preview").err()
    }
}

pub fn get_type(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    node: i32
) -> Result<ResponseBox, AppError> {
    let node: db::Inode =
        super::get_node_and_validate(span, &info.0, node, db).ok_or(AppError::BadRequest("Unknown node"))?;

    get_reply(&dto::responses::Type {
        statusCode: 200,
        _type: mime_guess::from_path(std::path::Path::new(&node.name)).first_or_octet_stream().to_string()
    })
}

pub fn move_node(
    span: &Span,
    _: &mut Request,
    db: &mut DBConnection,
    info: UserInfo,
    data: dto::requests::MoveNode
) -> Result<ResponseBox, AppError> {
    let guard_lock = DBConnection::get_lock(info.0.id);
    let _guard = guard_lock.write(span);

    let target =
        super::get_node_and_validate(span, &info.0, data.target, db).ok_or(AppError::BadRequest("Invalid target"))?;
    if target.is_file {
        return AppError::BadRequest("Invalid target").err();
    }

    let mut nodes = data
        .nodes
        .iter()
        .map(|v| super::get_node_and_validate(span, &info.0, *v, db))
        .into_iter()
        .collect::<Option<Vec<db::Inode>>>()
        .ok_or(AppError::BadRequest("Invalid node"))?;

    for parent in super::get_node_path(span, target.clone(), db) {
        if nodes.contains(&parent) {
            return AppError::BadRequest("Can't move node into one of it's subfolders").err();
        }
    }
    for child in db.get_children(span, target.id) {
        if nodes.iter().any(|n| n.name == child.name) {
            return AppError::BadRequest("Can't overwrite existing file").err();
        }
    }

    nodes.iter_mut().for_each(|n| db.move_node(span, n, target.id));

    get_reply(&dto::responses::Success { statusCode: 200 })
}