keep/src/modes/save.rs

use anyhow::{Context, Result, anyhow};
use gethostname::gethostname;
use is_terminal::IsTerminal;
use std::collections::HashMap;
use std::io::{self, Read, Write};

use clap::Command;
use clap::error::ErrorKind;
use log::debug;
use rusqlite::Connection;
use std::path::PathBuf;

use crate::compression_engine::get_compression_engine;
use crate::db::{self};
use crate::meta_plugin::{MetaPlugin, MetaPluginType, get_meta_plugin};
use crate::modes::common::{cmd_args_compression_type, cmd_args_digest_type, cmd_args_meta_plugin_types, get_meta_from_env, store_item_meta_value};
use chrono::Utc;

pub fn mode_save(
    cmd: &mut Command,
    args: &crate::Args,
    ids: &mut Vec<i64>,
    tags: &mut Vec<String>,
    conn: &mut Connection,
    data_path: PathBuf,
) -> Result<()> {
    if !ids.is_empty() {
        cmd.error(
            ErrorKind::InvalidValue,
            "ID given, you cannot supply IDs when using --save",
        )
        .exit();
    }

    if tags.is_empty() {
        tags.push("none".to_string());
    }

    let digest_type = cmd_args_digest_type(cmd, &args);
    debug!("MAIN: Digest type: {:?}", digest_type);
    
    let compression_type = cmd_args_compression_type(cmd, &args);
    debug!("MAIN: Compression type: {:?}", compression_type);
    let compression_engine =
        get_compression_engine(compression_type.clone()).expect("Unable to get compression engine");

    // Start with meta plugin types from command line
    let mut meta_plugin_types: Vec<MetaPluginType> = cmd_args_meta_plugin_types(cmd, &args);
    debug!("MAIN: Meta plugin types: {:?}", meta_plugin_types);
    
    // Convert digest type to meta plugin type and add to the list if needed
    let digest_meta_plugin_type = match digest_type {
        crate::meta_plugin::MetaPluginType::DigestSha256 => Some(MetaPluginType::DigestSha256),
        crate::meta_plugin::MetaPluginType::DigestMd5 => Some(MetaPluginType::DigestMd5),
        _ => None,
    };
    
    // Add digest meta plugin to the list if needed
    if let Some(digest_plugin_type) = digest_meta_plugin_type {
        if !meta_plugin_types.contains(&digest_plugin_type) {
            meta_plugin_types.push(digest_plugin_type);
        }
    }

    // Initialize meta_plugins with MetaPlugin instances for each MetaPluginType
    let mut meta_plugins: Vec<Box<dyn MetaPlugin>> = meta_plugin_types
        .iter()
        .map(|meta_plugin_type| get_meta_plugin(meta_plugin_type.clone()))
        .collect();
    
    // Check for unsupported meta plugins, warn the user, and remove them from the list
    let mut i = 0;
    meta_plugins.retain(|meta_plugin| {
        let is_supported = meta_plugin.is_supported();
        if !is_supported {
            // We need to get the meta name for the warning message
            // Since we can't mutably borrow meta_plugin here, we create a temporary one
            let meta_plugin_type = meta_plugin_types[i].clone();
            let mut temp_plugin = get_meta_plugin(meta_plugin_type);
            eprintln!("Warning: Meta plugin '{}' is enabled but not supported on this system", temp_plugin.meta_name());
        }
        i += 1;
        is_supported
    });

    let mut item = db::Item {
        id: None,
        ts: Utc::now(),
        size: None,
        compression: compression_type.to_string(),
    };

    let id = db::insert_item(conn, item.clone())?;
    item.id = Some(id);
    debug!("MAIN: Added item {:?}", item.clone());

    if !args.options.quiet {
        if std::io::stderr().is_terminal() {
            let mut t = term::stderr().unwrap();
            t.reset().unwrap_or(());
            t.attr(term::Attr::Bold).unwrap_or(());
            write!(t, "KEEP:").unwrap_or(());
            t.reset().unwrap_or(());
            write!(t, " New item ").unwrap_or(());
            t.attr(term::Attr::Bold).unwrap_or(());
            write!(t, "{id}")?;
            t.reset().unwrap_or(());
            write!(t, "  tags: ")?;
            t.attr(term::Attr::Bold).unwrap_or(());
            write!(t, "{}", tags.join(" "))?;
            t.reset().unwrap_or(());
            writeln!(t)?;
            std::io::stderr().flush()?;
        } else {
            let mut t = std::io::stderr();
            writeln!(t, "KEEP: New item: {}  tags: {:?}", id, tags)?;
        }
    }

    db::set_item_tags(conn, item.clone(), tags)?;

    let mut item_meta: HashMap<String, String> = get_meta_from_env();

    if let Ok(hostname) = gethostname().into_string() {
        if !item_meta.contains_key("hostname") {
            item_meta.insert("hostname".to_string(), hostname);
        }
    }

    for item in args.item.meta.iter() {
        let item = item.clone();
        item_meta.insert(item.key, item.value);
    }

    for kv in item_meta.iter() {
        let meta = db::Meta {
            id: item.id.unwrap(),
            name: kv.0.to_string(),
            value: kv.1.to_string(),
        };
        db::store_meta(conn, meta)?;
    }

    let item_id = item.id.ok_or_else(|| anyhow!("Item missing ID"))?;
    let mut item_path = data_path.clone();
    item_path.push(item_id.to_string());

    let mut stdin = io::stdin().lock();
    let mut stdout = io::stdout().lock();
    let mut buffer = [0; libc::BUFSIZ as usize];

    let mut item_out: Box<dyn Write> =
        compression_engine
            .create(item_path.clone())
            .context(anyhow!(
                "Unable to write file {:?} using compression {:?}",
                item_path,
                compression_type
            ))?;

    debug!("MAIN: Starting IO loop");
    loop {
        let n = stdin.read(&mut buffer[..libc::BUFSIZ as usize])?;
        item.size = match item.size {
            None => Some(n as i64),
            Some(prev_n) => Some(prev_n + n as i64),
        };

        if n == 0 {
            debug!("MAIN: EOF on STDIN");
            break;
        }

        debug!("MAIN: Loop - {:?} bytes", item.size);

        stdout.write_all(&buffer[..n])?;
        item_out.write_all(&buffer[..n])?;

        for meta_plugin in meta_plugins.iter_mut() {
            meta_plugin.update(&buffer[..n]);
        }
    }
    debug!("MAIN: Ending IO loop after {:?} bytes", item.size);

    stdout.flush()?;
    item_out.flush()?;

    for meta_plugin in meta_plugins.iter_mut() {
        let meta_name = meta_plugin.meta_name();

        match meta_plugin.finalize() {
            Ok(meta_value) => {
                if let Err(e) = store_item_meta_value(conn, item.clone(), meta_name.clone(), meta_value) {
                    eprintln!("Warning: Failed to store meta value for {}: {}", meta_name, e);
                }
            }
            Err(e) => {
                eprintln!("Warning: Failed to finalize meta plugin {}: {}", meta_name, e);
            }
        }
    }

    db::update_item(conn, item.clone())?;

    Ok(())
}