Documentation
use anyhow::{Context, Result};
use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
use pbkdf2::pbkdf2_hmac;
use rand::RngCore;
use sha2::Sha256;
use std::{
    fs,
    io::{self, Cursor, Write},
    path::{Path, PathBuf},
};
use walkdir::WalkDir;

pub const MAX_FILE_SIZE: u64 = 1000 * 1024 * 1024;
const ENCRYPT_MAGIC: &[u8] = b"XTOOLENC1";
const SALT_LEN: usize = 16;
const NONCE_LEN: usize = 12;
const PBKDF2_ITERS: u32 = 100_000;
pub const XTOOL_FILE_SUFFIX: &str = ".xtool_file";
pub const XTOOL_DIR_SUFFIX: &str = ".xtool_dir";

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ArchiveHint {
    File,
    Dir,
    None,
}

pub fn compress_directory(dir: &Path) -> Result<(PathBuf, String, u64)> {
    if !dir.exists() || !dir.is_dir() {
        return Err(anyhow::anyhow!("Directory not found: {}", dir.display()));
    }

    let base_name = dir
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or("archive");
    let zip_name = format!("{}{}", strip_xtool_suffix(base_name), XTOOL_DIR_SUFFIX);

    let tmp = tempfile::Builder::new()
        .prefix("xtool_upload_")
        .suffix(".zip")
        .tempfile()
        .context("Failed to create temp file")?;
    let mut writer = zip::ZipWriter::new(tmp.as_file());
    let options = zip::write::FileOptions::<()>::default()
        .compression_method(zip::CompressionMethod::Deflated)
        .unix_permissions(0o644);

    let base = dir.canonicalize().context("Failed to canonicalize path")?;

    for entry in WalkDir::new(&base) {
        let entry = entry.context("Failed to walk directory")?;
        let path = entry.path();
        let rel = path
            .strip_prefix(&base)
            .context("Failed to compute relative path")?;
        let name = rel.to_string_lossy().replace('\\', "/");
        if name.is_empty() {
            continue;
        }

        if path.is_dir() {
            writer
                .add_directory(name, options)
                .context("Failed to add directory to archive")?;
        } else if path.is_file() {
            writer
                .start_file(name, options)
                .context("Failed to add file to archive")?;
            let mut file = fs::File::open(path)
                .with_context(|| format!("Failed to open file: {}", path.display()))?;
            io::copy(&mut file, &mut writer)
                .context("Failed to write file to archive")?;
        }
    }

    writer.finish().context("Failed to finalize archive")?;
    tmp.as_file().sync_all().ok();

    let (file, path) = tmp.keep().context("Failed to keep temp file")?;
    let size = file
        .metadata()
        .context("Failed to read archive metadata")?
        .len();
    drop(file);

    Ok((path, zip_name, size))
}

pub fn compress_file(file_path: &Path) -> Result<(PathBuf, String, u64)> {
    if !file_path.exists() || !file_path.is_file() {
        return Err(anyhow::anyhow!("File not found: {}", file_path.display()));
    }

    let file_name = file_path
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or("file.bin")
        .to_string();
    let clean_name = strip_xtool_suffix(&file_name);
    let zip_name = format!("{}{}", clean_name, XTOOL_FILE_SUFFIX);

    let tmp = tempfile::Builder::new()
        .prefix("xtool_upload_")
        .suffix(".zip")
        .tempfile()
        .context("Failed to create temp file")?;

    let mut writer = zip::ZipWriter::new(tmp.as_file());
    let options = zip::write::FileOptions::<()>::default()
        .compression_method(zip::CompressionMethod::Deflated)
        .unix_permissions(0o644);

    writer
        .start_file(&file_name, options)
        .context("Failed to add file to archive")?;

    let mut file = fs::File::open(file_path)
        .with_context(|| format!("Failed to open file: {}", file_path.display()))?;
    io::copy(&mut file, &mut writer).context("Failed to write file to archive")?;

    writer.finish().context("Failed to finalize archive")?;
    tmp.as_file().sync_all().ok();

    let (file, path) = tmp.keep().context("Failed to keep temp file")?;
    let size = file
        .metadata()
        .context("Failed to read archive metadata")?
        .len();
    drop(file);

    Ok((path, zip_name, size))
}

pub fn compress_path(path: &Path) -> Result<(PathBuf, String, u64)> {
    if path.is_dir() {
        compress_directory(path)
    } else {
        compress_file(path)
    }
}

pub fn write_temp_zip(bytes: &[u8]) -> Result<PathBuf> {
    let mut tmp = tempfile::Builder::new()
        .prefix("xtool_download_")
        .suffix(".zip")
        .tempfile()
        .context("Failed to create temp file")?;
    tmp.write_all(bytes)
        .context("Failed to write temp archive")?;
    let (_file, path) = tmp.keep().context("Failed to keep temp file")?;
    Ok(path)
}

pub fn detect_archive_hint(filename: &str) -> (String, ArchiveHint) {
    if let Some(stripped) = filename.strip_suffix(XTOOL_FILE_SUFFIX) {
        return (stripped.to_string(), ArchiveHint::File);
    }
    if let Some(stripped) = filename.strip_suffix(XTOOL_DIR_SUFFIX) {
        return (stripped.to_string(), ArchiveHint::Dir);
    }
    (filename.to_string(), ArchiveHint::None)
}

