cindy-macros 0.2.1

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{FnArg, Item, ItemFn, Pat, PatIdent, parse_macro_input};

/// Given an argument type, if it's `impl Into<T>` return `Some(T)`; otherwise
/// `None` (the type is used as-is for the concrete raw fn).
fn into_target(ty: &syn::Type) -> Option<syn::Type> {
    let syn::Type::ImplTrait(it) = ty else {
        return None;
    };
    for bound in &it.bounds {
        let syn::TypeParamBound::Trait(tb) = bound else {
            continue;
        };
        let seg = tb.path.segments.last()?;
        if seg.ident != "Into" {
            continue;
        }
        if let syn::PathArguments::AngleBracketed(args) = &seg.arguments
            && let Some(syn::GenericArgument::Type(inner)) = args.args.first()
        {
            return Some(inner.clone());
        }
    }
    None
}

/// Shorthand for deriving `Debug` and `serde::{Deserialize, Serialize}`
/// without an explicit `serde` dependency.
///
/// `Debug` is included because inventory `vars` types need it (so
/// `cindy inventory` can render them). Don't also list `Debug` in a
/// manual `#[derive(...)]` on the same item — that's a duplicate impl.
#[proc_macro_attribute]
pub fn wire(_args: TokenStream, input: TokenStream) -> TokenStream {
    let item = parse_macro_input!(input as Item);

    quote! {
        #[derive(
            ::std::fmt::Debug,
            cindy::__reexports::serde::Serialize,
            cindy::__reexports::serde::Deserialize
        )]
        #[serde(crate = "cindy::__reexports::serde")]
        #item
    }
    .into()
}

