graphix-shell 0.8.0

use anyhow::{bail, Context, Result};
use arcstr::ArcStr;
use clap::{Parser, Subcommand};
use enumflags2::BitFlags;
use flexi_logger::{FileSpec, Logger};
use graphix_compiler::{
    expr::{ModuleResolver, Source},
    CFlag,
};
use graphix_package::{GraphixPM, MainThreadHandle, PackageId};
use graphix_rt::NoExt;
use graphix_shell::{Mode, ShellBuilder};
use log::info;
use netidx::{
    config::Config,
    path::Path,
    publisher::{BindCfg, DesiredAuth, Publisher, PublisherBuilder},
    subscriber::{Subscriber, SubscriberBuilder},
    InternalOnly,
};
use std::{path::PathBuf, str::FromStr, sync::OnceLock, time::Duration};

#[derive(Debug, Clone, Copy)]
enum RawFlag {
    Unhandled,
    NoUnhandled,
    Unused,
    NoUnused,
    Error,
    NoError,
}

impl FromStr for RawFlag {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s {
            "unhandled" => Ok(Self::Unhandled),
            "no-unhandled" => Ok(Self::NoUnhandled),
            "unused" => Ok(Self::Unused),
            "no-unused" => Ok(Self::NoUnused),
            "error" => Ok(Self::Error),
            "no-error" => Ok(Self::NoError),
            s => bail!("invalid flag {s}"),
        }
    }
}

impl RawFlag {
    fn as_flags(flags: &[RawFlag]) -> (BitFlags<CFlag>, BitFlags<CFlag>) {
        let mut enable = BitFlags::empty();
        let mut disable = BitFlags::empty();
        for fl in flags {
            match fl {
                Self::Unhandled => enable.insert(CFlag::WarnUnhandled),
                Self::NoUnhandled => disable.insert(CFlag::WarnUnhandled),
                Self::Unused => enable.insert(CFlag::WarnUnused),
                Self::NoUnused => disable.insert(CFlag::WarnUnused),
                Self::Error => enable.insert(CFlag::WarningsAreErrors),
                Self::NoError => disable.insert(CFlag::WarningsAreErrors),
            }
        }
        (enable, disable)
    }
}

#[derive(Subcommand)]
enum PackageAction {
    /// Add packages to the graphix runtime
    Add {
        /// Package names to add (with optional @version suffix)
        packages: Vec<String>,
        /// Skip crates.io validation (for packages from alternative registries or local sources)
        #[arg(long)]
        skip_crates_io_check: bool,
        /// Use a local path dependency instead of a crates.io version
        #[arg(long)]
        path: Option<PathBuf>,
    },
    /// Remove packages from the graphix runtime
    Remove {
        /// Package names to remove
        packages: Vec<String>,
    },
    /// Search crates.io for graphix packages
    Search {
        /// Search query
        query: String,
    },
    /// List installed packages
    List,
    /// Rebuild graphix from the current packages.toml
    Rebuild,
    /// Create a new graphix package
    Create {
        /// Package name (will be prefixed with graphix-package- if needed)
        name: String,
        /// Directory to create the package in
        #[arg(long, default_value = ".")]
        dir: PathBuf,
    },
    /// Update graphix to the latest version
    Update,
    /// Build a standalone graphix binary from the package in the current directory
    BuildStandalone,
}

#[derive(Subcommand)]
enum Command {
    /// Manage graphix packages
    Package {
        #[command(subcommand)]
        action: PackageAction,
    },
    /// Run the Language Server Protocol server (communicates over stdio).
    /// Editors typically launch this automatically; you don't usually run it directly.
    Lsp,
}

