use crate::{Repository, arq, errors::ArqError, utils::ReadExt};
use aes::cipher::{BlockDecryptMut, KeyIvInit, block_padding::Pkcs7};
use byteorder::{BigEndian, ReadBytesExt};
use chrono::{DateTime, Utc};
use display_json::DisplayAsJsonPretty;
use hmac::{Hmac, Mac};
use log::trace;
use pbkdf2::pbkdf2_hmac_array;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{
    collections::{HashMap, VecDeque},
    fs,
    io::{self, Read, Write},
    path::{Path, PathBuf},
    str::{self},
};

pub type Result<T> = std::result::Result<T, ArqError>;

const ENCRYPTED_OBJECT_HEADER: &[u8; 4] = b"ARQO";

#[derive(Debug, Serialize, Deserialize, DisplayAsJsonPretty)]
pub struct BackupPlan {
    #[serde(rename = "planUUID")]
    pub uuid: String,
    pub name: String,
    #[serde(rename = "backupFolderPlansByUUID")]
    pub backup_folders: HashMap<String, BackupFolder>,
}

impl BackupPlan {
    pub fn from_reader<R: Read>(reader: R, key_set: &KeySet) -> arq::Result<Self> {
        Ok(serde_json::from_slice(&decrypt_object(reader, key_set)?)?)
    }
}

#[derive(Debug, Serialize, Deserialize, DisplayAsJsonPretty)]
pub struct BackupFolder {
    #[serde(alias = "backupFolderUUID")]
    pub uuid: String,
    pub name: String,
    #[serde(rename = "localPath")]
    pub local_path: String,
}

impl BackupFolder {
    pub fn from_reader<R: Read>(reader: R, key_set: &KeySet) -> arq::Result<Self> {
        Ok(serde_json::from_slice(&decrypt_object(reader, key_set)?)?)
    }
}

#[derive(Debug, Serialize, Deserialize, DisplayAsJsonPretty)]
pub struct BackupRecord {
    #[serde(rename = "volumeName")]
    pub volume_name: Option<String>,
    #[serde(rename = "backupFolderUUID")]
    pub backup_folder_uuid: String,
    pub version: u64,
    #[serde(rename = "arqVersion")]
    pub arq_version: Option<String>,
    #[serde(rename = "node")]
    pub root: Node,
    #[serde(rename = "creationDate", with = "chrono::serde::ts_seconds")]
    pub creation_time: DateTime<Utc>,
}

impl BackupRecord {
    pub fn from_reader<R: Read>(reader: R, key_set: &KeySet) -> arq::Result<Self> {
        Ok(serde_json::from_slice(&decrypt_lz4decompress(
            reader, key_set,
        )?)?)
    }
}

#[derive(Debug, Serialize, Deserialize, DisplayAsJsonPretty, Clone)]
pub struct Node {
    #[serde(default)]
    pub name: String,

    #[serde(rename = "isTree")]
    pub is_tree: bool,

    #[serde(rename = "itemSize")]
    pub item_size: u64,

    #[serde(rename = "containedFilesCount")]
    pub contained_files_count: u64,

    #[serde(rename = "treeBlobLoc")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tree_loc: Option<BlobLoc>,

    #[serde(rename = "dataBlobLoc")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data_locs: Option<Vec<BlobLoc>>,

    // unclear

    // 1 = macOS, ? = windows
    #[serde(rename = "computerOSType")]
    pub computer_os_type: u32,
    pub deleted: bool,

    #[serde(rename = "creationTime_sec", with = "chrono::serde::ts_seconds_option")]
    #[serde(skip_serializing_if = "Option::is_none")]
    /// This is the birth time of the inode
    pub creation_time: Option<DateTime<Utc>>,
    #[serde(rename = "changeTime_sec")]
    pub ctime_sec: Option<i64>,
    #[serde(rename = "changeTime_nsec")]
    pub ctime_nsec: Option<u64>,
    #[serde(rename = "modificationTime_sec")]
    pub mtime_sec: Option<i64>,
    #[serde(rename = "modificationTime_nsec")]
    pub mtime_nsec: Option<u64>,