/// Specifies an entrypoint.
///
/// - On `feature = "orchestrator"`
///   - Compiles to normal program.
/// - On `feature = "remote"`
///   - Runs a RPC loop instead.
///
/// The user's `main` may be declared in one of two shapes:
///
/// 1. `async fn main() -> Result<()>` - legacy, no host context.
/// 2. `async fn main(host: cindy::Host<V>) -> Result<()>` - the macro
///    reads the `CINDY_HOST_CONTEXT` env var (JSON serialized
///    `cindy::Host<V>` produced by the CLI for this particular target)
///    and hands it pre-parsed to the body. Any `Serialize +
///    DeserializeOwned` `V` works; use `cindy::Host<()>` if you only
///    need `name`/`tags`.
#[proc_macro_attribute]
pub fn main(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);

    let (attrs, vis, sig, block) = (&input.attrs, &input.vis, &input.sig, &input.block);
    let output = &sig.output;
    let inputs = &sig.inputs;

    let invoke_user_main = match inputs.len() {
        0 => quote! {
            cindy::__reexports::tokio::spawn(__user_main())
        },
        1 => {
            let host_type = match &inputs[0] {
                FnArg::Typed(pt) => &pt.ty,
                FnArg::Receiver(_) => {
                    panic!("`#[cindy::main]` cannot take `self`");
                }
            };
            quote! {
                {
                    let __host_json = ::std::env::var("CINDY_HOST_CONTEXT").expect(
                        "CINDY_HOST_CONTEXT not set!\n\
                         The orchestrator process is meant to be launched by `cindy` command line tool. \
                         If you're trying to run the binary directly, set CINDY_HOST_CONTEXT to a \
                         JSON-serialised `cindy::Host<V>` first."
                    );
                    let __host: #host_type =
                        cindy::__reexports::serde_json::from_str(&__host_json)
                            .expect("CINDY_HOST_CONTEXT was not valid JSON for the declared `cindy::Host<V>` type");
                    cindy::__reexports::tokio::spawn(__user_main(__host))
                }
            }
        }
        _ => panic!(
            "`#[cindy::main]` accepts at most one parameter (the host context, `cindy::Host<V>`)"
        ),
    };

    quote! {
        #(#attrs)*
        #[cindy::__reexports::tokio::main(crate = "cindy::__reexports::tokio")]
        #vis async fn main() {
            let (rpc_in, rpc_out) = cindy::common::quarantine_stdio();

            #[cfg(feature = "orchestrator")]
            if ::std::env::var_os("CINDY_DUMP_INVENTORY").is_some() {
                let entries: ::std::vec::Vec<&cindy::inventory::RegisteredInventory> =
                    cindy::__reexports::inventory::iter::<cindy::inventory::RegisteredInventory>
                        .into_iter()
                        .collect();
                let dump: cindy::inventory::InventoryDump = match entries.as_slice() {
                    [one] => (one.dump)().await,
                    _ => {
                        ::std::eprintln!(
                            "There must be exactly 1 `#[cindy::inventory]` registered."
                        );
                        ::std::process::exit(2);
                    }
                };
                let bytes = cindy::__reexports::serde_json::to_vec(&dump)
                    .expect("Failed to serialize inventory dump to JSON");
                {
                    use cindy::__reexports::tokio::io::AsyncWriteExt as _;
                    let mut out = rpc_out;
                    out.write_all(&bytes).await.expect("Failed to write inventory");
                    out.flush().await.expect("Failed to flush inventory");
                }
                ::std::process::exit(0);
            }

            // Vault enumeration: print the set of vaults referenced by
            // `secret!` invocations compiled into this binary (the
            // "secrets in code" half of the preflight) as a JSON array,
            // then exit. The CLI runs this once during preflight to
            // know which DEKs every participating machine must hold.
            #[cfg(feature = "orchestrator")]
            if ::std::env::var_os("CINDY_DUMP_VAULTS").is_some() {
                use cindy::__reexports::tokio::io::AsyncWriteExt as _;
                let vaults = cindy::secret::registered_vaults();
                let bytes = cindy::__reexports::serde_json::to_vec(&vaults)
                    .expect("Failed to serialise vault list");
                let mut out = rpc_out;
                out.write_all(&bytes).await.expect("Failed to write vault list");
                out.flush().await.expect("Failed to flush vault list");
                ::std::process::exit(0);
            }

            #[cfg(feature = "orchestrator")]
            if ::std::env::var_os("CINDY_SEAL_SECRETS").is_some() {
                use cindy::__reexports::tokio::io::AsyncWriteExt as _;
                let mut out = rpc_out;
                let mut failed = false;
                for pending in cindy::__reexports::inventory::iter::<cindy::secret::PendingSecret>() {
                    let plaintext = (pending.serialize)();
                    // Refuse to bootstrap a vault on the fly. If the
                    // user didn't `cindy secret vault create <name>`
                    // before running seal, that's almost certainly a
                    // mistake — they probably meant to copy an
                    // existing team-shared key file into place. Bail
                    // with the (already actionable) error from the
                    // keychain layer instead of generating a fresh
                    // key whose ciphertext nobody else can decrypt.
                    let dek = match cindy::secret::keychain::get_dek(pending.vault) {
                        Ok(d) => d,
                        Err(e) => {
                            ::std::eprintln!(
                                "cindy secret seal: couldn't load DEK for vault `{}` \
                                 (referenced from {}:{}:{}): {e:#}",
                                pending.vault, pending.file, pending.line, pending.column,
                            );
                            failed = true;
                            continue;
                        }
                    };
                    let ciphertext = match cindy::secret::crypto::seal(&dek, &plaintext) {
                        Ok(c) => c,
                        Err(e) => {
                            ::std::eprintln!(
                                "cindy secret seal: encryption failed for {}:{}:{} ({e:#})",
                                pending.file, pending.line, pending.column,
                            );
                            failed = true;
                            continue;
                        }
                    };
                    use cindy::__reexports::base64::Engine as _;
                    let b64 = cindy::__reexports::base64::engine::general_purpose::STANDARD
                        .encode(&ciphertext);
                    let line = cindy::__reexports::serde_json::json!({
                        "file":       pending.file,
                        "line":       pending.line,
                        "column":     pending.column,
                        "vault":      pending.vault,
                        "ciphertext": b64,
                    });
                    let mut bytes = cindy::__reexports::serde_json::to_vec(&line)
                        .expect("Failed to serialise seal record");
                    bytes.push(b'\n');
                    out.write_all(&bytes).await.expect("Failed to write seal record");
                }
                out.flush().await.expect("Failed to flush seal records");
                ::std::process::exit(if failed { 2 } else { 0 });
            }

            // Remote-only builds skip the user's body entirely; the worker
            // process just runs the RPC dispatch loop. Vault DEKs arrive
            // over the channel's first frame (the handshake) and the
            // worker-side vault preflight runs there, inside
            // `cindy::remote::rpc` — keys never touch this process's
            // argv/env, so they're invisible in the target's `ps`.
            #[cfg(all(feature = "remote", not(feature = "orchestrator")))]
            {
                cindy::remote::rpc(rpc_in, rpc_out).await;
                ::std::process::exit(0);
            }

            // Orchestrator builds (and dual-feature LSP builds) set up the
            // dispatch channel and then execute the user's main body. Using a
            // *positive* cfg here — rather than `not(remote)` — keeps the
            // body in scope when both features are enabled at once, which is
            // how rust-analyzer typically evaluates this workspace.
            #[cfg(feature = "orchestrator")]
            {
                async fn __user_main(#inputs) #output #block

                // The CLI marshals this run's vault DEKs into
                // `CINDY_VAULT_KEYS` on *this* (local) process — env is
                // fine here since it's the operator's own machine, not a
                // target. We install them locally for orchestrator-side
                // reveals, and forward the same map to the worker over
                // the RPC handshake (so they never hit the target's env).
                let __cindy_vault_keys = match cindy::secret::keychain::decode_env_keys() {
                    ::std::result::Result::Ok(m) => m.unwrap_or_default(),
                    ::std::result::Result::Err(e) => {
                        ::std::eprintln!("\x1b[31m{:?}\x1b[0m", e);
                        ::std::process::exit(1);
                    }
                };
                if let ::std::result::Result::Err(e) =
                    cindy::secret::keychain::install_raw_keys(__cindy_vault_keys.clone())
                {
                    ::std::eprintln!("\x1b[31m{:?}\x1b[0m", e);
                    ::std::process::exit(1);
                }

                // Vault preflight: before any RPC or play work, confirm
                // this orchestrator can load every vault key it might
                // need — the union of in-code `secret!` vaults and the
                // sealed-secret vaults present in this host's context.
                // Fail loudly and early instead of partway through.
                if let ::std::result::Result::Err(e) = cindy::secret::preflight(
                    "the orchestrator",
                    ::std::env::var("CINDY_HOST_CONTEXT").ok().as_deref(),
                ) {
                    ::std::eprintln!("\x1b[31m{:?}\x1b[0m", e);
                    ::std::process::exit(1);
                }

                let (tx, rx) = cindy::__reexports::tokio::sync::mpsc::unbounded_channel();
                cindy::orchestrator::ORCHESTRATOR_TX
                    .set(tx)
                    .expect("ORCHESTRATOR_TX already set");
                cindy::__reexports::tokio::spawn(cindy::orchestrator::rpc(
                    rx, rpc_in, rpc_out, __cindy_vault_keys,
                ));
                match #invoke_user_main.await {
                    Ok(Ok(_)) => {
                        ::std::process::exit(0);
                    }
                    Ok(Err(e)) => {
                        ::std::eprintln!("\x1b[31m{:?}\x1b[0m", e);
                        ::std::process::exit(1);
                    }
                    Err(_) => {
                        ::std::process::exit(1);
                    }
                };
            }
        }
    }
    .into()
}

