extern crate libc;
extern crate log;
extern crate tokio;
extern crate zstd;

use std::ffi::OsString;
use std::fs::{File, OpenOptions};
use std::io::Write;
use std::os::unix::fs::PermissionsExt;
use std::os::unix::io::AsRawFd;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

use tokio::prelude::*;

use log::{info, warn};
use log::{Level, Log, Metadata, Record, SetLoggerError};

/// Kernel logger implementation
pub struct KernelLog {
    kmsg: Mutex<File>,
    level: Level,
}

impl KernelLog {
    /// Create new kernel logger
    pub fn new() -> KernelLog {
        KernelLog::with_level(Level::Trace)
    }

    /// Create new kernel logger with error level filter
    pub fn with_level(level: Level) -> KernelLog {
        KernelLog {
            kmsg: Mutex::new(OpenOptions::new().write(true).open("/dev/kmsg").unwrap()),
            level: level,
        }
    }

    pub fn init_with_level(level: Level) -> Result<(), SetLoggerError> {
        let logger = KernelLog::with_level(level);
        log::set_boxed_logger(Box::new(logger))
    }

    pub fn init() -> Result<(), SetLoggerError> {
        KernelLog::init_with_level(Level::Trace)
    }
}

impl Log for KernelLog {
    fn enabled(&self, meta: &Metadata) -> bool {
        meta.level() <= self.level
    }

    fn log(&self, record: &Record) {
        if !self.enabled(record.metadata()) {
            return;
        }

        let level: u8 = match record.level() {
            Level::Error => 3,
            Level::Warn => 4,
            Level::Info => 5,
            Level::Debug => 6,
            Level::Trace => 7,
        };

        let mut buf = Vec::new();
        writeln!(
            buf,
            "<{}>{}[{}]: {}",
            level,
            record.target(),
            unsafe { libc::getpid() },
            record.args()
        )
        .unwrap();

        let mut kmsg = self.kmsg.lock().unwrap();
        let _ = kmsg.write(&buf);
    }

    fn flush(&self) {
        let mut kmsg = self.kmsg.lock().unwrap();
        let _ = kmsg.flush();
    }
}

fn app_path() -> String {
    match std::env::current_exe() {
        Ok(exe) => String::from(exe.to_str().unwrap()),
        Err(_) => String::from("rustdump"),
    }
}

fn app_name() -> String {
    match std::env::current_exe() {
        Ok(exe) => String::from(Path::new(&exe).file_name().unwrap().to_str().unwrap()),
        Err(_) => String::from("rustdump"),
    }
}

fn usage() -> ! {
    eprintln!(
        "Usage: set /proc/sys/kernel/core_pattern to '|{} %P %u %g %s %t %c %h %e",
        app_path()
    );
    std::process::exit(1);
}

fn parse_str(opts: &[OsString], i: usize) -> &str {
    match opts[i].to_str() {
        Some(v) => v,
        None => {
            eprintln!("FATAL: can't decode argument {}", i);
            usage();
        }
    }
}

fn parse_int(opts: &[OsString], i: usize) -> u64 {
    match parse_str(opts, i).parse::<u64>() {
        Ok(v) => v,
        Err(e) => {
            eprintln!("FATAL: can't parse argument {}: {:?}", i, e);
            usage();
        }
    }
}

// default pattern %P %u %g %s %t %c %h %e
//                  1  2  3  4  5  6  7  8
fn process_args() -> (String, u64, u64) {
    match std::env::args_os().len() {
        9 => {
            let args: Vec<_> = std::env::args_os().collect();
            let pid = parse_int(&args, 1);
            let uid = parse_int(&args, 2);
            let gid = parse_int(&args, 3);
            let _sig = parse_int(&args, 4);
            let ts = parse_str(&args, 5);
            let _lim = parse_int(&args, 6);
            let host = parse_str(&args, 7);
            let exe = parse_str(&args, 8);
            (format!("{}.{}.{}.{}.zst", exe, pid, host, ts), uid, gid)
        }
        _ => usage(),
    }
}

fn dump(dest: std::fs::File) {
    let stdin = tokio::io::stdin();
    let z = zstd::Encoder::new(
        tokio::fs::File::from_std(dest),
        zstd::DEFAULT_COMPRESSION_LEVEL,
    )
    .unwrap();
    let future = tokio::io::copy(stdin, z)
        .map_err(|e| {
            warn!("dump failed: {}", e);
        })
        .and_then(|copied| {
            info!("processed {} bytes", copied.0);
            Ok(())
        });

    tokio::run(future);
}

fn main() {
    unsafe {
        libc::prctl(libc::PR_SET_DUMPABLE, 0);
    }

    // If argument parsing is not successful, exit and don't initialize anything
    let (name, uid, gid) = process_args();

    // Setup kernel logger as a default logger
    KernelLog::init().unwrap();
    log::set_max_level(log::LevelFilter::Trace);

    let mut corepath = PathBuf::from("/var/core");
    // Ensure target directory exists
    std::fs::create_dir_all(&corepath).unwrap();

    // Append filename and open destination file
    corepath.push(name);
    let dest = std::fs::File::create(corepath).unwrap();

    // Adjust file ownership and permissions
    unsafe {
        let fd = dest.as_raw_fd();
        libc::fchmod(fd, 0o640);
        libc::fchown(fd, uid as u32, gid as u32);
    }

    // Compress and write stdin to file
    dump(dest);
}