mod arq;
mod errors;
mod repository;
mod utils;

use crate::{
    arq::{
        BackupFolder, BackupPlan, BackupRecord, KeySet, NodeIter, decrypt_lz4decompress,
        decrypt_object,
    },
    errors::ArqError,
    repository::Repository,
    utils::{bytes_to_mb, to_hex_string},
};
use chrono::{DateTime, Local};
use clap::{Parser, Subcommand};
use inquire::{InquireError, Password, validator::Validation};
use log::{debug, trace};
use shlex::bytes;
use std::{
    fs,
    io::{IsTerminal, Read},
    path::{Path, PathBuf},
    process::Command,
    str::{self},
    time::Instant,
};

const ENCRYPTED_KEY_SET: &str = "encryptedkeyset.dat";

/// Restore files from backups created by the Arq backup program in version 6 or 7.
#[derive(Parser)]
#[command(version)]
struct Cli {
    /// Password to the repository
    #[arg(long)]
    password: Option<String>,
    /// Path to file containing the repository password (whitespace will be truncated)
    #[arg(long, short)]
    password_file: Option<PathBuf>,
    /// Shell command to obtain the repository password (whitespace will be truncated)
    #[arg(long)]
    password_command: Option<String>,

    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Subcommand)]
enum Commands {
    List {
        /// The path to the Arq backup repository
        #[arg(short, long)]
        repository: PathBuf,

        #[arg(short, long)]
        folder: Option<PathBuf>,

        #[arg(long)]
        record: Option<u32>,

        #[arg(long)]
        subtree: Option<PathBuf>,
    },
    /// Restore
    Restore {
        /// Path to the Arq backup repository
        #[arg(short, long)]
        repository: PathBuf,

        /// Backup folder to restore
        #[arg(short, long)]
        folder: PathBuf,

        /// Backup record to restore (as timestamp)
        #[arg(long)]
        record: u32,

        /// Subfolder within a backup folder to restore
        #[arg(long)]
        subtree: Option<PathBuf>,

        /// Location to restore to
        target: PathBuf,

        /// Don't actually create directory or files when restoring
        #[arg(long)]
        dry_run: bool,
    },
    /// Inspect various files/objects in the repository for debugging purposes
    Inspect {
        /// File to inspect
        file: PathBuf,

        /// Only decrypt the file provided and output plaintext
        #[arg(short, long)]
        raw: bool,
    },
}

fn main() {
    env_logger::init();

    let cli = Cli::parse();

    let password = ensure_password(&cli.password, &cli.password_file, &cli.password_command);

    match &cli.command {
        Some(Commands::List {
            repository,

            folder,
            record,
            subtree,
        }) => {
            if let Err(err) = cmd_list(
                password,
                repository,
                folder.clone(),
                *record,
                subtree.clone(),
            ) {
                eprintln!("Error: {}", err);
                std::process::exit(1);
            }
        }
        Some(Commands::Inspect { file, raw }) => {
            if let Err(err) = cmd_inspect(password, file, raw) {
                eprintln!("Error: {}", err);
                std::process::exit(1);
            }
        }
        Some(Commands::Restore {
            repository,
            folder,
            subtree,
            record,
            target,
            dry_run,
        }) => {
            if let Err(err) = cmd_restore(
                password,
                repository,
                folder.clone(),
                record,
                subtree.clone(),
                target,
                dry_run,
            ) {
                eprintln!("Error: {}", err);
                std::process::exit(1);
            }
        }
        None => {}
    }
}

