//! `cindy secret …` subcommand implementation.
//!
//! V1 model — file-based vault keys, à la `ansible-vault`'s
//! `--vault-password-file`. Each vault has a 32-byte data-encryption
//! key (DEK) stored at `keys/<name>.dek`. The directory should be
//! `.gitignore`d; the file gets distributed to teammates out-of-band.
//!
//! Subcommands:
//!
//!   * `cindy secret vault create <name>` — write a fresh random DEK to
//!     `keys/<name>.dek` (mode 0600).
//!   * `cindy secret seal`   — encrypt every `cindy::secret!(...)` in
//!     source, rewriting the file in-place. Auto-creates missing
//!     vaults with a warning.
//!   * `cindy secret unseal` — reverse: turn every
//!     `Secret::sealed_b64("vault", "...")` back into
//!     `cindy::secret!("vault", <original tokens>)`. Uses the source
//!     tokens that `seal` captured via `stringify!()` and embedded in
//!     the ciphertext, so you get back the exact value expression you
//!     originally wrote (post-`cargo fmt`).

use std::path::{Path, PathBuf};

use base64::Engine as _;
use cindy::secret::SealedPayload;
use cindy::secret::crypto;
use cindy::secret::keychain;
use eyre::{ContextCompat as _, WrapErr as _};

#[derive(Debug, clap::Subcommand)]
pub enum SecretCommand {
    /// Vault management.
    #[clap(subcommand)]
    Vault(VaultCommand),

    /// Encrypt every `cindy::secret!(...)` invocation in source.
    ///
    /// Compiles the orchestrator binary, runs it with
    /// `CINDY_SEAL_SECRETS=1` to collect (file, line, column, vault,
    /// ciphertext) records for every registered pending secret, then
    /// rewrites each source file in place by replacing the
    /// `secret!(...)` token tree with `cindy::Secret::sealed_b64(...)`.
    /// Missing vaults are auto-created with a warning.
    Seal,

    /// Reverse of `seal`: turn every `Secret::sealed_b64(...)` call
    /// in source back into the original `cindy::secret!(...)` macro
    /// invocation, decrypting via the local key files.
    Unseal,
}

#[derive(Debug, clap::Subcommand)]
pub enum VaultCommand {
    /// Generate a fresh 32-byte random DEK for `<name>` and write it
    /// to `keys/<name>.dek` (mode 0600). The directory should be
    /// `.gitignore`d. Errors if the file already exists.
    Create {
        /// Vault name. Lower-case, no slashes; appears verbatim as
        /// `keys/<name>.dek` and inside `secret!("<name>", ...)`.
        name: String,
    },
}

/// True iff the requested subcommand needs the orchestrator binary
/// compiled. Lets `main` skip the `cargo build` for the cheap
/// subcommands.
pub fn requires_orchestrator(cmd: &SecretCommand) -> bool {
    matches!(cmd, SecretCommand::Seal)
}

pub async fn dispatch(cmd: SecretCommand, orchestrator_path: Option<&Path>) -> eyre::Result<()> {
    match cmd {
        SecretCommand::Vault(VaultCommand::Create { name }) => vault_create(&name),
        SecretCommand::Seal => {
            let orch = orchestrator_path
                .expect("`Seal` requires an orchestrator path (see `requires_orchestrator`)");
            seal_all(orch).await
        }
        SecretCommand::Unseal => unseal_all(),
    }
}

// ---------------------------------------------------------------------
// vault create
// ---------------------------------------------------------------------

fn vault_create(name: &str) -> eyre::Result<()> {
    validate_vault_name(name)?;
    let path = keychain::dek_path(name);
    if path.exists() {
        eyre::bail!(
            "{} already exists. Refusing to overwrite an existing vault key. \
             If you really mean to rotate, delete the file by hand first.",
            path.display(),
        );
    }
    write_new_dek(&path)?;
    println!("created vault `{name}` at {}", path.display());
    println!(
        "(remember to add `{}/` to .gitignore)",
        keychain::keys_dir().display()
    );
    Ok(())
}

fn validate_vault_name(name: &str) -> eyre::Result<()> {
    if name.is_empty() {
        eyre::bail!("vault name must not be empty");
    }
    if name.contains(['/', '\\', '\0']) {
        eyre::bail!("vault name {name:?} contains invalid characters");
    }
    Ok(())
}