/// Registers an inventory function.
///
/// Exactly one `#[cindy::inventory]` should exist per binary. The function
/// may be named in any way, can be sync or async, and may return either
/// `cindy::Inventory<V>` or `cindy::Result<cindy::Inventory<V>>`.
///
/// ```rust,ignore
/// #[cindy::wire]
/// struct MyVars {
///     var_a: u64,
///     var_b: String,
/// }
///
/// #[cindy::inventory]
/// async fn inventory() -> cindy::Result<cindy::Inventory<MyVars>> {
///     Ok(cindy::Inventory {
///         hosts: vec![
///             cindy::Host {
///                 name: "host-01".into(),
///                 tags: cindy::tags!["env:production", "continent:eu", "router"],
///                 vars: MyVars { var_a: 42, var_b: "hello".to_string() },
///             },
///         ],
///     })
/// }
/// ```
#[proc_macro_attribute]
pub fn inventory(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);

    let (attrs, vis, sig, block) = (&input.attrs, &input.vis, &input.sig, &input.block);
    let function_ident = &sig.ident;
    let asyncness = &sig.asyncness;
    let output = &sig.output;
    let inputs = &sig.inputs;

    if !inputs.is_empty() {
        panic!("`#[cindy::inventory]` functions must take no arguments");
    }

    let invocation = if asyncness.is_some() {
        quote! { #function_ident().await }
    } else {
        quote! {
            match cindy::__reexports::tokio::task::spawn_blocking(move || #function_ident())
                .await
            {
                Ok(v) => v,
                Err(je) => ::std::panic::resume_unwind(je.into_panic()),
            }
        }
    };

    quote! {
        #(#attrs)*
        #vis #asyncness fn #function_ident () #output #block

        cindy::__reexports::inventory::submit! {
            cindy::inventory::RegisteredInventory {
                dump: || ::std::boxed::Box::pin(async move {
                    cindy::inventory::IntoInventoryDump::into_inventory_dump(#invocation)
                }),
            }
        }
    }
    .into()
}

/// Marks a function as "running on the remote host".
///
/// Each `#[cindy::remote] fn foo(args) -> T` expands to a small bundle of
/// items at the same path:
///
/// 1. **`fn foo(args) -> cindy::orchestrator::Future<T>`** — the RPC shim
///    used by `#[cindy::main]` and other orchestrator-side code. Returns a
///    future that resolves once the worker has replied. (Lives in the
///    value namespace.)
/// 2. **`enum foo {} + impl foo { fn inner(args) -> T { ..user body.. } }`** -
///    the actual body, exposed at `foo::inner(args)`. Used inside other
///    `#[remote]` bodies on the worker to call siblings directly without
///    going through the RPC channel. The enum is uninhabited; it exists
///    purely so the path syntax `foo::inner` works in the type namespace,
///    side-by-side with `fn foo` in the value namespace.
#[proc_macro_attribute]
pub fn remote(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);

    let (attrs, vis, sig, block) = (&input.attrs, &input.vis, &input.sig, &input.block);

    if !sig.generics.params.is_empty() {
        panic!("Generics not allowed. Remote functions cannot be generic.");
    }

    let function_ident = &sig.ident;
    let asyncness = &sig.asyncness;
    let inputs = &sig.inputs;
    let return_type = match &sig.output {
        syn::ReturnType::Default => quote! { () },
        syn::ReturnType::Type(_, ty) => quote! { #ty },
    };

    let mut arg_idents = vec![];
    let mut arg_types = vec![];
    for input_arg in &sig.inputs {
        match input_arg {
            FnArg::Receiver(..) => panic!("Argument `self` not allowed"),
            FnArg::Typed(pat_type) => match &*pat_type.pat {
                Pat::Ident(PatIdent { ident, .. }) => {
                    arg_idents.push(ident);
                    arg_types.push(&pat_type.ty)
                }
                other => panic!("Only standard named arguments are supported. Found: {other:?}"),
            },
        }
    }

    let type_signature = arg_types
        .iter()
        .map(|ty| quote! { #ty }.to_string().replace(' ', ""))
        .collect::<Vec<String>>()
        .join(",");
    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default();
    let absolute_span_path = std::path::PathBuf::from(proc_macro::Span::call_site().file());
    let relative_path = absolute_span_path
        .strip_prefix(&manifest_dir)
        .unwrap_or(&absolute_span_path)
        .to_string_lossy()
        .to_string();
    let crate_name = std::env::var("CARGO_PKG_NAME").unwrap_or_default();

    let remote_fn_id = format!(
        "::{}::{}::{}({})",
        crate_name, relative_path, function_ident, type_signature
    );

    let invocation = if asyncness.is_some() {
        quote! { #function_ident::inner( #(#arg_idents),* ).await }
    } else {
        quote! {
            match cindy::__reexports::tokio::task::spawn_blocking(move || {
                #function_ident::inner( #(#arg_idents),* )
            })
            .await
            {
                Ok(v) => v,
                Err(je) => ::std::panic::resume_unwind(je.into_panic()),
            }
        }
    };

    let outer_docstring = format!(
        "This function can be called from the orchestrator.

To run the remote-side version of this function (e.g. to call it from another remote function),
see [`{function_ident}::inner`]."
    );
    let inner_docstring = format!(
        "This function can be called from another remote functions.

For documentation about the actual function, please refer to [`{function_ident}`]."
    );
    quote! {
        #[allow(non_camel_case_types)]
        #[doc(hidden)]
        #vis enum #function_ident {}

        #[doc(hidden)]
        impl #function_ident {
            #[doc = #inner_docstring]
            #(#attrs)*
            pub #asyncness fn inner (#inputs) -> #return_type #block
        }

        //
        //
        #[cfg(feature = "orchestrator")]
        #[doc = #outer_docstring]
        #(#attrs)*
        #vis fn #function_ident (#inputs) -> cindy::orchestrator::Future<#return_type> {
            let uuid = cindy::__reexports::uuid::Uuid::new_v4();
            let payload = cindy::common::RemoteFnPayload {
                uuid,
                fn_id: #remote_fn_id.to_string(),
                data: cindy::__reexports::postcard::to_allocvec(&( #(#arg_idents),* ))
                    .expect("Failed to serialize args"),
            };
            let (tx, rx) = cindy::__reexports::tokio::sync::oneshot::channel();
            cindy::orchestrator::ORCHESTRATOR_TX
                .get()
                .expect("ORCHESTRATOR_TX not set")
                .send(cindy::orchestrator::OutboundRegistration { payload, tx })
                .expect("Orchestrator channel closed");
            cindy::orchestrator::Future::new(rx)
        }

        #[cfg(feature = "remote")]
        cindy::__reexports::inventory::submit! {
            cindy::remote::RemoteFn {
                id: #remote_fn_id,
                function: |args_bytes| {
                    let ( #(#arg_idents),* ): ( #(#arg_types),* ) =
                        cindy::__reexports::postcard::from_bytes(&args_bytes)
                        .expect("Failed to deserialize args");

                    ::std::boxed::Box::pin(async move {
                        let result = #invocation;
                        cindy::__reexports::postcard::to_allocvec(&result)
                            .expect("Failed to serialize return value")
                    })
                },
            }
        }
    }
    .into()
}

/// Defines a builtin "action": one author-written function with ergonomic
/// `impl Into<T>` arguments and a real body, from which the macro generates
/// both the wire-level `#[remote]` entry point *and* the ergonomic layer.
///
/// `#[remote]` requires concrete, postcard-serializable argument types (no
/// generics), but call sites want `impl Into<T>` so they can pass `&str`,
/// literals, etc. Those can't be one function, so historically each action
/// was two hand-written items: a concrete `*_raw` `#[remote]` fn plus a
/// body-less convenience shim. `#[action]` collapses that to one: you write
/// the `impl Into` fn with its body once, and the macro emits:
///
/// 1. **`<name>_raw(<concrete args>)`** — a `#[remote]` fn whose body is
///    *your* body (the `impl Into<T>` params rewritten to their concrete `T`).
///    This is the real, registered remote entry point.
/// 2. **`<name>(<impl Into args>)`** (orchestrator) — converts each arg with
///    `.into()` and `.await`s the `<name>_raw` RPC shim (runs on the worker).
/// 3. **`<name>::inner(<impl Into args>)`** (both features) — converts and
///    calls `<name>_raw::inner(..)` directly, in-process (no RPC).
///
/// Asyncness is taken from the annotated fn — async fns produce an async
/// `_raw`/`::inner` (the body is `.await`-capable); no marker needed.
///
/// ```ignore
/// #[crate::action]
/// pub async fn restart(name: impl Into<String>) -> crate::Result<super::Return> {
///     apply(State { name, runtime: Some(RuntimeAction::Restarted), ..Default::default() }).await
/// }
/// ```
#[proc_macro_attribute]
pub fn action(_args: TokenStream, input: TokenStream) -> TokenStream {
    let func = parse_macro_input!(input as ItemFn);
    let (attrs, vis, sig, block) = (&func.attrs, &func.vis, &func.sig, &func.block);

    if !sig.generics.params.is_empty() {
        panic!("`#[action]` functions cannot have generic parameters (use `impl Into<T>` args)");
    }

    let ident = &sig.ident;
    let raw_ident = format_ident!("{ident}_raw");
    let is_async = sig.asyncness.is_some();
    let (maybe_async, maybe_await) = if is_async {
        (quote! { async }, quote! { .await })
    } else {
        (quote! {}, quote! {})
    };
    let return_type = match &sig.output {
        syn::ReturnType::Default => quote! { () },
        syn::ReturnType::Type(_, ty) => quote! { #ty },
    };

    // Build, in lockstep:
    //  - `ergonomic_inputs`: the original `impl Into<T>` params (the shim sig),
    //  - `raw_inputs`: the same params with each `impl Into<T>` lowered to `T`
    //    (the concrete `#[remote]` sig),
    //  - `arg_idents`: param names, forwarded as `ident.into()` from the shims.
    let mut ergonomic_inputs = vec![];
    let mut raw_inputs = vec![];
    let mut arg_idents = vec![];
    for input_arg in &sig.inputs {
        let FnArg::Typed(pat_type) = input_arg else {
            panic!("`#[action]` functions cannot take `self`");
        };
        let Pat::Ident(PatIdent { ident, .. }) = &*pat_type.pat else {
            panic!("only standard named arguments are supported in `#[action]` fns");
        };
        arg_idents.push(ident.clone());
        ergonomic_inputs.push(input_arg.clone());

        let concrete_ty = into_target(&pat_type.ty).unwrap_or_else(|| (*pat_type.ty).clone());
        raw_inputs.push(quote! { #ident: #concrete_ty });
    }

    quote! {
        // (1) The real, wire-level entry point: a `#[remote]` fn with concrete
        // args, carrying the author's body verbatim.
        #[doc(hidden)]
        #[cindy::remote]
        #vis #maybe_async fn #raw_ident (#(#raw_inputs),*) -> #return_type #block

        // (2) Orchestrator shim: convert args, drive the remote worker.
        #[cfg(feature = "orchestrator")]
        #(#attrs)*
        #vis async fn #ident (#(#ergonomic_inputs),*) -> #return_type {
            #raw_ident ( #(#arg_idents.into()),* ).await
        }

        // (3) Local/in-process variant (`#ident::inner`), under BOTH features:
        // worker builtins call siblings with it, and orchestrator code uses it
        // to act on the *local* machine without RPC. Lives in the type
        // namespace so it coexists with the `fn #ident` shim.
        #[allow(non_camel_case_types)]
        #vis enum #ident {}

        impl #ident {
            #(#attrs)*
            #vis #maybe_async fn inner (#(#ergonomic_inputs),*) -> #return_type {
                #raw_ident ::inner( #(#arg_idents.into()),* ) #maybe_await
            }
        }
    }
    .into()
}