fn ensure_password(
    password: &Option<String>,
    password_file: &Option<PathBuf>,
    password_command: &Option<String>,
) -> String {
    // direct password has highest priority
    if let Some(pw) = password {
        return pw.clone();
    }

    // next, check for password-file
    if let Some(file) = password_file {
        match fs::read_to_string(file) {
            Ok(p) => return p.trim_end().to_string(),
            Err(err) => {
                eprintln!("Error reading password file: {}", err);
                std::process::exit(1);
            }
        }
    }

    // next, check for password-command
    if let Some(command) = password_command {
        let args: Vec<String> = bytes::Shlex::new(command.as_bytes())
            .map(|arg| String::from_utf8(arg).unwrap())
            .collect();

        let program = &args[0];
        let args = &args[1..];

        match Command::new(program).args(args).output() {
            Ok(output) => {
                if output.status.success()
                    && let Ok(password) = String::from_utf8(output.stdout)
                {
                    return password.trim().to_string();
                } else {
                    eprintln!("Password command exited with non-zero {}", output.status);
                }
            }
            Err(err) => {
                eprintln!("Password command {:?} errored: {err}", command);
            }
        }
        std::process::exit(1)
    }

    // next, ask for the password, if STDIN is a TTY
    if !std::io::stdin().is_terminal() {
        eprintln!("Non-interactive session and no --password or --password-file provided!");
        std::process::exit(1);
    }
    match Password::new("Backup encryption password: ")
        .without_confirmation()
        .with_validator(|input: &str| {
            if input.is_empty() {
                Ok(Validation::Invalid("Password must be non-empty".into()))
            } else {
                Ok(Validation::Valid)
            }
        })
        .prompt()
    {
        Ok(password) => password,
        Err(InquireError::OperationCanceled | InquireError::OperationInterrupted) => {
            eprintln!("Aborted");
            std::process::exit(1)
        }
        Err(err) => {
            println!("An error happened when asking for your key, try again later: {err}");
            std::process::exit(1)
        }
    }
}

fn ensure_backup_folder_uuid(input: &Path) -> String {
    match input.to_str() {
        Some(path) => path.to_string(),
        None => panic!("backup folder is not a valid UTF-8 string!?"),
    }
}

fn cmd_list(
    password: String,
    repository: &Path,
    folder: Option<PathBuf>,
    record: Option<u32>,
    subtree: Option<PathBuf>,
) -> arq::Result<()> {
    let repo = Repository::build(repository, &password).unwrap_or_else(|err| {
        match err {
            err @ ArqError::KeySetMacError(_) => {
                eprintln!("{}", err)
            }
            _ => {
                eprintln!("Not a Arq backup repository/plan? {}", err);
            }
        }

        std::process::exit(1)
    });

    let plan = repo.load_plan().unwrap_or_else(|err| {
        eprintln!("Could not load backup plan: {err}");
        std::process::exit(1);
    });
    println!("Backup plan: {} ({})", plan.name, plan.uuid);

    if let Some(folder) = folder {
        if let Ok(backup_folder) = repo.load_folder(&ensure_backup_folder_uuid(&folder)) {
            println!("Backup folder: {}", backup_folder.local_path);

            if let Some(record) = record {
                let rec = repo.load_record(&backup_folder.uuid, &record)?;

                println!("Backup record: {} ({})", rec.creation_time, record);
                println!();

                let backup_local_path = PathBuf::from(&backup_folder.local_path);

                let (base_path, start_node) = if let Some(subtree) = subtree {
                    let (start_node, p) =
                        match repo.find_subtree(&backup_local_path, rec.root, &subtree) {
                            Ok(Some((n, p))) => (n, p),
                            Ok(None) => {
                                eprintln!(
                                    "Could not find directory {} in backup!",
                                    subtree.display()
                                );
                                std::process::exit(1);
                            }
                            Err(e) => {
                                eprintln!("An error occurred while looking up sub tree: {e}",);
                                std::process::exit(1);
                            }
                        };

                    (backup_local_path.join(p), start_node)
                } else {
                    (backup_local_path, rec.root)
                };

                for n in NodeIter::new(&repo, start_node, true) {
                    let (node, mut path) = n?;
                    path.push(node.name.clone());

                    if node.is_tree {
                        println!("{}/", base_path.join(path).display());
                    } else {
                        println!("{}", base_path.join(path).display());
                    }
                }
            } else {
                println!(
                    "Backup records for folder \"{}\" ({}):\n",
                    backup_folder.local_path, backup_folder.uuid
                );

                for record_time in repo
                    .backup_record_timestamps(&backup_folder.uuid)?
                    .iter()
                    .rev()
                {
                    println!(
                        "{:?} ({})",
                        DateTime::from_timestamp_secs(*record_time as i64)
                            .expect("Cannot parse backup record timestamp!")
                            .with_timezone(&Local),
                        record_time
                    );
                }
            }
        } else {
            println!("Backup folder {} not found in repository", folder.display());
            std::process::exit(1);
        }
    } else {
        println!("Backup folders in {}:\n", repo.uuid);
        repo.backup_folders()?.iter().for_each(|folder| {
            println!("{}: {}", folder.uuid, folder.local_path);
        });
    }

    Ok(())
}