    #[serde(rename = "userName")]
    pub user_name: Option<String>,
    #[serde(rename = "mac_st_uid")]
    pub uid: u32,
    #[serde(rename = "groupName")]
    pub group_name: Option<String>,
    #[serde(rename = "mac_st_gid")]
    pub gid: u32,
    #[serde(rename = "mac_st_mode")]
    pub mode: u32,
}

impl Node {
    pub fn from_binary(mut reader: impl Read, name: &str, tree_version: u32) -> arq::Result<Node> {
        // [Bool:isTree]
        let is_tree = reader.read_bool()?;
        // [BlobLoc:treeBlobLoc] /* present if isTree is true */
        let tree_loc = if is_tree {
            Some(BlobLoc::from_binary(&mut reader, tree_version)?)
        } else {
            None
        };

        // [UInt32:computerOSType]
        let computer_os_type = reader.read_u32::<BigEndian>()?;
        // [UInt64:dataBlobLocsCount]
        let data_locs_count = reader.read_u64::<BigEndian>()?;

        let data_locs = (data_locs_count > 0)
            .then(|| {
                (0..data_locs_count)
                    .map(|_| BlobLoc::from_binary(&mut reader, tree_version))
                    .collect()
            })
            .transpose()?;

        // [Bool:aclBlobLocIsNotNil]
        let acl_blob_loc_is_not_nil = reader.read_bool()?;
        if acl_blob_loc_is_not_nil {
            // [BlobLoc:aclBlobLoc] /* present if aclBlobLocIsNotNil is true */
            let _ = BlobLoc::from_binary(&mut reader, tree_version);
        }

        // [UInt64:xattrsBlobLocCount]
        let xattrs_locs_count = reader.read_u64::<BigEndian>()?;
        for _ in 1..=xattrs_locs_count {
            // [BlobLoc:xattrsBlobLoc]
            let _ = BlobLoc::from_binary(&mut reader, tree_version);
        }
        // [UInt64:itemSize]
        let item_size = reader.read_u64::<BigEndian>()?;
        // [UInt64:containedFilesCount]
        let contained_files_count = reader.read_u64::<BigEndian>()?;
        // [Int64:mtime_sec]
        let mtime_sec = reader.read_i64::<BigEndian>()?;
        // [Int64:mtime_nsec]
        let mtime_nsec = reader.read_i64::<BigEndian>()? as u64;
        // [Int64:ctime_sec]
        let ctime_sec = reader.read_i64::<BigEndian>()?;
        // [Int64:ctime_nsec]
        let ctime_nsec = reader.read_i64::<BigEndian>()? as u64;
        // [Int64:create_time_sec]
        let create_time_sec = reader.read_i64::<BigEndian>()?;
        // [Int64:create_time_nsec]
        let create_time_nsec = reader.read_i64::<BigEndian>()? as u32;
        // [String:username]
        let user_name = reader.read_string()?;
        // [String:groupName]
        let group_name = reader.read_string()?;
        // [Bool:deleted]
        let deleted = reader.read_bool()?;
        // [Int32:mac_st_dev]
        reader.read_i32::<BigEndian>()?;
        // [UInt64:mac_st_ino]
        reader.read_u64::<BigEndian>()?;
        // [UInt32:mac_st_mode]
        let mode = reader.read_u32::<BigEndian>()?;
        // [UInt32:mac_st_nlink]
        reader.read_u32::<BigEndian>()?;
        // [UInt32:mac_st_uid]
        let uid = reader.read_u32::<BigEndian>()?;
        // [UInt32:mac_st_gid]
        let gid = reader.read_u32::<BigEndian>()?;
        // [Int32:mac_st_rdev]
        reader.read_i32::<BigEndian>()?;
        // [Int32:mac_st_flags]
        reader.read_i32::<BigEndian>()?;
        // [UInt32:win_attrs]
        reader.read_u32::<BigEndian>()?;

        if tree_version >= 2 {
            // [UInt32:win_reparse_tag] /* if Tree version >= 2 */
            reader.read_u32::<BigEndian>()?;
            // [Bool:win_reparse_point_is_directory] /* if Tree version >= 2 */
            reader.read_bool()?;
        }

        Ok(Node {
            name: name.to_string(),
            is_tree,
            item_size,
            contained_files_count,
            creation_time: DateTime::from_timestamp(create_time_sec, create_time_nsec),
            ctime_sec: Some(ctime_sec),
            ctime_nsec: Some(ctime_nsec),
            mtime_sec: Some(mtime_sec),
            mtime_nsec: Some(mtime_nsec),
            tree_loc,
            data_locs,
            user_name,
            uid,
            group_name,
            gid,
            deleted,
            computer_os_type,
            mode,
        })
    }

