//! Encrypt workflow

use std::collections::HashSet;
use std::env;
use std::fs::{self, File};
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use chrono::{DateTime, Utc};
use flate2::{write::GzEncoder, Compression};
use sequoia_openpgp::{
    self as openpgp,
    crypto::{KeyPair, Password},
    policy::StandardPolicy,
    serialize::stream::{Encryptor, LiteralWriter, Message, Recipient, Signer},
};
use walkdir::WalkDir;
use zip::{write::FileOptions, CompressionMethod, ZipWriter};

use crate::checksum::{generate_checksums_file_content, Sha256ChecksumWriter};
use crate::crypt;
use crate::dpkg;
use crate::error::Error;
use crate::filesystem::{check_space, get_common_path};
use crate::utils::{get_max_cpu, to_human_readable_size, Progress, ProgressReader};

/// Run encryption using information provided in [`EncryptOpts`]
pub trait Encrypt {
    fn encrypt<P: AsRef<Path>>(
        &self,
        opts: &EncryptOpts<P, impl Fn(f32) -> Result<(), Error>>,
    ) -> Result<Option<String>, Error>;
}

/// Destination type for data packages
pub enum Destination<'a, Cb: Fn() -> Result<String, Error>> {
    Local,
    Sftp(crate::sftp::Remote<'a, Cb>),
}

impl<Cb: Fn() -> Result<String, Error>> Encrypt for Destination<'_, Cb> {
    fn encrypt<P: AsRef<Path>>(
        &self,
        opts: &EncryptOpts<P, impl Fn(f32) -> Result<(), Error>>,
    ) -> Result<Option<String>, Error> {
        log::debug!("Extract signing and encryption certificates");
        let policy = StandardPolicy::new();
        let signer_key = opts
            .sender
            .map(|cert| crypt::get_signing_cert(cert, opts.password.as_ref(), &policy))
            .transpose()?;
        let encryption_keys = crypt::get_encryption_keys(opts.recipients, &policy)?;
        match self {
            Self::Local => encrypt_local(opts, signer_key, &encryption_keys),
            Self::Sftp(remote) => encrypt_sftp(opts, signer_key, &encryption_keys, remote),
        }
    }
}

/// Options required by the encrypt workflow.
pub struct EncryptOpts<'a, P: AsRef<Path>, Cb: Fn(f32) -> Result<(), Error>> {
    /// Input files for encryption.
    pub files: &'a [P],
    /// Public PGP keys of recipients.
    pub recipients: &'a [&'a str],
    /// Private PGP key for sender (used for signing data).
    pub sender: Option<&'a str>,
    /// Password for decrypting the sender's key.
    pub password: Option<Password>,
    /// Output path for the encrypted data package.
    pub output: Option<&'a str>,
    /// Compression level from 0 (no compression, fastest) to 9 (highest
    /// compression, slowest).
    pub compression_level: Option<u32>,
    /// Check the input, don't run encryption.
    pub dry_run: bool,
    /// Ignore warnings (e.g. not sufficient storage) and proceed with
    /// encryption.
    pub force: bool,
    /// The maximum number of CPU cores used by the workflow (minimum: 2).
    pub max_cpu: Option<usize>,
    /// Report encryption progress using this callback.
    pub progress: Option<Cb>,
    /// Purpose, see [`dpkg::Purpose`] for more details.
    pub purpose: Option<dpkg::Purpose>,
    /// Data transfer ID, see [`dpkg::PkgMetadata::transfer_id`] for more details.
    pub transfer_id: Option<u32>,
}

/// Creates a data package with encrypted and optionally compressed data along
/// with checksums and metadata.
///
/// The data package can be stored in either the local filesystem or on a remote
/// server.
pub fn encrypt<P: AsRef<Path>, T: Encrypt>(
    opts: &EncryptOpts<P, impl Fn(f32) -> Result<(), Error>>,
    dest: &T,
) -> Result<Option<String>, Error> {
    dest.encrypt(opts)
}

