// Copyright 2016 Kitware, Inc.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Utilities for director-based tools.
//!
//! Tools written using this crate usually have other tasks related to management of the job files.
//! These functions are meant to be used in the tools so that these tasks are built into the tool
//! rather than managed using external scripts.

extern crate chrono;
use self::chrono::UTC;

extern crate lzma;
use self::lzma::LzmaWriter;

extern crate rand;
use self::rand::Rng;

extern crate serde_json;
use self::serde_json::Value;

extern crate tar;
use self::tar::Builder;

extern crate tempdir;
use self::tempdir::TempDir;

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

error_chain! { }

fn empty_job<K>(kind: K) -> Value
    where K: ToString,
{
    let mut job = BTreeMap::new();
    job.insert("kind".to_string(), Value::String(kind.to_string()));
    job.insert("data".to_string(), Value::Object(BTreeMap::new()));

    Value::Object(job)
}

fn write_job(queue: &Path, data: Value) -> Result<()> {
    let rndpart = rand::thread_rng()
        .gen_ascii_chars()
        .take(12)
        .collect::<String>();
    let filename = format!("{:?}-{}.json", UTC::now(), rndpart);
    let job_file = queue.join(&filename);
    let mut file = try!(File::create(&job_file).chain_err(|| "failed to create the job file"));
    try!(serde_json::to_writer(&mut file, &data)
        .chain_err(|| format!("failed to write the job to {:?}", filename)));

    Ok(())
}

/// Write a restart job to the given queue.
pub fn restart<Q>(queue: Q) -> Result<()>
    where Q: AsRef<Path>,
{
    let job = empty_job("watchdog:restart");
    write_job(queue.as_ref(), job)
}

/// Write an exit job to the given queue.
pub fn exit<Q>(queue: Q) -> Result<()>
    where Q: AsRef<Path>,
{
    let job = empty_job("watchdog:exit");
    write_job(queue.as_ref(), job)
}

const LZMA_COMPRESSION: u32 = 6;

/// Archive the jobs in the given queue into a tarball in the output directory.
///
/// Each subdirectory, `accept`, `fail`, and `reject` will be archived separately.
pub fn archive_queue<Q, O>(queue: Q, output: O) -> Result<()>
    where Q: AsRef<Path>,
          O: AsRef<Path>,
{
    for result in &["accept", "fail", "reject"] {
        let (filename, file) = try!(archive_file(output.as_ref(), result));
        let opt_writer = try!(archive_directory(queue.as_ref(), output.as_ref(), result, file));
        if let Some(mut writer) = opt_writer {
            try!(writer.finish()
                .chain_err(|| format!("failed to finish archive stream for {}", result)));
        } else {
            try!(fs::remove_file(&filename)
                .chain_err(|| format!("failed to delete file {:?}", filename)));
        }
    }

    Ok(())
}

fn archive_file(path: &Path, result: &str) -> Result<(PathBuf, LzmaWriter<File>)> {
    let now = UTC::now();
    let filepath = path.join(format!("{}-{}.tar.xz", now.to_rfc3339(), result));
    let file = try!(File::create(&filepath)
        .chain_err(|| format!("failed to create output file {:?}", filepath)));
    let writer = try!(LzmaWriter::new_compressor(file, LZMA_COMPRESSION)
        .chain_err(|| "failed to construct LZMA writer"));

    Ok((filepath, writer))
}

fn archive_directory<O>(path: &Path, workdir: &Path, subdir: &str, output: O) -> Result<Option<O>>
    where O: Write,
{
    let tempdir = try!(TempDir::new_in(workdir, "archive-jobs")
        .chain_err(|| "failed to create temporary directory"));
    let targetdir = tempdir.path().join(subdir);
    try!(fs::create_dir_all(&targetdir)
        .chain_err(|| "failed to create target directory"));
    let mut archive = Builder::new(output);
    try!(archive.append_dir(subdir, &targetdir)
        .chain_err(|| format!("failed to append directory to {}", subdir)));
    let entries = try!(fs::read_dir(path.join(subdir))
        .chain_err(|| "failed to read input directory"));
    let mut is_empty = true;
    for entry in entries {
        is_empty = false;
        let entry = try!(entry.chain_err(|| "failed to read entry"));
        let path = entry.path();
        let file_name = path.file_name().unwrap();
        let target_path = targetdir.join(&file_name);
        try!(fs::rename(&path, &target_path)
            .chain_err(|| format!("failed to move input file {:?}", path)));
        let mut added_file = try!(File::open(&target_path)
            .chain_err(|| format!("failed to open job file {:?}", &target_path)));
        try!(archive.append_file(format!("{}/{}", subdir, file_name.to_string_lossy()), &mut added_file)
            .chain_err(|| format!("failed to append {:?} to the archive", target_path)));
    }

    if is_empty {
        return Ok(None);
    }

    Ok(Some(try!(archive.into_inner()
        .chain_err(|| format!("failed to finish TAR stream for {}", subdir)))))
}