    pub fn restore(
        &self,
        repo: &Repository,
        target_path: &Path,
        dry_run: bool,
    ) -> arq::Result<[u8; 32]> {
        if self.is_tree {
            trace!("Creating directory {:?} to {:?}", self.name, target_path);

            if !dry_run {
                fs::create_dir_all(target_path)?
            };

            Ok([0u8; 32])
        } else {
            trace!("Restoring file {}", target_path.display());

            let mut hasher = Sha256::new();

            let file: &mut dyn Write = if dry_run {
                &mut io::sink()
            } else {
                &mut fs::File::create_new(target_path)?
            };

            let Some(data_locs) = self.data_locs.as_ref() else {
                return Err(ArqError::MissingDataLocs);
            };

            for (i, data_loc) in data_locs.iter().enumerate() {
                trace!(
                    "Loading data blob {}/{} ({} Bytes compressed): {}",
                    i + 1,
                    data_locs.len(),
                    data_loc.length,
                    data_loc.relative_path
                );
                let buf = &repo.data_from_pack(data_loc)?;
                hasher.update(buf);
                if !dry_run {
                    file.write_all(buf)?;
                }
            }

            let hash = hasher.finalize();

            Ok(hash.into())
        }
    }
}

#[derive(Debug, Serialize, Deserialize, DisplayAsJsonPretty, Clone, PartialEq, Eq, Hash)]
pub struct BlobLoc {
    pub offset: u64,
    pub length: u64,
    #[serde(rename = "relativePath")]
    pub relative_path: String,

    #[serde(rename = "blobIdentifier")]
    pub blob_id: String,

    // 0 = none
    // 1 = GZip (Arq 5/6)
    // 2 = LZ4 (Arq 7)
    #[serde(rename = "compressionType")]
    pub compression_type: u32,
    // true for new data, false for "very old data".
    // It's unclear what is considered "very old" data, or what this
    // actually means.
    #[serde(rename = "stretchEncryptionKey")]
    pub stretch_encryption_key: bool,
    // From the documentation:
    //     A “blob” is just a chunk of data stored either in a pack file
    //    (in the blobpacks or largeblobpacks subdirectory) or as a standalone
    //    file in the “standardobjects” (or “standardiaobjects” etc depending
    //    on the storage location type and storage class chosen).
    //
    // This probably means, that if is_packed=true, the entire file at
    // relative_path should be used.
    // TODO: Find some real-world examples and test them.
    #[serde(rename = "isPacked")]
    pub is_packed: bool,
}