/// Generates an output path for the data package.
fn get_output_path<P: AsRef<Path>>(
    output: Option<P>,
    timestamp: Option<DateTime<Utc>>,
) -> Result<PathBuf, Error> {
    let default_name = dpkg::generate_package_name(timestamp);
    let (output_dir, output_name) = match output {
        Some(p) => {
            let p = PathBuf::from(p.as_ref());
            if !p.is_absolute() {
                return Err(Error(format!(
                    "Output path must be absolute (provided path {:?})",
                    p
                )));
            }
            if p.is_dir() {
                (p, default_name)
            } else {
                let d = p
                    .parent()
                    .ok_or("Unable to find directory for the given output")?
                    .to_path_buf();
                let mut f = p
                    .file_name()
                    .ok_or("Unable to find file name in the given output")?
                    .to_string_lossy()
                    .to_string();
                if !d.exists() {
                    return Err(Error(format!("Output directory does not exist {:?})", d)));
                }
                if !f.ends_with(".zip") {
                    f.push_str(".zip");
                }
                (d, f)
            }
        }
        None => (env::current_dir()?, default_name),
    };
    // TODO: check if output directory is writable
    Ok(output_dir.join(output_name))
}

/// Returns paths where the common prefix (to all paths) is stripped.
fn get_archive_paths(files: &[PathBuf]) -> Result<Vec<PathBuf>, Error> {
    let parent_folders = files
        .iter()
        .map(|f| {
            f.parent()
                .ok_or(format!("Unable to find parent folder of {:?}", f))
        })
        .collect::<Result<Vec<_>, _>>()?;
    let root_dir = get_common_path(&parent_folders)?;
    let archive_paths = files
        .iter()
        .map(|f| Ok(Path::new(dpkg::CONTENT_FOLDER).join(f.strip_prefix(&root_dir)?)))
        .collect::<Result<Vec<PathBuf>, Error>>()?;
    Ok(archive_paths)
}

/// Initializes an returns a PGP message, which can be used for writing data into it.
fn init_message<'a, W: io::Write + Send + Sync, S: openpgp::crypto::Signer + Send + Sync + 'a>(
    writer: &'a mut W,
    encryption_keys: &'a [openpgp::packet::Key<
        openpgp::packet::key::PublicParts,
        openpgp::packet::key::UnspecifiedRole,
    >],
    signer: Option<S>,
) -> Result<Message<'a>, Error> {
    let message = Message::new(writer);
    let message =
        Encryptor::for_recipients(message, encryption_keys.iter().map(Recipient::from)).build()?;
    let message = match signer {
        Some(x) => Signer::new(message, x).build()?,
        None => message,
    };
    let message = LiteralWriter::new(message).build()?;

    Ok(message)
}

/// Adds a file to the inner tar archive containing checksums of the individual data files.
fn add_checksum_file<W: io::Write>(
    archive: &mut tar::Builder<W>,
    content: &[u8],
) -> Result<(), Error> {
    let mut header = tar::Header::new_gnu();
    header.set_entry_type(tar::EntryType::file());
    header.set_size(content.len().try_into()?);
    header.set_mtime(Utc::now().timestamp().try_into()?);
    header.set_mode(0o644);
    header.set_cksum();
    archive.append_data(&mut header, dpkg::CHECKSUM_FILE, content)?;
    Ok(())
}

/// Writes files to the inner tar file with optional compression.
fn add_files_to_archive<W: io::Write, P: AsRef<Path>, Cb: Fn(f32) -> Result<(), Error>>(
    writer: &mut W,
    file_path: &[(P, P)],
    max_cpu: usize,
    compression_level: Option<u32>,
    progress: &mut Option<Progress<Cb>>,
) -> Result<(), Error> {
    let enc = GzEncoder::new(
        writer,
        Compression::new(compression_level.unwrap_or_else(|| Compression::default().level())),
    );
    let mut archive = tar::Builder::new(enc);
    if let Some(pg) = progress {
        for (f, p) in file_path {
            let f = fs::File::open(f)?;
            let mut header = tar::Header::new_gnu();
            let metadata = f.metadata()?;
            header.set_metadata(&metadata);
            header.set_cksum();
            let reader = ProgressReader::new(f, |p| pg.update(p));
            archive.append_data(&mut header, p, reader)?;
        }
    } else {
        for (f, p) in file_path {
            archive.append_path_with_name(f, p)?;
        }
    }
    add_checksum_file(
        &mut archive,
        generate_checksums_file_content(file_path, max_cpu)?.as_bytes(),
    )?;
    archive.finish()?;
    Ok(())
}