fn write_new_dek(path: &Path) -> eyre::Result<()> {
    use std::io::Write as _;

    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("Creating {}", parent.display()))?;
    }

    let dek = crypto::generate_dek();

    let mut opts = std::fs::OpenOptions::new();
    opts.write(true).create_new(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt as _;
        opts.mode(0o600);
    }
    let mut f = opts
        .open(path)
        .with_context(|| format!("Creating {} (refusing to overwrite)", path.display()))?;
    f.write_all(dek.as_slice())
        .with_context(|| format!("Writing DEK to {}", path.display()))?;
    Ok(())
}

// ---------------------------------------------------------------------
// seal: run orchestrator with CINDY_SEAL_SECRETS, rewrite source in place
// ---------------------------------------------------------------------

#[derive(Debug, serde::Deserialize)]
struct SealRecord {
    file: String,
    line: u32,
    column: u32,
    vault: String,
    ciphertext: String,
}

async fn seal_all(orchestrator_path: &Path) -> eyre::Result<()> {
    tracing::info!("Asking orchestrator to seal pending secrets...");

    let output = tokio::process::Command::new(orchestrator_path)
        .env("CINDY_SEAL_SECRETS", "1")
        .stdin(std::process::Stdio::null())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::inherit())
        .output()
        .await
        .context("Couldn't spawn orchestrator in seal mode")?;
    if !output.status.success() {
        eyre::bail!(
            "orchestrator seal mode exited with {:?} (see stderr above for the per-secret reason; \
             missing vaults are bootstrapped with `cindy secret vault create <name>`)",
            output.status
        );
    }

    let stdout =
        std::str::from_utf8(&output.stdout).context("orchestrator emitted non-UTF8 on stdout")?;
    let mut records: Vec<SealRecord> = Vec::new();
    for (i, line) in stdout.lines().enumerate() {
        if line.trim().is_empty() {
            continue;
        }
        let rec: SealRecord = serde_json::from_str(line).with_context(|| {
            format!(
                "orchestrator seal line {} is not valid JSON: {line:?}",
                i + 1
            )
        })?;
        records.push(rec);
    }
    if records.is_empty() {
        tracing::info!("No pending secrets found; nothing to seal.");
        return Ok(());
    }
    tracing::info!(count = records.len(), "Sealing pending secrets...");

    let mut by_file: std::collections::BTreeMap<String, Vec<SealRecord>> =
        std::collections::BTreeMap::new();
    for r in records {
        by_file.entry(r.file.clone()).or_default().push(r);
    }
    let mut affected = Vec::new();
    for (file, recs) in by_file {
        let path = PathBuf::from(&file);
        rewrite_file_seal(&path, &recs)?;
        affected.push(path);
    }

    cargo_fmt(&affected);
    Ok(())
}

/// Find each `secret!(...)` macro invocation matching a seal record
/// and splice in a `Secret::sealed_b64(...)` call.
fn rewrite_file_seal(path: &Path, records: &[SealRecord]) -> eyre::Result<()> {
    let original = std::fs::read_to_string(path)
        .with_context(|| format!("Reading source file {}", path.display()))?;

    let line_byte_starts = compute_line_starts(&original);

    let tokens: proc_macro2::TokenStream = original.parse().map_err(|e| {
        eyre::eyre!(
            "Tokenising {} while looking for `secret!` calls: {e}",
            path.display(),
        )
    })?;

    let mut macros: Vec<MacroSite> = Vec::new();
    find_secret_macros(&tokens, &original, &line_byte_starts, &mut macros);

    let mut planned: Vec<(usize, usize, String)> = Vec::new();
    for rec in records {
        let m = macros
            .iter()
            .find(|m| m.line == rec.line && m.column == rec.column)
            .with_context(|| {
                format!(
                    "Couldn't locate `secret!` macro at {}:{}:{} \u{2014} did the source \
                     change between seal-prepare and seal-apply?",
                    path.display(),
                    rec.line,
                    rec.column,
                )
            })?;
        let replacement = format!(
            "::cindy::Secret::sealed_b64({:?}, {:?})",
            rec.vault, rec.ciphertext
        );
        planned.push((m.byte_start, m.byte_end, replacement));
    }

    apply_planned_replacements(path, &original, &mut planned)?;
    tracing::info!(
        file = %path.display(),
        count = records.len(),
        "sealed secrets in source"
    );
    Ok(())
}

// ---------------------------------------------------------------------
// unseal: replace Secret::sealed_b64(...) with cindy::secret!(...)
// ---------------------------------------------------------------------