fn cmd_inspect(password: String, file: &Path, raw: &bool) -> arq::Result<()> {
    if !file.exists() {
        println!("{:?} not found", file);
        std::process::exit(1);
    }

    if file.is_dir() {
        println!("Currently only files can be inspected");
        std::process::exit(1);
    }

    fn find_keyset_file(file: &Path) -> Option<PathBuf> {
        if file.ends_with(ENCRYPTED_KEY_SET) {
            return Some(file.to_path_buf());
        }

        let mut current = file.parent()?;
        loop {
            if current.join(ENCRYPTED_KEY_SET).exists() {
                return Some(current.join(ENCRYPTED_KEY_SET));
            }

            current = current.parent()?;
        }
    }

    let keyset_path = match find_keyset_file(file) {
        Some(path) => path,
        None => {
            println!("Could not find encryptedkeyset.dat containing the encryption keys.");
            std::process::exit(1);
        }
    };

    debug!("Found encryptedkeyset.dat at {:?}", keyset_path);

    let key_set = KeySet::decrypt_parse_key_set(fs::File::open(keyset_path)?, &password)?;

    let mut fh = fs::File::open(file)?;

    if file.ends_with("backupplan.json") {
        if *raw {
            let json = decrypt_object(fh, &key_set)?;
            println!("{}", str::from_utf8(&json)?);
        } else {
            println!("Type: backup plan");
            let obj = BackupPlan::from_reader(fh, &key_set)?;
            println!("{}", obj);
        }
    } else if file.ends_with("backupfolder.json") {
        if *raw {
            let json = decrypt_object(fh, &key_set)?;
            println!("{}", str::from_utf8(&json)?);
        } else {
            println!("Type: backup folder");
            let obj = BackupFolder::from_reader(fh, &key_set)?;
            println!("{}", obj);
        }
    } else if file.extension().unwrap_or_default() == "backuprecord" {
        if *raw {
            let json = decrypt_lz4decompress(fh, &key_set)?;
            println!("{}", str::from_utf8(&json)?);
        } else {
            println!("Type: backup record");
            let obj = BackupRecord::from_reader(fh, &key_set)?;
            println!("{}", obj);
        }
    } else if file.extension().unwrap_or_default() == "snapshot" {
        let json = decrypt_lz4decompress(fh, &key_set)?;
        println!("{}", str::from_utf8(&json)?);
    } else if file.ends_with(ENCRYPTED_KEY_SET) {
        println!("Type: key set");
        println!(
            "Master encryption key (hex): {}",
            to_hex_string(&key_set.enc_key)
        );
        println!(
            "Master HMAC key (hex): {}",
            to_hex_string(&key_set.hmac_key)
        );
    } else if file.ends_with("backupfolders.json") {
        let mut buf = Vec::new();
        fh.read_to_end(&mut buf)?;

        println!("{}", str::from_utf8(&buf)?);
    } else {
        eprintln!("Unknown file/object type!");
        std::process::exit(1);
    }

    Ok(())
}