#[derive(Parser)]
#[command(version, about, trailing_var_arg = true)]
struct Params {
    #[command(subcommand)]
    command: Option<Command>,
    /// enable logging and put the log in the specified directory. You
    /// should also set the RUST_LOG enviornment variable. e.g. RUST_LOG=debug
    #[arg(long)]
    log_dir: Option<PathBuf>,
    /// path to the netidx config to load, otherwise the default will
    /// be loaded (unless --no-netidx is specified)
    #[arg(long)]
    config: Option<PathBuf>,
    /// the desired netidx auth mechanism to use, otherwise use the config default
    #[arg(long)]
    auth: Option<DesiredAuth>,
    /// the kerberos user principal name to use for netidx, otherwise
    /// the default from the current user's cached tickets, only valid
    /// if using kerberos auth
    #[arg(long)]
    upn: Option<String>,
    /// the netidx nerberos service princial name, otherwise the
    /// default from the current user's cached ticket, only valid if
    /// using kerberos auth
    #[arg(long)]
    spn: Option<String>,
    /// the netidx tls identity to use, otherwise use the configured
    /// default, only valid if using tls auth.
    #[arg(long)]
    identity: Option<String>,
    /// specify the netidx publisher bind address.
    #[arg(long)]
    bind: Option<BindCfg>,
    /// drop subscribers if they don't consume published values with
    /// the specifed timeout (in seconds).
    #[arg(long)]
    publish_timeout: Option<u64>,
    /// module resolution from netidx should fail if we can't
    /// subscribe to the module before the timeout expires. Default,
    /// wait forever.
    #[arg(long)]
    resolve_timeout: Option<u64>,
    /// disable netidx, net functions will only work internally
    #[arg(short, long)]
    no_netidx: bool,
    /// do not attempt to run the init module
    #[arg(short = 'i', long)]
    no_init: bool,
    /// do not execute the program, just veryify that it compiles and
    /// type checks.
    #[arg(long = "check")]
    check: bool,
    /// run the program in the specified file instead of starting the REPL
    file: Option<ArcStr>,
    /// enable or disable compiler flags. Currently supported flags are,
    /// - unhandled, no-unhandled: warn about unhandled ? operators (default)
    /// - unused, no-unused: warn about unused variables (default)
    /// - error, no-error makes warnings errors
    ///
    /// the no- variant turns the flag off. If both are specifed the no- variant
    /// always wins
    #[arg(short = 'W')]
    warn: Vec<RawFlag>,
    /// arguments passed to the graphix program (everything after the filename)
    #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
    program_args: Vec<String>,
}

impl Params {
    async fn get_pub_sub(&self, cfg: Result<Config>) -> Result<(Publisher, Subscriber)> {
        let res = async {
            let cfg = cfg?;
            let auth = match &self.auth {
                None => cfg.default_auth(),
                Some(a) => a.clone(),
            };
            let publisher = PublisherBuilder::new(cfg.clone())
                .bind_cfg(self.bind)
                .build()
                .await
                .context("creating publisher")?;
            let subscriber = SubscriberBuilder::new(cfg)
                .desired_auth(auth)
                .build()
                .context("creating subscriber")?;
            Ok::<_, anyhow::Error>((publisher, subscriber))
        };
        match res.await {
            Ok(ps) => Ok(ps),
            Err(e) => {
                eprintln!("netidx initialization failed {e:?}");
                eprintln!("netidx will be process internal only");
                eprintln!("to fix this see https://netidx.github.io/netidx-book");
                static NETIDX: OnceLock<InternalOnly> = OnceLock::new();
                if let Err(_) = NETIDX.set(InternalOnly::new().await?) {
                    panic!("BUG: NETIDX static set multiple times")
                }
                let env = NETIDX.get().unwrap();
                Ok((env.publisher().clone(), env.subscriber().clone()))
            }
        }
    }
}

fn parse_package_arg(s: &str) -> PackageId {
    match s.split_once('@') {
        Some((name, version)) => PackageId::new(name, Some(version)),
        None => PackageId::new(s, None),
    }
}

#[tokio::main]
async fn handle_package(action: PackageAction) -> Result<()> {
    match action {
        PackageAction::Add { packages, skip_crates_io_check, path } => {
            let pm = GraphixPM::new().await?;
            let ids: Vec<_> = if let Some(ref path) = path {
                packages
                    .iter()
                    .map(|s| {
                        let name = s.split('@').next().unwrap();
                        PackageId::with_path(name, path.clone())
                    })
                    .collect()
            } else {
                packages.iter().map(|s| parse_package_arg(s)).collect()
            };
            pm.add_packages(&ids, skip_crates_io_check).await
        }
        PackageAction::Remove { packages } => {
            let pm = GraphixPM::new().await?;
            let ids: Vec<_> = packages.iter().map(|s| PackageId::new(s, None)).collect();
            pm.remove_packages(&ids).await
        }
        PackageAction::Search { query } => {
            let pm = GraphixPM::new().await?;
            pm.search(&query).await
        }
        PackageAction::List => {
            let pm = GraphixPM::new().await?;
            pm.list().await
        }
        PackageAction::Rebuild => {
            let pm = GraphixPM::new().await?;
            pm.do_rebuild().await
        }
        PackageAction::Create { name, dir } => {
            let full_name = if name.starts_with("graphix-package-") {
                name
            } else {
                format!("graphix-package-{name}")
            };
            graphix_package::create_package(&dir, &full_name).await
        }
        PackageAction::Update => {
            let pm = GraphixPM::new().await?;
            pm.update().await
        }
        PackageAction::BuildStandalone => {
            let pm = GraphixPM::new().await?;
            let cwd = std::env::current_dir().context("getting current directory")?;
            pm.build_standalone(&cwd, None).await
        }
    }
}