fn unseal_all() -> eyre::Result<()> {
    let src_root = Path::new("src");
    if !src_root.exists() {
        eyre::bail!(
            "no `src/` directory in the current working directory \u{2014} \
             run `cindy secret unseal` from your project root"
        );
    }
    let files = walk_rs_files(src_root);

    let mut affected = Vec::new();
    let mut total = 0usize;
    for file in &files {
        let count = unseal_file(file)?;
        if count > 0 {
            tracing::info!(
                file = %file.display(),
                count,
                "unsealed secrets in source"
            );
            affected.push(file.clone());
            total += count;
        }
    }

    if total == 0 {
        tracing::info!("No `Secret::sealed_b64(...)` calls found; nothing to unseal.");
        return Ok(());
    }

    cargo_fmt(&affected);
    tracing::info!(
        count = total,
        "unsealed total secrets (re-run `cindy secret seal` to put them back)"
    );
    Ok(())
}

fn unseal_file(path: &Path) -> eyre::Result<usize> {
    let original = std::fs::read_to_string(path)
        .with_context(|| format!("Reading source file {}", path.display()))?;
    let line_starts = compute_line_starts(&original);

    let tokens: proc_macro2::TokenStream = original.parse().map_err(|e| {
        eyre::eyre!(
            "Tokenising {} while looking for `Secret::sealed_b64` calls: {e}",
            path.display(),
        )
    })?;

    let mut sites: Vec<SealedSite> = Vec::new();
    find_sealed_calls(&tokens, &original, &line_starts, &mut sites);
    if sites.is_empty() {
        return Ok(0);
    }

    let mut planned: Vec<(usize, usize, String)> = Vec::new();
    for site in &sites {
        let dek = keychain::get_dek(&site.vault).map_err(|e| {
            eyre::eyre!(
                "{}:{}:{}: loading DEK for vault `{}`: {e}",
                path.display(),
                site.line,
                site.column,
                site.vault,
            )
        })?;
        let ciphertext = base64::engine::general_purpose::STANDARD
            .decode(site.ciphertext_b64.as_bytes())
            .with_context(|| {
                format!(
                    "{}:{}:{}: ciphertext for vault `{}` is not valid base64",
                    path.display(),
                    site.line,
                    site.column,
                    site.vault,
                )
            })?;
        let plaintext = crypto::unseal(&dek, &ciphertext).map_err(|e| {
            eyre::eyre!(
                "{}:{}:{}: decrypting vault `{}`: {e}",
                path.display(),
                site.line,
                site.column,
                site.vault,
            )
        })?;
        let payload: SealedPayload = postcard::from_bytes(&plaintext).map_err(|e| {
            eyre::eyre!(
                "{}:{}:{}: SealedPayload didn't deserialise for vault `{}`: {e}. \
                 This blob was probably produced by an older cindy version; \
                 rewrite it as `cindy::secret!(...)` by hand and re-run `cindy secret seal`.",
                path.display(),
                site.line,
                site.column,
                site.vault,
            )
        })?;

        let replacement = format!("::cindy::secret!({:?}, {})", site.vault, payload.source);
        planned.push((site.byte_start, site.byte_end, replacement));
    }

    apply_planned_replacements(path, &original, &mut planned)?;
    Ok(sites.len())
}

// ---------------------------------------------------------------------
// shared helpers
// ---------------------------------------------------------------------

fn apply_planned_replacements(
    path: &Path,
    original: &str,
    planned: &mut Vec<(usize, usize, String)>,
) -> eyre::Result<()> {
    // Reverse byte order so earlier offsets stay stable while we
    // splice tail-first.
    planned.sort_by_key(|(s, _, _)| std::cmp::Reverse(*s));

    let mut new = original.to_owned();
    for (start, end, replacement) in planned.iter() {
        new.replace_range(*start..*end, replacement);
    }

    if new != original {
        std::fs::write(path, &new)
            .with_context(|| format!("Writing rewritten source to {}", path.display()))?;
    }
    Ok(())
}

/// Recursively collect every `.rs` file under `root`, skipping
/// `target/` directories.
fn walk_rs_files(root: &Path) -> Vec<PathBuf> {
    let mut out = Vec::new();
    walk_recursive(root, &mut out);
    out
}

fn walk_recursive(dir: &Path, out: &mut Vec<PathBuf>) {
    let Ok(entries) = std::fs::read_dir(dir) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            if path.file_name().and_then(|s| s.to_str()) == Some("target") {
                continue;
            }
            walk_recursive(&path, out);
        } else if path.extension().map(|e| e == "rs").unwrap_or(false) {
            out.push(path);
        }
    }
}