pub fn unzip_single_from_bytes(bytes: &[u8], output_path: &Path) -> Result<()> {
    let cursor = Cursor::new(bytes);
    let mut archive = zip::ZipArchive::new(cursor).context("Failed to read archive")?;
    if archive.len() == 0 {
        return Err(anyhow::anyhow!("Archive is empty"));
    }
    let mut entry = archive.by_index(0).context("Failed to read archive entry")?;

    if let Some(parent) = output_path.parent() {
        if !parent.as_os_str().is_empty() {
            fs::create_dir_all(parent)
                .with_context(|| format!("Failed to create directory: {}", parent.display()))?;
        }
    }

    let mut outfile = fs::File::create(output_path)
        .with_context(|| format!("Failed to create file: {}", output_path.display()))?;
    io::copy(&mut entry, &mut outfile).context("Failed to extract file")?;
    Ok(())
}

pub fn encrypt_zip_file(zip_path: &Path, key: &str) -> Result<u64> {
    let bytes = fs::read(zip_path)
        .with_context(|| format!("Failed to read archive: {}", zip_path.display()))?;
    let encrypted = encrypt_zip_bytes(&bytes, key)?;
    fs::write(zip_path, &encrypted)
        .with_context(|| format!("Failed to write encrypted archive: {}", zip_path.display()))?;
    Ok(encrypted.len() as u64)
}

pub fn is_encrypted_zip(bytes: &[u8]) -> bool {
    bytes.len() > ENCRYPT_MAGIC.len() + SALT_LEN + NONCE_LEN
        && bytes.starts_with(ENCRYPT_MAGIC)
}

pub fn decrypt_zip_bytes(bytes: &[u8], key: &str) -> Result<Vec<u8>> {
    if !is_encrypted_zip(bytes) {
        return Err(anyhow::anyhow!("Archive is not encrypted"));
    }

    let header_len = ENCRYPT_MAGIC.len();
    let salt_start = header_len;
    let salt_end = salt_start + SALT_LEN;
    let nonce_start = salt_end;
    let nonce_end = nonce_start + NONCE_LEN;

    let salt = &bytes[salt_start..salt_end];
    let nonce = &bytes[nonce_start..nonce_end];
    let ciphertext = &bytes[nonce_end..];

    let mut key_bytes = [0u8; 32];
    pbkdf2_hmac::<Sha256>(key.as_bytes(), salt, PBKDF2_ITERS, &mut key_bytes);

    let cipher = Aes256Gcm::new_from_slice(&key_bytes)
        .context("Failed to initialize cipher")?;
    let nonce = Nonce::from_slice(nonce);
    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| anyhow::anyhow!("Decrypt failed (bad key or corrupted data)"))
}

fn encrypt_zip_bytes(bytes: &[u8], key: &str) -> Result<Vec<u8>> {
    let mut salt = [0u8; SALT_LEN];
    let mut rng = rand::rng();
    rng.fill_bytes(&mut salt);

    let mut key_bytes = [0u8; 32];
    pbkdf2_hmac::<Sha256>(key.as_bytes(), &salt, PBKDF2_ITERS, &mut key_bytes);

    let cipher = Aes256Gcm::new_from_slice(&key_bytes)
        .context("Failed to initialize cipher")?;

    let mut nonce_bytes = [0u8; NONCE_LEN];
    rng.fill_bytes(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);

    let mut out = Vec::with_capacity(ENCRYPT_MAGIC.len() + SALT_LEN + NONCE_LEN + bytes.len() + 16);
    out.extend_from_slice(ENCRYPT_MAGIC);
    out.extend_from_slice(&salt);
    out.extend_from_slice(&nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, bytes)
        .map_err(|_| anyhow::anyhow!("Encrypt failed"))?;
    out.extend_from_slice(&ciphertext);
    Ok(out)
}

pub fn unzip_to_dir(zip_path: &Path, output_dir: &Path) -> Result<()> {
    let file = fs::File::open(zip_path)
        .with_context(|| format!("Failed to open archive: {}", zip_path.display()))?;
    let mut archive = zip::ZipArchive::new(file).context("Failed to read archive")?;

    fs::create_dir_all(output_dir)
        .with_context(|| format!("Failed to create directory: {}", output_dir.display()))?;

    for i in 0..archive.len() {
        let mut entry = archive.by_index(i).context("Failed to read archive entry")?;
        let out_path = output_dir.join(entry.name());

        if entry.name().ends_with('/') {
            fs::create_dir_all(&out_path)
                .with_context(|| format!("Failed to create directory: {}", out_path.display()))?;
        } else {
            if let Some(parent) = out_path.parent() {
                fs::create_dir_all(parent).with_context(|| {
                    format!("Failed to create directory: {}", parent.display())
                })?;
            }
            let mut outfile = fs::File::create(&out_path)
                .with_context(|| format!("Failed to create file: {}", out_path.display()))?;
            io::copy(&mut entry, &mut outfile)
                .context("Failed to extract file")?;
        }
    }

    Ok(())
}

pub fn resolve_output_path(output: Option<&Path>, filename: &str) -> PathBuf {
    match output {
        Some(path) if path.exists() && path.is_dir() => path.join(filename),
        Some(path) => path.to_path_buf(),
        None => PathBuf::from(filename),
    }
}

pub fn resolve_output_dir(output: Option<&Path>, filename: &str) -> Result<PathBuf> {
    if let Some(path) = output {
        if path.exists() && path.is_file() {
            return Err(anyhow::anyhow!(
                "Output path must be a directory for archives"
            ));
        }
        return Ok(path.to_path_buf());
    }

    let (mut stem, _) = detect_archive_hint(filename);
    stem = stem.trim_end_matches(".zip").to_string();
    if stem.is_empty() {
        Ok(PathBuf::from("xtool_download"))
    } else {
        Ok(PathBuf::from(stem))
    }
}

fn strip_xtool_suffix(name: &str) -> &str {
    if let Some(stripped) = name.strip_suffix(XTOOL_FILE_SUFFIX) {
        return stripped;
    }
    if let Some(stripped) = name.strip_suffix(XTOOL_DIR_SUFFIX) {
        return stripped;
    }
    name
}