fn tokio_main(
    p: Params,
    cfg: Result<Config>,
    run_on_main: MainThreadHandle,
) -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .context("building tokio runtime")?;
    rt.block_on(async move {
        if let Some(dir) = &p.log_dir {
            let _ = Logger::try_with_env()
                .context("initializing log")?
                .log_to_file(
                    FileSpec::default()
                        .directory(dir)
                        .basename("graphix")
                        .use_timestamp(false),
                )
                .start()
                .context("starting log")?;
        }
        info!("graphix shell starting");
        let mut _internal = None;
        let (publisher, subscriber) = if p.no_netidx {
            let i = InternalOnly::new().await?;
            let (p, s) = (i.publisher().clone(), i.subscriber().clone());
            _internal = Some(i);
            (p, s)
        } else {
            p.get_pub_sub(cfg).await?
        };
        let mut shell = ShellBuilder::<NoExt>::default();
        let program_args: Vec<ArcStr> =
            p.program_args.iter().map(|s| ArcStr::from(s.as_str())).collect();
        shell = shell.program_args(program_args);
        shell = shell.no_init(p.no_init);
        if let Some(t) = p.publish_timeout {
            shell = shell.publish_timeout(Duration::from_secs(t));
        }
        if let Some(t) = p.resolve_timeout {
            shell = shell.resolve_timeout(Duration::from_secs(t));
        }
        if p.file.is_none() && p.check {
            bail!("check mode requires a file to check")
        }
        if let Some(f) = &p.file {
            let source = match f.strip_prefix("netidx:") {
                Some(path) => {
                    shell = shell.module_resolvers(vec![ModuleResolver::Netidx {
                        subscriber: subscriber.clone(),
                        base: netidx::path::Path::from(ArcStr::from(path)),
                        timeout: None,
                    }]);
                    Source::Netidx(Path::from(ArcStr::from(path)))
                }
                None => {
                    let path = PathBuf::from(&**f).canonicalize()?;
                    let path = if path.is_dir() { path.join("main.gx") } else { path };
                    match path.parent() {
                        Some(p) if p.as_os_str().is_empty() => (),
                        None => (),
                        Some(p) => {
                            let p = PathBuf::from(p);
                            shell = shell.module_resolvers(vec![ModuleResolver::Files {
                                base: p,
                                overrides: None,
                            }]);
                        }
                    };
                    Source::File(path)
                }
            };
            let mode = if p.check { Mode::Check(source) } else { Mode::Script(source) };
            shell = shell.mode(mode);
        }
        let (enable, disable) = RawFlag::as_flags(&p.warn);
        shell
            .publisher(publisher)
            .subscriber(subscriber)
            .enable_flags(enable)
            .disable_flags(disable)
            .build()?
            .run(run_on_main)
            .await
    })
}

fn main() -> Result<()> {
    Config::maybe_run_machine_local_resolver()?;
    let p = Params::parse();
    match p.command {
        Some(Command::Package { action }) => return handle_package(action),
        Some(Command::Lsp) => return graphix_shell::lsp_backend::run(),
        None => (),
    }
    let cfg = match &p.config {
        None => Config::load_default_or_local_only(),
        Some(p) => Config::load(p),
    };
    let (handle, main_rx) = MainThreadHandle::new();
    let tokio_handle = std::thread::Builder::new()
        .name("graphix-tokio".into())
        .spawn(move || tokio_main(p, cfg, handle))
        .expect("failed to spawn tokio thread");
    while let Ok(f) = main_rx.recv() {
        f();
    }
    tokio_handle.join().map_err(|_| anyhow::anyhow!("tokio thread panicked"))?
}