GeneIT/geneit_backend/src/controllers/data_controller.rs

use crate::connections::s3_connection;
use crate::constants;
use crate::controllers::HttpResult;
use crate::controllers::couples_controller::CoupleRequest;
use crate::controllers::members_controller::MemberRequest;
use crate::extractors::family_extractor::{FamilyInPath, FamilyInPathWithAdminMembership};
use crate::models::{CoupleID, MemberID, PhotoID};
use crate::services::photos_service::UploadedFile;
use crate::services::{couples_service, members_service, photos_service};
use actix_multipart::form::MultipartForm;
use actix_multipart::form::tempfile::TempFile;
use actix_web::HttpResponse;
use mime_guess::Mime;
use std::collections::HashMap;
use std::io;
use std::io::{Cursor, Read, Write};
use zip::write::SimpleFileOptions;
use zip::{CompressionMethod, ZipArchive};

const MEMBERS_FILE: &str = "members.json";
const COUPLES_FILE: &str = "couples.json";
const PHOTOS_DIR: &str = "photos/";

#[derive(serde::Deserialize)]
struct ImportMemberRequest {
    id: MemberID,
    photo_id: Option<PhotoID>,
    #[serde(flatten)]
    data: MemberRequest,
}

#[derive(serde::Deserialize)]
struct ImportCoupleRequest {
    photo_id: Option<PhotoID>,
    #[serde(flatten)]
    data: CoupleRequest,
}

/// Export whole family data
pub async fn export_family(f: FamilyInPath) -> HttpResult {
    let files_opt = SimpleFileOptions::default().compression_method(CompressionMethod::Bzip2);

    let members = members_service::get_all_of_family(f.family_id()).await?;
    let couples = couples_service::get_all_of_family(f.family_id()).await?;

    let buff = Vec::with_capacity(1000000);
    let mut zip_file = zip::ZipWriter::new(Cursor::new(buff));

    // Add main files
    zip_file.start_file(MEMBERS_FILE, files_opt)?;
    zip_file.write_all(serde_json::to_string(&members)?.as_bytes())?;

    zip_file.start_file(COUPLES_FILE, files_opt)?;
    zip_file.write_all(serde_json::to_string(&couples)?.as_bytes())?;

    // Add photos
    let mut photos = Vec::new();
    for member in &members {
        if let Some(id) = member.photo_id() {
            photos.push(id);
        }
    }

    for couple in &couples {
        if let Some(id) = couple.photo_id() {
            photos.push(id);
        }
    }

    for id in photos {
        let photo = photos_service::get_by_id(id).await?;
        let ext = photo.mime_extension().unwrap_or("bad");
        let file = s3_connection::get_file(&photo.photo_path()).await?;

        zip_file.start_file(format!("{PHOTOS_DIR}{}.{ext}", id.0), files_opt)?;
        zip_file.write_all(&file)?;
    }

    let buff = zip_file.finish()?.into_inner();

    Ok(HttpResponse::Ok()
        .content_type("application/zip")
        .body(buff))
}

#[derive(Debug, MultipartForm)]
pub struct UploadFamilyDataForm {
    #[multipart(rename = "archive")]
    archive: TempFile,
}

struct Photo {
    path: String,
    mimetype: Mime,
}

enum PhotoTarget {
    Member(MemberID),
    Couple(CoupleID),
}

struct PhotoToProcess {
    id: PhotoID,
    target: PhotoTarget,
}

