use crate::search::{SearchEngine, SearchResult};
use include_dir::{Dir, include_dir};
use crate::metadata::{self, ModuleMetadata};
#[cfg(not(feature = "db"))]
use crate::search::BM25Index;
#[cfg(feature = "db")]
use crate::search_fts5::FTS5Index;
static STDLIB_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/stdlib");
#[derive(Debug, Clone, PartialEq)]
pub enum ModuleSource {
BuiltIn,
Project,
Global,
}
#[derive(Debug, Clone)]
pub struct DiscoveredModule {
pub module_name: String,
pub source: ModuleSource,
pub metadata: ModuleMetadata,
pub lua_source: String,
}
const BUILTINS: &[(&str, &str, &[&str])] = &[
(
"http",
"HTTP client and server: get, post, put, patch, delete, serve",
&[
"http", "client", "server", "request", "response", "headers", "endpoint", "api",
"webhook", "rest",
],
),
(
"json",
"JSON serialization: parse and encode",
&[
"json",
"serialization",
"deserialize",
"stringify",
"parse",
"encode",
"format",
],
),
(
"yaml",
"YAML serialization: parse and encode",
&[
"yaml",
"serialization",
"deserialize",
"parse",
"encode",
"format",
],
),
(
"toml",
"TOML serialization: parse and encode",
&[
"toml",
"serialization",
"deserialize",
"parse",
"encode",
"configuration",
],
),
(
"fs",
"Filesystem: read and write files",
&["fs", "filesystem", "file", "read", "write", "io", "path"],
),
(
"crypto",
"Cryptography: jwt_sign, hash, hmac, random",
&[
"crypto",
"jwt",
"signature",
"hash",
"hmac",
"encryption",
"random",
"security",
"password",
"signing",
"rsa",
"sha256",
],
),
(
"base64",
"Base64 encoding and decoding",
&["base64", "encoding", "decode", "encode", "binary"],
),
(
"regex",
"Regular expressions: match, find, find_all, replace",
&[
"regex",
"pattern",
"match",
"find",
"replace",
"regular-expression",
"regexp",
],
),
(
"db",
"Database: connect, query, execute, close (Postgres, MySQL, SQLite)",
&[
"db",
"database",
"sql",
"postgres",
"mysql",
"sqlite",
"connection",
"query",
"execute",
],
),
(
"ws",
"WebSocket: connect, send, recv, close",
&[
"ws",
"websocket",
"connection",
"message",
"streaming",
"realtime",
"socket",
],
),
(
"template",
"Jinja2-compatible templates: render file or string",
&[
"template",
"jinja2",
"rendering",
"string-template",
"mustache",
"render",
],
),
(
"async",
"Async tasks: spawn, spawn_interval, await, cancel",
&[
"async",
"asynchronous",
"task",
"coroutine",
"concurrent",
"spawn",
"interval",
],
),
(
"assert",
"Assertions: eq, gt, lt, contains, not_nil, matches",
&[
"assert",
"assertion",
"test",
"validation",
"comparison",
"check",
"verify",
],
),
(
"log",
"Logging: info, warn, error",
&[
"log", "logging", "output", "debug", "error", "warning", "info", "trace",
],
),
(
"env",
"Environment variables: get",
&["env", "environment", "variable", "configuration", "config"],
),
(
"sleep",
"Sleep for N seconds",
&["sleep", "delay", "pause", "wait", "time"],
),
(
"time",
"Unix timestamp in seconds",
&["time", "timestamp", "unix", "epoch", "clock", "datetime"],
),
];
pub fn discover_modules() -> Vec<DiscoveredModule> {
let mut modules = Vec::new();
discover_filesystem_modules(
std::path::Path::new("./modules"),
ModuleSource::Project,
&mut modules,
);
let global_path = resolve_global_modules_path();
if let Some(path) = global_path {
discover_filesystem_modules(&path, ModuleSource::Global, &mut modules);
}
discover_embedded_stdlib(&mut modules);
discover_rust_builtins(&mut modules);
modules
}
pub fn build_index(modules: &[DiscoveredModule]) -> Box<dyn SearchEngine> {
#[cfg(feature = "db")]
{
let mut idx = FTS5Index::new();
for m in modules {
idx.add_document(
&m.module_name,
&[
("keywords", &m.metadata.keywords.join(" "), 3.0),
("module_name", &m.module_name, 2.0),
("description", &m.metadata.description, 1.0),
("functions", &m.metadata.auto_functions.join(" "), 1.0),
],
);
}
Box::new(idx)
}
#[cfg(not(feature = "db"))]
{
let mut idx = BM25Index::new();
for m in modules {
idx.add_document(
&m.module_name,
&[
("keywords", &m.metadata.keywords.join(" "), 3.0),
("module_name", &m.module_name, 2.0),
("description", &m.metadata.description, 1.0),
("functions", &m.metadata.auto_functions.join(" "), 1.0),
],
);
}
Box::new(idx)
}
}
pub fn search_modules(query: &str, limit: usize) -> Vec<SearchResult> {
let modules = discover_modules();
let index = build_index(&modules);
index.search(query, limit)
}
fn resolve_global_modules_path() -> Option<std::path::PathBuf> {
if let Ok(custom) = std::env::var(crate::lua::MODULES_PATH_ENV) {
return Some(std::path::PathBuf::from(custom));
}
if let Ok(home) = std::env::var("HOME") {
return Some(std::path::Path::new(&home).join(".assay/modules"));
}
None
}
fn discover_filesystem_modules(
dir: &std::path::Path,
source: ModuleSource,
modules: &mut Vec<DiscoveredModule>,
) {
let entries = match std::fs::read_dir(dir) {
Ok(entries) => entries,
Err(_) => return, };
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("lua") {
continue;
}
let lua_source = match std::fs::read_to_string(&path) {
Ok(s) => s,
Err(_) => continue,
};
let stem = path
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or_default();
let module_name = format!("assay.{stem}");
let meta = metadata::parse_metadata(&lua_source);
modules.push(DiscoveredModule {
module_name,
source: source.clone(),
metadata: meta,
lua_source,
});
}
}
fn discover_embedded_stdlib(modules: &mut Vec<DiscoveredModule>) {
for file in STDLIB_DIR.files() {
let path = file.path();
if path.extension().and_then(|e| e.to_str()) != Some("lua") {
continue;
}
let lua_source = match file.contents_utf8() {
Some(s) => s,
None => continue,
};
let stem = path
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or_default();
let module_name = format!("assay.{stem}");
let meta = metadata::parse_metadata(lua_source);
modules.push(DiscoveredModule {
module_name,
source: ModuleSource::BuiltIn,
metadata: meta,
lua_source: lua_source.to_string(),
});
}
}
fn discover_rust_builtins(modules: &mut Vec<DiscoveredModule>) {
for &(name, description, kw) in BUILTINS {
modules.push(DiscoveredModule {
module_name: name.to_string(),
source: ModuleSource::BuiltIn,
lua_source: String::new(),
metadata: ModuleMetadata {
module_name: name.to_string(),
description: description.to_string(),
keywords: kw.iter().map(|k| k.to_string()).collect(),
..Default::default()
},
});
}
}