GeneIT/geneit_backend/src/controllers/data_controller.rs

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

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

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

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

/// Export whole family data
pub async fn export_family(f: FamilyInPath) -> HttpResult {
    let files_opt = FileOptions::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/{}.{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,
}

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

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

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

        req_member_data.to_member(member, true).await?;
    }

    // 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 couple_mapping = HashMap::new();
    for c in &mut couples_list {
        // Map wife and husband
        if let Some(i) = c.data.wife {
            c.data.wife = members_id_mapping.get(&i).copied();
        }

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

        let mut db_couple = couples_service::create(f.family_id()).await?;
        couple_mapping.insert(c.id, db_couple.id());

        c.data.clone().to_couple(&mut db_couple, true).await?;

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

    // Insert photos
    // TODO

    Ok(HttpResponse::Ok().body("go on"))
}

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