/// Import whole family data
pub async fn import_family(
    f: FamilyInPathWithAdminMembership,
    MultipartForm(form): MultipartForm<UploadFamilyDataForm>,
) -> HttpResult {
    let mut zip = ZipArchive::new(form.archive.file)?;

    // Pre-process photos list
    let mut photos = HashMap::new();
    for file in zip.file_names() {
        let (id, ext) = match file.strip_prefix(PHOTOS_DIR).map(|f| f.split_once('.')) {
            Some(Some((id, ext))) => (id, ext),
            _ => continue,
        };

        photos.insert(
            PhotoID(id.parse()?),
            Photo {
                path: file.to_string(),
                mimetype: mime_guess::from_ext(ext).first_or_octet_stream(),
            },
        );
    }

    // Parse general information
    let members_list = serde_json::from_slice::<Vec<ImportMemberRequest>>(&read_zip_file(
        &mut zip,
        MEMBERS_FILE,
    )?)?;
    let mut couples_list = serde_json::from_slice::<Vec<ImportCoupleRequest>>(&read_zip_file(
        &mut zip,
        COUPLES_FILE,
    )?)?;

    // Delete all existing members
    members_service::delete_all_family(f.family_id()).await?;
    couples_service::delete_all_family(f.family_id()).await?;

    // Create empty members set
    let mut mapped_req_members = HashMap::new();

    let mut members_id_mapping = HashMap::new();
    let mut rev_members_id_mapping = HashMap::new();
    let mut new_members = Vec::with_capacity(members_list.len());

    let mut photos_to_insert = Vec::with_capacity(photos.len());

    for req_m in members_list {
        // Create member entry in database
        let new_m = members_service::create(f.family_id()).await?;

        // Map new member ID with request id
        members_id_mapping.insert(req_m.id, new_m.id());
        rev_members_id_mapping.insert(new_m.id(), req_m.id);

        // Save new member structure
        new_members.push(new_m);

        // Save request member information, mapped with its old id
        mapped_req_members.insert(req_m.id, req_m);
    }

    // Set member information, checking for eventual loops
    for member in &mut new_members {
        let db_id = member.id();
        let req_id = *rev_members_id_mapping.get(&db_id).unwrap();

        let req_member = mapped_req_members.get(&req_id).unwrap();

        // Map mother and father id and extract member information
        let mut req_member_data = req_member.data.clone();
        if let Some(i) = req_member_data.father {
            req_member_data.father = members_id_mapping.get(&i).copied();
        }
        if let Some(i) = req_member_data.mother {
            req_member_data.mother = members_id_mapping.get(&i).copied();
        }

        if let Err(e) = req_member_data.to_member(member).await {
            log::error!("Error while processing import (member {req_id:?}) - {e}");
            return Ok(
                HttpResponse::BadRequest().json(format!("Failed to validate member {req_id:?}!"))
            );
        }

        if let Some(id) = req_member.photo_id {
            photos_to_insert.push(PhotoToProcess {
                id,
                target: PhotoTarget::Member(db_id),
            })
        }
    }

    // Check for loops
    let members_slice = new_members.iter().collect::<Vec<_>>();
    if members_service::loop_detection::detect_loop(&members_slice) {
        return Ok(HttpResponse::BadRequest().body("Loop detected in members relationships!"));
    }

    // Save member information
    for member in &mut new_members {
        members_service::update(member).await?;
    }

    // Extract and insert couples information
    let mut new_couples = HashMap::new();
    for req_couple in &mut couples_list {
        // Map wife and husband
        if let Some(i) = req_couple.data.wife {
            req_couple.data.wife = members_id_mapping.get(&i).copied();
        }

        if let Some(i) = req_couple.data.husband {
            req_couple.data.husband = members_id_mapping.get(&i).copied();
        }

        let mut db_couple = couples_service::create(f.family_id()).await?;

        req_couple.data.clone().to_couple(&mut db_couple).await?;

        couples_service::update(&mut db_couple).await?;

        if let Some(id) = req_couple.photo_id {
            photos_to_insert.push(PhotoToProcess {
                id,
                target: PhotoTarget::Couple(db_couple.id()),
            })
        }

        new_couples.insert(db_couple.id(), db_couple);
    }

    // Insert member photos
    for photo_to_process in photos_to_insert {
        let photo = match photos.get(&photo_to_process.id) {
            None => continue,
            Some(photo) => photo,
        };

        let mut file = zip.by_name(&photo.path)?;
        if file.size() > constants::PHOTOS_MAX_SIZE as u64 {
            return Ok(HttpResponse::BadRequest().body("File is too large!"));
        }
        let mut buff = Vec::with_capacity(file.size() as usize);
        file.read_to_end(&mut buff)?;

        let photo = photos_service::finalize_upload(UploadedFile::from_memory(
            &buff,
            Some(photo.mimetype.clone()),
        )?)
        .await?;

        // Update appropriate database entry
        match photo_to_process.target {
            PhotoTarget::Member(member_id) => {
                let member = new_members
                    .iter_mut()
                    .find(|m| m.id().eq(&member_id))
                    .unwrap();
                member.set_photo_id(Some(photo.id()));
                members_service::update(member).await?;
            }
            PhotoTarget::Couple(couple_id) => {
                let couple = new_couples.get_mut(&couple_id).unwrap();
                couple.set_photo_id(Some(photo.id()));
                couples_service::update(couple).await?;
            }
        }
    }

    Ok(HttpResponse::Accepted().finish())
}

fn read_zip_file<R: Read + io::Seek>(
    archive: &mut ZipArchive<R>,
    file: &str,
) -> anyhow::Result<Vec<u8>> {
    let mut entry = archive.by_name(file)?;
    assert!(entry.size() < constants::PHOTOS_MAX_SIZE as u64);
    let mut buff = Vec::with_capacity(entry.size() as usize);
    entry.read_to_end(&mut buff)?;

    Ok(buff)
}