// From the documentation:
//    [String:blobIdentifier] /* can't be null */
//    [Bool:isPacked]
//    [String:relativePath]
//    [UInt64:offset]
//    [UInt64:length
//    [Bool:stretchEncryptionKey]
//    [UInt32:compressionType]
//
// HOWEVER there seems to be an extra byte-wide item, between
// isPacked and relativePath, at least for tree version > 1.
impl BlobLoc {
    #[allow(unused)]
    pub const COMPRESSION_TYPE_NONE: u32 = 0;
    #[allow(unused)]
    pub const COMPRESSION_TYPE_GZIP: u32 = 1;
    pub const COMPRESSION_TYPE_LZ4: u32 = 2;

    pub fn from_binary<R: Read>(mut reader: R, tree_version: u32) -> arq::Result<Self> {
        // [String:blobIdentifier] /* can't be null */
        let Some(blob_id) = reader.read_string()? else {
            return Err(ArqError::BlobIdMissing);
        };
        // [Bool:isPacked]
        let is_packed = reader.read_bool()?;

        // There seems to be an extra, undocumented Byte in the Arq 7 data format.
        // tree version 1 does not seem to have this extra byte.
        if tree_version > 1 {
            let unknown = &mut [0u8; 1];
            reader.read_exact(unknown)?;
        }

        // [String:relativePath]
        let relative_path = reader.read_string_default()?;
        // [UInt64:offset]
        let offset = reader.read_u64::<BigEndian>()?;
        // [UInt64:length
        let length = reader.read_u64::<BigEndian>()?;
        // [Bool:stretchEncryptionKey]
        let stretch_encryption_key = reader.read_bool()?;
        // [UInt32:compressionType]
        let compression_type = reader.read_u32::<BigEndian>()?;

        // type=0 (no compression) is only used by very old Arq versions and probably
        // won't be implemented ever.
        // TODO: Get some test data for type=2 (GZip) and implement this.
        if compression_type != Self::COMPRESSION_TYPE_LZ4 {
            return Err(ArqError::UnsupportedCompressionType(compression_type));
        }

        Ok(Self {
            offset,
            length,
            relative_path,
            blob_id,
            compression_type,
            stretch_encryption_key,
            is_packed,
        })
    }
}

#[derive(Debug, Serialize, Deserialize, DisplayAsJsonPretty)]
pub struct Tree {
    pub version: u32,
    pub child_count: u64,
    pub nodes: Vec<Node>,
}

impl Tree {
    pub fn from_binary<R: Read>(mut reader: R) -> arq::Result<Self> {
        let version = reader.read_u32::<BigEndian>()?;
        let child_count = reader.read_u64::<BigEndian>()?;
        trace!("Tree: version={}, child_count={}", version, child_count);

        let mut nodes: Vec<Node> = vec![];

        for n in 1..=child_count {
            let name = reader.read_string_default()?;

            trace!("Child-Node: n={}/{}, name={}", n, child_count, name);
            nodes.push(Node::from_binary(&mut reader, &name, version)?);
        }

        Ok(Tree {
            version,
            child_count,
            nodes,
        })
    }
}

#[derive(Debug, Serialize, DisplayAsJsonPretty)]
pub struct KeySet {
    /// master encryption key
    pub enc_key: [u8; 32],
    /// master HMAC key
    pub hmac_key: [u8; 32],
}

impl KeySet {
    const ROUNDS: u32 = 200_000;
    const KEY_SET_VERSION: u32 = 3;
    const KEY_SET_HEADER: &[u8; 25] = b"ARQ_ENCRYPTED_MASTER_KEYS";

    pub fn decrypt_parse_key_set<R: Read>(reader: R, password: &str) -> arq::Result<KeySet> {
        KeySet::from_binary(&KeySet::decrypt(reader, password)?[..])
    }