fn cmd_restore(
    password: String,
    repository: &Path,
    folder: PathBuf,
    record: &u32,
    subtree: Option<PathBuf>,
    target: &Path,
    dry_run: &bool,
) -> arq::Result<()> {
    let repo = Repository::build(repository, &password).unwrap_or_else(|err| {
        match err {
            err @ ArqError::KeySetMacError(_) => {
                eprintln!("{}", err)
            }
            _ => {
                eprintln!("Not a Arq backup repository/plan? {}", err);
            }
        }

        std::process::exit(1)
    });

    let plan = repo.load_plan().unwrap_or_else(|err| {
        eprintln!("Could not load backup plan {}", err);
        std::process::exit(1)
    });

    let backup_folder = repo
        .load_folder(&ensure_backup_folder_uuid(&folder))
        .unwrap_or_else(|err| {
            eprintln!("Error: {}", err);
            std::process::exit(1)
        });

    let backup_record = repo.load_record(&backup_folder.uuid, record)?;

    println!("Backup plan: {} ({})", plan.name, plan.uuid);
    println!(
        "Backup folder: {} ({})",
        backup_folder.local_path, backup_folder.uuid
    );
    println!(
        "Backup record: {} ({})",
        backup_record.creation_time, record
    );

    if let Some(arq_version) = backup_record.arq_version
        && arq_version.split(".").collect::<Vec<&str>>()[0]
            .to_string()
            .parse::<u32>()
            .expect("non-integer in Arq version")
            < 6
    {
        eprintln!(
            "Unsupported Arq backup record version ({}). Only Arq version 6+ is supported!",
            arq_version
        );
        std::process::exit(1);
    }

    let local_path = PathBuf::from(&backup_folder.local_path);

    let (start_node, base_path) = if let Some(subtree) = subtree {
        match repo.find_subtree(&local_path, backup_record.root, &subtree) {
            Ok(Some((n, p))) => (n, p),
            Ok(None) => {
                eprintln!("Could not find directory {} in backup!", subtree.display());
                std::process::exit(1);
            }
            Err(e) => {
                eprintln!("An error occurred while looking up sub tree: {}", e);
                std::process::exit(1);
            }
        }
    } else {
        (backup_record.root, local_path)
    };

    println!();
    println!("Restoring {} to {}", base_path.display(), target.display());
    println!(
        "Size : {mb:.2} MB ({bytes} Bytes)",
        mb = bytes_to_mb(start_node.item_size),
        bytes = start_node.item_size
    );
    println!("Files: {}", start_node.contained_files_count);
    println!();

    let mut restored_bytes = 0;
    let mut restored_files = 0;
    let mut restored_dirs = 0;
    let started = Instant::now();

    let mut failures: Vec<(PathBuf, ArqError)> = vec![];

    for n in NodeIter::new(&repo, start_node, true) {
        let (node, mut path) = n?;
        path.push(node.name.clone());

        let backup_path = base_path.join(&path);
        let restore_path = target.join(&path);

        println!(
            "{} -> {}",
            backup_path.to_string_lossy(),
            restore_path.to_string_lossy()
        );

        match node.restore(&repo, &restore_path, *dry_run) {
            Ok(_hash) => {
                if !node.is_tree {
                    restored_files += 1;
                    restored_bytes += node.item_size;
                } else {
                    restored_dirs += 1;
                }
                // println!("{}: {}", path.to_string_lossy(), to_hex_string(&hash));
                // println!("{}: {}", path.to_string_lossy(), node);
            }
            Err(e) => {
                eprintln!("Could not restore {}: {e}", restore_path.to_string_lossy());

                failures.push((backup_path, e));
            }
        }
    }

    println!(
        "\nDONE!\n\nRead {read_bytes:.2} MB to restore {restored_bytes:.2} MB in {restored_files} files across {restored_dirs} directories. Took {duration:.3} sec.",
        read_bytes = bytes_to_mb(repo.used_data_pack_bytes()),
        restored_bytes = bytes_to_mb(restored_bytes),
        duration = started.elapsed().as_secs_f32()
    );

    if !failures.is_empty() {
        println!("Some errors occurred while restoring!\n\n");

        for (path, err) in failures {
            println!("{path}: {err}", path = path.to_string_lossy())
        }
    }

    debug!("Used tree packs: {}", repo.used_tree_pack_files().len());
    trace!("Used tree packs: {:?}", repo.used_tree_pack_files());
    debug!("Used tree packs: {}", repo.used_data_pack_files().len());
    trace!("User data packs: {:?}", repo.used_data_pack_files());

    Ok(())
}
