// 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.

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

extern crate itertools;
use self::itertools::Itertools;

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

extern crate serde_json;
use self::serde_json::{Map, Value};

use super::error::*;
use super::handler::{Handler, HandlerResult};
use super::watcher::Watcher;

use std::collections::hash_map::{Entry, HashMap};
use std::ffi::OsStr;
use std::fmt::{self, Debug};
use std::fs::{self, File};
use std::io::Write;
use std::path::{Path, PathBuf};

/// Dispatch jobs to registered handlers.
///
/// Jobs are sorted into `accept`, `reject`, and `fail` directories based on whether they were
/// accepted by the relevant handler. Once handled, a `.stamp` file containing the timestamp of
/// when the job was completed is created beside the final location. In addition, rejected and
/// failed jobs have a `.reason` file describing what happened.
pub struct Director<'a> {
    accept: PathBuf,
    reject: PathBuf,
    fail: PathBuf,

    handlers: HashMap<String, &'a Handler>,
}

fn log_file_name<P>(path: &P) -> &OsStr
    where P: AsRef<Path>,
{
    path.as_ref().file_name().unwrap_or(OsStr::new("."))
}

impl<'a> Debug for Director<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "Director {{ accept: {:?}, reject: {:?}, fail: {:?}, handlers: {:?} }}",
               self.accept,
               self.reject,
               self.fail,
               self.handlers.keys().collect::<Vec<_>>())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
/// The result of running jobs.
pub enum RunResult {
    /// The director should continue running.
    Continue,
    /// The director should be restarted from the beginning.
    Restart,
    /// The director should be terminated.
    Done,
}

impl RunResult {
    fn should_exit(&self) -> bool {
        if let RunResult::Continue = *self {
            false
        } else {
            true
        }
    }

    fn is_done(&self) -> bool {
        if let RunResult::Done = *self {
            true
        } else {
            false
        }
    }
}

impl<'a> Director<'a> {
    /// Creates a new `Director`.
    ///
    /// The `accept` and `reject` directories are created as children of the given directory.
    pub fn new(root: &Path) -> Result<Self> {
        if !root.is_dir() {
            bail!(ErrorKind::NotADirectory(root.to_path_buf()));
        }

        info!(target: "director", "setting up a director in {:?}", root);

        let accept_dir = root.join("accept");
        let reject_dir = root.join("reject");
        let fail_dir = root.join("fail");
        try!(fs::create_dir_all(&accept_dir)
            .chain_err(|| "failed to create the 'accept' directory"));
        try!(fs::create_dir_all(&reject_dir)
            .chain_err(|| "failed to create the 'reject' directory"));
        try!(fs::create_dir_all(&fail_dir).chain_err(|| "failed to create the 'fail' directory"));

        Ok(Director {
            accept: accept_dir,
            reject: reject_dir,
            fail: fail_dir,

            handlers: HashMap::new(),
        })
    }

