// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{
    io::{Read, Write},
    path::Path,
};

use file_system::{File, OpenOptions, rename};
use kvproto::encryptionpb::EncryptedContent;
use protobuf::Message;
use rand::{RngCore, thread_rng};
use slog_global::error;
use tikv_util::time::Instant;

use crate::{Result, master_key::*, metrics::*};

mod header;
pub use header::*;

pub const TMP_FILE_SUFFIX: &str = "tmp";

/// An file encrypted by master key.
pub struct EncryptedFile<'a> {
    base: &'a Path,
    name: &'a str,
}

impl<'a> EncryptedFile<'a> {
    /// New an `EncryptedFile`.
    ///
    /// It's different from `file_system::File`, it does not hold a reference
    /// to the file or open the file, util we actually read or write.
    pub fn new(base: &'a Path, name: &'a str) -> EncryptedFile<'a> {
        EncryptedFile { base, name }
    }

    /// Read and decrypt the file. Caller need to handle the NotFound io error
    /// in case file not exists.
    pub fn read(&self, master_key: &dyn Backend) -> Result<Vec<u8>> {
        let start = Instant::now();
        let res = OpenOptions::new()
            .read(true)
            .open(self.base.join(self.name));
        match res {
            Ok(mut f) => {
                let mut buf = Vec::new();
                f.read_to_end(&mut buf)?;
                let (_, content, _) = Header::parse(&buf)?;
                let mut encrypted_content = EncryptedContent::default();
                encrypted_content.merge_from_bytes(content)?;
                let plaintext = master_key.decrypt(&encrypted_content)?;

                ENCRYPT_DECRPTION_FILE_HISTOGRAM
                    .with_label_values(&[self.name, "read"])
                    .observe(start.saturating_elapsed().as_secs_f64());

                Ok(plaintext)
            }
            Err(e) => Err(e.into()),
        }
    }

    pub fn write(&self, plaintext_content: &[u8], master_key: &dyn Backend) -> Result<()> {
        let start = Instant::now();
        // Write to a tmp file.
        // TODO what if a tmp file already exists?
        let origin_path = self.base.join(self.name);
        let mut tmp_path = origin_path.clone();
        tmp_path.set_extension(format!("{}.{}", thread_rng().next_u64(), TMP_FILE_SUFFIX));
        let mut tmp_file = OpenOptions::new()
            .create(true)
            .write(true)
            .open(&tmp_path)
            .map_err(|e| {
                error!(
                    "EncryptedFile::write open failed";
                    "path" => %tmp_path.display(),
                    "error" => %e,
                );
                e
            })?;

        // Encrypt the content.
        let encrypted_content = master_key
            .encrypt(plaintext_content)?
            .write_to_bytes()
            .unwrap();
        let header = Header::new(&encrypted_content, Version::V1);
        tmp_file.write_all(&header.to_bytes())?;
        tmp_file.write_all(&encrypted_content)?;
        tmp_file.sync_all()?;

        // Replace old file with the tmp file aomticlly.
        rename(tmp_path, origin_path)?;
        let base_dir = File::open(self.base)?;
        base_dir.sync_all()?;

        ENCRYPT_DECRPTION_FILE_HISTOGRAM
            .with_label_values(&[self.name, "write"])
            .observe(start.saturating_elapsed().as_secs_f64());

        // TODO GC broken temp files if necessary.
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::io::ErrorKind;

    use matches::assert_matches;

    use super::*;
    use crate::Error;

    #[test]
    fn test_open_write() {
        let tmp = tempfile::TempDir::new().unwrap();
        let file = EncryptedFile::new(tmp.path(), "encrypted");
        assert_eq!(file.base, tmp.path());
        assert_eq!(file.name, "encrypted");
        let ret = file.read(&PlaintextBackend::default());
        assert_matches!(ret, Err(Error::Io(_)));
        if let Err(Error::Io(e)) = file.read(&PlaintextBackend::default()) {
            assert_eq!(ErrorKind::NotFound, e.kind());
        }

        let content = b"test content";
        file.write(content, &PlaintextBackend::default()).unwrap();

        let file = EncryptedFile::new(tmp.path(), "encrypted");
        assert_eq!(file.read(&PlaintextBackend::default()).unwrap(), content);
    }
}