/// Best-effort `cargo fmt -- <files>`. Failures (rustfmt missing, the
/// user has unparseable code we just touched, …) are logged but don't
/// stop the run — the rewrite itself already landed.
fn cargo_fmt(files: &[PathBuf]) {
    if files.is_empty() {
        return;
    }
    // `cargo fmt -- <file>` forwards `<file>` to rustfmt directly,
    // bypassing cargo's "format the whole package" logic so we only
    // touch the files we rewrote.
    let mut cmd = std::process::Command::new("cargo");
    cmd.arg("fmt").arg("--");
    for f in files {
        cmd.arg(f);
    }
    match cmd.status() {
        Ok(s) if s.success() => {
            tracing::info!(count = files.len(), "ran `cargo fmt` on affected files");
        }
        Ok(s) => {
            tracing::warn!(
                "`cargo fmt` exited with {s:?}; rewritten files may be poorly formatted"
            );
        }
        Err(e) => {
            tracing::warn!("couldn't spawn `cargo fmt` ({e}); rewritten files left as-is");
        }
    }
}

fn compute_line_starts(s: &str) -> Vec<usize> {
    let mut starts = vec![0usize];
    for (i, b) in s.bytes().enumerate() {
        if b == b'\n' {
            starts.push(i + 1);
        }
    }
    starts
}

fn line_col_to_byte(line_starts: &[usize], source: &str, line: u32, column: u32) -> Option<usize> {
    let line_idx = line.checked_sub(1)? as usize;
    let line_start = *line_starts.get(line_idx)?;
    let line_end = line_starts
        .get(line_idx + 1)
        .copied()
        .unwrap_or(source.len());
    let line_slice = source.get(line_start..line_end)?;

    let col_chars = column.checked_sub(1)? as usize;
    let mut chars_seen = 0;
    let mut byte_off = 0;
    for c in line_slice.chars() {
        if chars_seen == col_chars {
            return Some(line_start + byte_off);
        }
        chars_seen += 1;
        byte_off += c.len_utf8();
    }
    if chars_seen == col_chars {
        Some(line_start + byte_off)
    } else {
        None
    }
}

struct MacroSite {
    line: u32,
    column: u32,
    byte_start: usize,
    byte_end: usize,
}

struct SealedSite {
    line: u32,
    column: u32,
    byte_start: usize,
    byte_end: usize,
    vault: String,
    ciphertext_b64: String,
}

/// Walk `<path>secret!<group>` invocations including nested ones.
fn find_secret_macros(
    stream: &proc_macro2::TokenStream,
    source: &str,
    line_starts: &[usize],
    out: &mut Vec<MacroSite>,
) {
    use proc_macro2::{Spacing, TokenTree};

    let trees: Vec<TokenTree> = stream.clone().into_iter().collect();
    let mut i = 0;
    while i < trees.len() {
        if let TokenTree::Ident(ident) = &trees[i]
            && ident == "secret"
            && let Some(TokenTree::Punct(bang)) = trees.get(i + 1)
            && bang.as_char() == '!'
            && let Some(TokenTree::Group(grp)) = trees.get(i + 2)
        {
            let mut path_start = i;
            while path_start >= 3 {
                if let Some(TokenTree::Punct(p2)) = trees.get(path_start - 1)
                    && p2.as_char() == ':'
                    && let Some(TokenTree::Punct(p1)) = trees.get(path_start - 2)
                    && p1.as_char() == ':'
                    && p1.spacing() == Spacing::Joint
                    && let Some(TokenTree::Ident(_)) = trees.get(path_start - 3)
                {
                    path_start -= 3;
                } else {
                    break;
                }
            }
            // Optional leading `::` (the seal/unseal step emits it,
            // and `column!()` at the call site includes its position).
            if path_start >= 2 {
                if let Some(TokenTree::Punct(c2)) = trees.get(path_start - 1)
                    && c2.as_char() == ':'
                    && let Some(TokenTree::Punct(c1)) = trees.get(path_start - 2)
                    && c1.as_char() == ':'
                    && c1.spacing() == Spacing::Joint
                {
                    path_start -= 2;
                }
            }

            let start_lc = trees[path_start].span().start();
            let close_lc = grp.span_close().end();
            let line = start_lc.line as u32;
            let column = start_lc.column as u32 + 1;
            let byte_start = line_col_to_byte(line_starts, source, line, column).unwrap_or(0);
            let byte_end = line_col_to_byte(
                line_starts,
                source,
                close_lc.line as u32,
                close_lc.column as u32 + 1,
            )
            .unwrap_or(source.len());

            out.push(MacroSite {
                line,
                column,
                byte_start,
                byte_end,
            });
            i += 3;
            continue;
        }

        if let TokenTree::Group(grp) = &trees[i] {
            find_secret_macros(&grp.stream(), source, line_starts, out);
        }
        i += 1;
    }
}