/// Adds the metadata file to the data package (outer zip file).
fn add_metadata_file<W: io::Write + io::Seek>(
    archive: &mut ZipWriter<W>,
    metadata: &dpkg::PkgMetadata,
    signer_key: Option<KeyPair>,
) -> Result<(), Error> {
    let options = FileOptions::default().compression_method(CompressionMethod::Stored);
    archive.start_file(dpkg::METADATA_FILE, options)?;
    let metadata_json = serde_json::to_string(&metadata)?.as_bytes().to_owned();
    archive.write_all(&metadata_json)?;
    if let Some(key) = signer_key {
        archive.start_file(dpkg::METADATA_SIG_FILE, options)?;
        archive.write_all(&crypt::sign_detached(&metadata_json, key, false)?)?;
    }
    Ok(())
}

/// Processes input data files.
///
/// Returns:
///
/// - A vector pairs where the first element is the absolute path to the file
///   and the second element is the path inside the data package, i.e. a path
///   where the common prefix (for all files) is stripped.
/// - The combined size of all files.
///
/// Processing steps:
///
/// - Walk directories and return individual files.
/// - Normalize paths (make them absolute).
/// - Remove duplicated paths.
/// - Calculate the combined size of the input.
fn process_files<P: AsRef<Path>>(files: &[P]) -> Result<(Vec<(PathBuf, PathBuf)>, u64), Error> {
    let mut total_input_size = 0u64;
    let mut unique_files = HashSet::new();
    for entry in files.iter().map(WalkDir::new).into_iter().flatten() {
        let entry = entry?;
        if entry.file_type().is_file() {
            total_input_size += entry.metadata()?.len();
            unique_files.insert(entry.path().canonicalize()?.to_owned());
        }
    }
    let files = Vec::from_iter(unique_files);
    let archive_paths = get_archive_paths(&files)?;
    let file_archive_path = files
        .into_iter()
        .zip(archive_paths.into_iter())
        .collect::<Vec<(_, _)>>();
    Ok((file_archive_path, total_input_size))
}

/// Runs the compression and encryption stages to a given writer.
///
/// The end result of running those stages is the creation of the final data
/// package.
fn encrypt_to_writer<'a, P: AsRef<Path>, W: io::Write + io::Seek + Sync + Send>(
    output: W,
    opts: &EncryptOpts<P, impl Fn(f32) -> Result<(), Error>>,
    signer_key: Option<KeyPair>,
    encryption_keys: &'a [openpgp::packet::Key<
        openpgp::packet::key::PublicParts,
        openpgp::packet::key::UnspecifiedRole,
    >],
    file_path: &'a [(PathBuf, PathBuf)],
    total_input_size: u64,
    timestamp: DateTime<Utc>,
) -> Result<(), Error> {
    log::info!("Compress and encrypt input data [this can take a while]");
    let mut zip = ZipWriter::new(output);
    zip.start_file(
        dpkg::DATA_FILE_ENCRYPTED,
        FileOptions::default().compression_method(CompressionMethod::Stored),
    )?;
    let mut checksum_writer = Sha256ChecksumWriter::new(&mut zip);
    let mut message = init_message(&mut checksum_writer, encryption_keys, signer_key.clone())?;
    let mut progress = opts
        .progress
        .as_ref()
        .map(|cb| Progress::new(total_input_size as usize, cb));
    add_files_to_archive(
        &mut message,
        file_path,
        opts.max_cpu.unwrap_or_else(get_max_cpu),
        opts.compression_level,
        &mut progress,
    )?;
    message.finalize()?;
    let data_checksum = checksum_writer.get_hash();
    add_metadata_file(
        &mut zip,
        &dpkg::PkgMetadata {
            sender: signer_key
                .as_ref()
                .map_or_else(|| "".into(), |fp| fp.public().fingerprint().to_hex()),
            recipients: encryption_keys
                .iter()
                .map(|k| k.fingerprint().to_hex())
                .collect(),
            checksum: data_checksum,
            timestamp: timestamp.format(dpkg::METADATA_DATETIME_FORMAT).to_string(),
            version: dpkg::default_version(),
            checksum_algorithm: Default::default(),
            compression_algorithm: Default::default(),
            transfer_id: opts.transfer_id,
            purpose: opts.purpose,
        },
        signer_key,
    )?;
    zip.finish()?;
    Ok(())
}