    pub fn decrypt<R: Read>(mut reader: R, password: &str) -> arq::Result<Vec<u8>> {
        // == encryptedkeyset.dat encrypted data format
        // header     : 25 bytes, always "ARQ_ENCRYPTED_MASTER_KEYS"
        // salt       :  8 bytes
        // hmac       : 32 bytes
        // iv         : 16 bytes
        // ciphertext : <rest>

        // == plaintext data format
        // version                    : 4 bytes
        // encryption key length      : 8 bytes
        // encryption key             : $encryption_key_length bytes
        // hmac key length            : 8 bytes
        // hmac key                   : $hmac_key_length bytes
        // blob identifier salt length: 8 bytes
        // blob identifier salt       : $blob_identifier_salt_length bytes

        let mut header = [0; 25];
        reader.read_exact(&mut header)?;

        if header != *KeySet::KEY_SET_HEADER {
            return Err(ArqError::InvalidKeySetHeader);
        }

        let mut salt = [0; 8];
        reader.read_exact(&mut salt)?;

        let mut hmac = [0; 32];
        reader.read_exact(&mut hmac)?;

        let mut iv = [0; 16];
        reader.read_exact(&mut iv)?;

        let mut ciphertext = vec![];
        reader.read_to_end(&mut ciphertext)?;

        let key = pbkdf2_hmac_array::<Sha256, 64>(password.as_bytes(), &salt, KeySet::ROUNDS);

        let derived_enc_key = &key[0..32];
        let derived_hmac_key = &key[32..64];

        // verify HMAC
        let mut mac = Hmac::<Sha256>::new_from_slice(derived_hmac_key)
            .expect("session IV or key have incorrect length");
        mac.update(&iv);
        mac.update(&ciphertext);

        // match mac.verify_slice(&hmac) {
        //     Ok(_) => {}
        //     Err(err) => return Err(ArqError::KeySetMacError(err)),
        // }
        mac.verify_slice(&hmac).map_err(ArqError::KeySetMacError)?;
        // .expect("KeySet HMAC mismatch; incorrect Password?");

        // decrypt key set
        Ok(
            cbc::Decryptor::<aes::Aes256>::new(derived_enc_key.into(), &iv.into())
                .decrypt_padded_vec_mut::<Pkcs7>(&ciphertext)
                .expect("Error decrypting master key set!"),
        )
    }

    pub fn from_binary<R: Read>(mut reader: R) -> arq::Result<Self> {
        let version = reader.read_u32::<BigEndian>()?;
        if version != KeySet::KEY_SET_VERSION {
            return Err(ArqError::UnsupportedKeySetVersion(version));
        }

        let master_enc_key_len = reader.read_u64::<BigEndian>()?;
        let mut master_enc_key = vec![0u8; master_enc_key_len as usize];
        reader.read_exact(&mut master_enc_key)?;

        let hmac_key_len = reader.read_u64::<BigEndian>()?;
        let mut hmac_key = vec![0u8; hmac_key_len as usize];
        reader.read_exact(&mut hmac_key)?;

        Ok(KeySet {
            enc_key: master_enc_key.try_into().map_err(ArqError::InvalidKey)?,
            hmac_key: hmac_key.try_into().map_err(ArqError::InvalidKey)?,
        })
    }
}

type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