/// Walk `<path>Secret::sealed_b64(<group>)` invocations including
/// nested ones. The group is expected to be `("<vault>", "<b64>")` —
/// any call that doesn't parse cleanly is silently skipped (probably
/// hand-written and not produced by `cindy secret seal`).
fn find_sealed_calls(
    stream: &proc_macro2::TokenStream,
    source: &str,
    line_starts: &[usize],
    out: &mut Vec<SealedSite>,
) {
    use proc_macro2::{Spacing, TokenTree};

    let trees: Vec<TokenTree> = stream.clone().into_iter().collect();
    let mut i = 0;
    while i < trees.len() {
        // Need at least: Secret :: sealed_b64 ( ... )  (5 tokens).
        if let TokenTree::Ident(secret_ident) = &trees[i]
            && secret_ident == "Secret"
            && let Some(TokenTree::Punct(p1)) = trees.get(i + 1)
            && p1.as_char() == ':'
            && p1.spacing() == Spacing::Joint
            && let Some(TokenTree::Punct(p2)) = trees.get(i + 2)
            && p2.as_char() == ':'
            && let Some(TokenTree::Ident(method_ident)) = trees.get(i + 3)
            && method_ident == "sealed_b64"
            && let Some(TokenTree::Group(grp)) = trees.get(i + 4)
        {
            // Walk back over any path prefix: `(ident ::)*` before
            // `Secret`. Any of `::cindy::Secret`, `cindy::Secret`,
            // `crate::Secret`, etc. is fine.
            let mut path_start = i;
            while path_start >= 3 {
                if let Some(TokenTree::Punct(c2)) = trees.get(path_start - 1)
                    && c2.as_char() == ':'
                    && let Some(TokenTree::Punct(c1)) = trees.get(path_start - 2)
                    && c1.as_char() == ':'
                    && c1.spacing() == Spacing::Joint
                    && let Some(TokenTree::Ident(_)) = trees.get(path_start - 3)
                {
                    path_start -= 3;
                } else {
                    break;
                }
            }
            // Leading `::` is also legal (the seal step emits it).
            // Each `::` is two Joint puncts; allow one extra step.
            if path_start >= 2 {
                if let Some(TokenTree::Punct(c2)) = trees.get(path_start - 1)
                    && c2.as_char() == ':'
                    && let Some(TokenTree::Punct(c1)) = trees.get(path_start - 2)
                    && c1.as_char() == ':'
                    && c1.spacing() == Spacing::Joint
                {
                    path_start -= 2;
                }
            }

            if let Some((vault, ciphertext_b64)) = parse_two_string_literals(&grp.stream()) {
                let start_lc = trees[path_start].span().start();
                let close_lc = grp.span_close().end();
                let line = start_lc.line as u32;
                let column = start_lc.column as u32 + 1;
                let byte_start = line_col_to_byte(line_starts, source, line, column).unwrap_or(0);
                let byte_end = line_col_to_byte(
                    line_starts,
                    source,
                    close_lc.line as u32,
                    close_lc.column as u32 + 1,
                )
                .unwrap_or(source.len());

                out.push(SealedSite {
                    line,
                    column,
                    byte_start,
                    byte_end,
                    vault,
                    ciphertext_b64,
                });
                i += 5;
                continue;
            }
        }

        if let TokenTree::Group(grp) = &trees[i] {
            find_sealed_calls(&grp.stream(), source, line_starts, out);
        }
        i += 1;
    }
}

/// Parse `("vault-literal", "ciphertext-literal")` out of a token
/// group. Returns `None` if the shape doesn't match — we'd rather
/// silently skip a stray `Secret::sealed_b64(...)` than crash.
fn parse_two_string_literals(stream: &proc_macro2::TokenStream) -> Option<(String, String)> {
    use syn::parse::Parser as _;

    let parser = syn::punctuated::Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated;
    let punct = parser.parse2(stream.clone()).ok()?;
    let args: Vec<&syn::Expr> = punct.iter().collect();
    if args.len() != 2 {
        return None;
    }
    let extract = |e: &syn::Expr| -> Option<String> {
        if let syn::Expr::Lit(lit) = e
            && let syn::Lit::Str(s) = &lit.lit
        {
            Some(s.value())
        } else {
            None
        }
    };
    Some((extract(args[0])?, extract(args[1])?))
}