fn show_summary(location: &str, total_size: u64, package_size: u64) {
    log::info!(
        "Completed data encryption: {} (source: {}, output: {}, ratio: {:.2})",
        location,
        to_human_readable_size(total_size),
        to_human_readable_size(package_size),
        total_size as f32 / package_size as f32
    );
}
fn encrypt_local<'a, P: AsRef<Path>>(
    opts: &EncryptOpts<P, impl Fn(f32) -> Result<(), Error>>,
    signer_key: Option<KeyPair>,
    encryption_keys: &'a [openpgp::packet::Key<
        openpgp::packet::key::PublicParts,
        openpgp::packet::key::UnspecifiedRole,
    >],
) -> Result<Option<String>, Error> {
    let timestamp = Utc::now();
    log::info!("Verify files to encrypt");
    let (file_path, total_input_size) = process_files(opts.files)?;
    let output_path = get_output_path(opts.output, Some(timestamp))?;
    log::info!("Verify available disk space");
    check_space(
        total_input_size,
        output_path
            .parent()
            .ok_or("Unable to find output directory")?,
        opts.force,
    )?;
    if opts.dry_run {
        log::info!("Dry run completed successfully");
        return Ok(None);
    }
    let output = File::create(&output_path)?;
    encrypt_to_writer(
        output,
        opts,
        signer_key,
        encryption_keys,
        &file_path,
        total_input_size,
        timestamp,
    )?;
    let package_location = output_path.to_string_lossy().into_owned();
    show_summary(
        &package_location,
        total_input_size,
        fs::metadata(output_path)?.len(),
    );
    Ok(Some(package_location))
}

fn encrypt_sftp<'a, P: AsRef<Path>, Cb: Fn() -> Result<String, Error>>(
    opts: &EncryptOpts<P, impl Fn(f32) -> Result<(), Error>>,
    signer_key: Option<KeyPair>,
    encryption_keys: &'a [openpgp::packet::Key<
        openpgp::packet::key::PublicParts,
        openpgp::packet::key::UnspecifiedRole,
    >],
    remote: &crate::sftp::Remote<'a, Cb>,
) -> Result<Option<String>, Error> {
    let timestamp = Utc::now();
    log::info!("Verify files to encrypt");
    let (file_path, total_input_size) = process_files(opts.files)?;
    let sftp = remote.connect()?;
    if opts.dry_run {
        log::info!("Dry run completed successfully");
        return Ok(None);
    }
    let upload_dir = crate::sftp::UploadDir::new(remote.path, remote.envelope_dir, &sftp);
    upload_dir.create(None)?;
    let package_name = opts.output.map_or_else(
        || dpkg::generate_package_name(Some(timestamp)),
        |f| {
            PathBuf::from(f)
                .file_name()
                .expect("This should be a file name")
                .to_string_lossy()
                .into_owned()
        },
    );
    let dpkg_path = crate::sftp::DpkgPath::new(&upload_dir.path, &package_name, &sftp);
    let output = sftp.create(&dpkg_path.tmp)?;
    encrypt_to_writer(
        output,
        opts,
        signer_key,
        encryption_keys,
        &file_path,
        total_input_size,
        timestamp,
    )?;
    dpkg_path.finalize()?;
    upload_dir.finalize()?;
    let package_location = dpkg_path.path.to_string_lossy().into_owned();
    show_summary(
        &format!("{}:{}/{}", &remote.host, &remote.port, &package_location),
        total_input_size,
        sftp.stat(&dpkg_path.path)?.size.unwrap_or(0),
    );
    Ok(Some(package_location))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_get_output_path() -> Result<(), Error> {
        assert!(get_output_path(Some("output"), None).is_err());
        let ts = Utc::now();
        let pwd = env::current_dir()?;
        let tmp = env::temp_dir();
        assert_eq!(
            get_output_path::<String>(None, Some(ts))?,
            pwd.join(ts.format(dpkg::DATETIME_FORMAT).to_string() + ".zip")
        );
        assert_eq!(
            get_output_path(Some(tmp.join("package.zip")), None)?,
            tmp.join("package.zip")
        );
        assert_eq!(
            get_output_path(Some(tmp.join("package.foo")), None)?,
            tmp.join("package.foo.zip")
        );
        assert_eq!(
            get_output_path(Some(tmp.clone()), Some(ts))?,
            tmp.join(ts.format(dpkg::DATETIME_FORMAT).to_string() + ".zip")
        );
        Ok(())
    }
}