// Decrypt encrypted Arq object. Both the HMAC and encryption key from
// the key set are used to verify and decrypt a session initialization value
// as well as a session key. The session IV and key are then used to decrypt
// the actual payload in the object.
//
// All objects in a Arq backup are encrypted in this format, except the
// key set.
//
// The format of an encrypted object looks like this:
//   header                         : 4 bytes, always "ARQO"
//   HMAC SHA-256                   : 32 bytes
//   master IV                      : 16 bytes
//   encrypted data IV + session key: 64 bytes
//   ciphertext                     : <rest>
pub fn decrypt_object<R: Read>(mut reader: R, key_set: &KeySet) -> arq::Result<Vec<u8>> {
    let mut header = [0; 4];
    reader.read_exact(&mut header)?;

    if header != *ENCRYPTED_OBJECT_HEADER {
        return Err(ArqError::InvalidEncryptedObjectHeader);
    }

    let mut hmac = [0; 32];
    reader.read_exact(&mut hmac)?;

    let mut iv = [0; 16];
    reader.read_exact(&mut iv)?;

    let mut enc_data_iv_session_key = [0; 64];
    reader.read_exact(&mut enc_data_iv_session_key)?;

    let mut ciphertext = vec![];
    reader.read_to_end(&mut ciphertext)?;

    // verify HMAC
    let mut mac = Hmac::<Sha256>::new_from_slice(&key_set.hmac_key)
        .expect("session IV or key have incorrect length");
    mac.update(&iv);
    mac.update(&enc_data_iv_session_key);
    mac.update(&ciphertext);
    mac.verify_slice(&hmac).expect("HMAC mismatch?");

    // decrypt data iv and session key
    let data_iv_session_key = Aes256CbcDec::new(&key_set.enc_key.into(), &iv.into())
        .decrypt_padded_vec_mut::<Pkcs7>(&enc_data_iv_session_key)
        .expect("Error decrypting data iv + session key!");

    let data_iv = &data_iv_session_key[0..16];
    let session_key = &data_iv_session_key[16..48];

    // decrypt payload
    let plaintext = Aes256CbcDec::new(session_key.into(), data_iv.into())
        .decrypt_padded_vec_mut::<Pkcs7>(&ciphertext)
        .expect("Error decrypting payload!");

    Ok(plaintext)
}

pub fn decrypt_lz4decompress(reader: impl Read, key_set: &KeySet) -> arq::Result<Vec<u8>> {
    let decrypted_data = &decrypt_object(reader, key_set)?;

    // Arq uses LZ4 block format, with a 4-byte big-endian uncompressed size header.
    let uncompressed_size = u32::from_be_bytes(
        decrypted_data
            .get(0..4)
            .ok_or(ArqError::InvalidCompressedData)?
            .try_into()
            .unwrap(), // safe after get() succeeds
    );

    Ok(lz4_flex::decompress(
        &decrypted_data[4..],
        uncompressed_size as usize,
    )?)
}

/// NodeIter is a depth-first iterator over a tree of Nodes.
///
/// Traversal starts with "root" Node. New Nodes (via Trees) are loaded
/// lazily when needed. Items are yielded in lexical order, grouped by
/// files and directories (which first, depends on "directories_first").
pub struct NodeIter<'a> {
    repo: &'a Repository,
    directories_first: bool,

    queue: VecDeque<(Node, PathBuf)>,
}

impl<'a> NodeIter<'a> {
    pub fn new(repo: &'a Repository, root: Node, directories_first: bool) -> Self {
        NodeIter {
            repo,
            queue: vec![(root, PathBuf::from(""))].into(),
            directories_first,
        }
    }
}

impl<'a> Iterator for NodeIter<'a> {
    type Item = arq::Result<(Node, PathBuf)>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some((node, path)) = self.queue.pop_front() {
            if node.is_tree {
                let mut p = path.clone();
                p.push(node.name.clone());

                let Some(tree_loc) = node.tree_loc.as_ref() else {
                    return Some(Err(ArqError::MissingTreeLoc));
                };

                match self.repo.tree_from_pack(tree_loc) {
                    Ok(tree) => {
                        let (mut dirs, mut files): (Vec<Node>, Vec<Node>) =
                            tree.nodes.into_iter().partition(|node| node.is_tree);

                        dirs.sort_by(|a, b| b.name.cmp(&a.name));
                        files.sort_by(|a, b| b.name.cmp(&a.name));

                        for node in if self.directories_first {
                            files.into_iter().chain(dirs.into_iter())
                        } else {
                            dirs.into_iter().chain(files.into_iter())
                        } {
                            self.queue.push_front((node, p.clone()))
                        }
                    }
                    Err(e) => return Some(Err(e)),
                }
            }

            Some(Ok((node, path)))
        } else {
            None
        }
    }
}