    /// Add a handler for jobs labeled as `kind`.
    pub fn add_handler<K>(&mut self, kind: K, handler: &'a Handler) -> Result<()>
        where K: ToString,
    {
        match self.handlers.entry(kind.to_string()) {
            Entry::Occupied(o) => bail!(ErrorKind::DuplicateHandler(o.key().clone())),
            Entry::Vacant(v) => {
                debug!(target: "director", "adding handler '{}'", v.key());

                v.insert(handler);
                Ok(())
            },
        }
    }

    // Move the job file into the appropriate directory.
    fn tag(&self, target_dir: &Path, file: &Path) -> Result<PathBuf> {
        let mut target_path = target_dir.to_path_buf();
        target_path.push(file.file_name().unwrap());

        // Write the stamp file.
        let mut stamp_file = try!(File::create(target_path.with_extension("stamp"))
            .chain_err(|| {
                format!("failed to create a stamp file for {:?}",
                        log_file_name(&target_path))
            }));
        let time = UTC::now();
        try!(write!(stamp_file, "{}\n", time.to_string()).chain_err(|| {
            format!("failed to write the stamp file for {:?}",
                    log_file_name(&target_path))
        }));

        // Rename the file into the target path.
        try!(fs::rename(file, &target_path)
            .chain_err(|| format!("failed to move to {:?}", target_path)));

        Ok(target_path)
    }

    fn tag_with_reason<R>(&self, target_dir: &Path, file: &Path, reason: R) -> Result<()>
        where R: fmt::Display,
    {
        let target_file = try!(self.tag(target_dir, file));
        self.write_reason(&target_file, reason)
    }

    fn write_reason<R>(&self, target_file: &Path, reason: R) -> Result<()>
        where R: fmt::Display,
    {
        let mut reason_file = try!(File::create(target_file.with_extension("reason"))
            .chain_err(|| {
                format!("failed to create the reason file for {:?}",
                        log_file_name(&target_file))
            }));
        try!(write!(reason_file, "{}\n", reason).chain_err(|| {
            format!("failed to write the reason file for {:?}",
                    log_file_name(&target_file))
        }));

        Ok(())
    }

    /// Handle a single file.
    ///
    /// Returns `true` if the loop should terminate.
    pub fn run_one(&self, path: &Path) -> RunResult {
        match self.dispatch(path) {
            Ok(res) => res,
            Err(err) => {
                error!("error when handling {:?}: {}", log_file_name(&path), err);
                RunResult::Done
            },
        }
    }

    /// Handle a batch of files at once.
    ///
    /// Returns `true` if the loop should terminate.
    pub fn run(&self, paths: Vec<PathBuf>) -> RunResult {
        let mut result = RunResult::Continue;

        for path in paths {
            let one_result = self.run_one(&path);
            if one_result > result {
                result = one_result;
            }

            if result.is_done() {
                break;
            }
        }

        result
    }

    /// Watch a directory for new job files and act upon them.
    ///
    /// Upon startup, this searches the directory for existing files. These files are processed in
    /// lexigraphic order.
    pub fn watch_directory(&self, path: &Path) -> Result<RunResult> {
        let mut watcher = {
            let paths = try!(path.read_dir()
                .chain_err(|| format!("failed to list the queue directory {:?}", path)));
            let watcher = try!(Watcher::new(path).chain_err(|| {
                format!("failed to start the watcher in the queue directory {:?}",
                        path)
            }));

            let path_entries = paths.filter_map(|e| e.ok())
                .filter(|e| !e.path().is_dir())
                .map(|e| e.path())
                .sorted();
            let result = self.run(path_entries);
            if result.should_exit() {
                return Ok(result);
            }

            watcher
        };

        let loop_result;

        loop {
            let events = try!(watcher.events()
                .chain_err(|| format!("failed to watch the queue directory {:?}", path)));

            let path_entries = events.into_iter()
                .filter(|event| !event.is_dir())
                .map(|event| path.join(&event.name))
                .sorted();
            let result = self.run(path_entries);
            if result.should_exit() {
                loop_result = result;
                break;
            }
        }

        Ok(loop_result)
    }

    // Takes the file and directs it to the right location.
    fn dispatch(&self, file: &Path) -> Result<RunResult> {
        trace!(target: "director", "handling file {:?}", log_file_name(&file));

        let ext = file.extension();

        if ext.is_none() || ext.unwrap() != "json" {
            return Ok(RunResult::Continue);
        }

        let event_file = try!(File::open(file)
            .chain_err(|| format!("failed to open an input file: {:?}", file)));
        let mut payload: Value = if let Ok(payload) = serde_json::from_reader(event_file) {
            payload
        } else {
            return Ok(RunResult::Continue);
        };

        if !payload.is_object() {
            return Ok(RunResult::Continue);
        }

        let ret = match try!(self.handle(&payload)) {
            HandlerResult::Accept => {
                debug!("accepted {:?}", log_file_name(&file));
                try!(self.tag(&self.accept, file));
                RunResult::Continue
            },
            HandlerResult::Defer(ref reason) => {
                debug!("deferring {:?}: {}", log_file_name(&file), reason);
                let rndpart = rand::thread_rng()
                    .gen_ascii_chars()
                    .take(12)
                    .collect::<String>();
                let defer_job_file =
                    file.with_file_name(format!("{:?}-{}.json", UTC::now(), rndpart));
                Self::add_reason_to_object(&mut payload, file.to_string_lossy(), &reason);
                let mut defer_file = try!(File::create(&defer_job_file).chain_err(|| {
                    format!("failed to create the job for the deferred job {:?} as {:?}",
                            log_file_name(&file),
                            log_file_name(&defer_job_file))
                }));
                try!(serde_json::to_writer(&mut defer_file, &payload).chain_err(|| {
                    format!("failed to write the deferred job to {:?}",
                            log_file_name(&defer_job_file))
                }));
                try!(self.tag_with_reason(&self.reject, file, reason));
                RunResult::Continue
            },
            HandlerResult::Reject(ref reason) => {
                debug!("rejecting {:?}: {}", file, reason);
                try!(self.tag_with_reason(&self.reject, file, reason));
                RunResult::Continue
            },
            HandlerResult::Fail(ref reason) => {
                debug!("failed {:?}: {:?}", file, reason);
                try!(self.tag_with_reason(&self.fail, file, format!("{:?}", reason)));
                RunResult::Continue
            },
            HandlerResult::Restart => {
                info!(target: "director", "restarting via {:?}", log_file_name(&file));

                try!(self.tag(&self.accept, file));
                RunResult::Restart
            },
            HandlerResult::Done => {
                info!(target: "director", "completed via {:?}", log_file_name(&file));

                try!(self.tag(&self.accept, file));
                RunResult::Done
            },
        };

        trace!(target: "director", "handled file {:?}", log_file_name(&file));

        Ok(ret)
    }

    fn add_reason_to_object<N, R>(object: &mut Value, name: N, reason: R)
        where N: ToString,
              R: ToString,
    {
        let retry_map = object.as_object_mut()
            .unwrap() // Verified above.
            .entry("retry".to_owned())
            .or_insert_with(|| Value::Object(Map::new()));

        if !retry_map.is_object() {
            *retry_map = Value::Object(Map::new());
        }

        retry_map.as_object_mut()
            .unwrap() // Ensured above.
            .insert(name.to_string(), Value::String(reason.to_string()));
    }

    // Dispatch the contents of the file to the right handler.
    fn handle(&self, payload: &Value) -> Result<HandlerResult> {
        let kind = match payload.pointer("/kind") {
            Some(&Value::String(ref kind)) => kind,
            Some(_) => return Ok(HandlerResult::Reject("'kind' is not a string".to_string())),
            None => return Ok(HandlerResult::Reject("no 'kind'".to_string())),
        };
        let data = match payload.pointer("/data") {
            Some(data) => data,
            None => return Ok(HandlerResult::Reject("no 'data'".to_string())),
        };
        let retry_reasons = match payload.pointer("/retry") {
            Some(&Value::Object(ref reasons)) => {
                let retry_reasons = reasons.iter()
                    .map(|(_, reason)| {
                        reason.as_str()
                            .map(ToString::to_string)
                            .ok_or_else(|| {
                                HandlerResult::Reject("retry reason is not a string".to_string())
                            })
                    })
                    .collect::<::std::result::Result<Vec<_>, HandlerResult>>();

                match retry_reasons {
                    Ok(reasons) => reasons,
                    Err(reject) => return Ok(reject),
                }
            },
            Some(_) => return Ok(HandlerResult::Reject("'retry' is not an object".to_string())),
            None => vec![],
        };

        match self.handlers.get(kind) {
            Some(handler) => handler.handle_retry(kind, data, retry_reasons),
            None => Ok(HandlerResult::Reject(format!("no handler for kind '{}'", kind))),
        }
    }
}
