use mlua::{Function, Lua, Result as LuaResult, Table, Value};
use pasta_core::registry::{SceneRegistry, WordDefRegistry};
use std::collections::HashMap;
#[derive(Debug)]
pub struct WordCollectionEntry {
pub key: String,
pub values: Vec<String>,
pub is_local: bool,
pub scene_name: Option<String>,
pub actor_name: Option<String>,
}
pub fn collect_scenes(lua: &Lua) -> LuaResult<Vec<(String, String)>> {
let scene_module: Table = lua.load("return require('pasta.scene')").eval()?;
let get_all_scenes: Function = scene_module.get("get_all_scenes")?;
let registry: Table = get_all_scenes.call(())?;
let mut scenes = Vec::new();
for pair in registry.pairs::<String, Table>() {
let (global_name, scene_table) = pair?;
for entry in scene_table.pairs::<String, Value>() {
let (local_name, _value) = entry?;
if local_name == "__global_name__" {
continue;
}
scenes.push((global_name.clone(), local_name));
}
}
if scenes.is_empty() {
tracing::warn!("Scene registry is empty");
}
Ok(scenes)
}
fn collect_word_entries(
word_map: &Table,
is_local: bool,
scene_name: Option<String>,
actor_name: Option<String>,
out: &mut Vec<WordCollectionEntry>,
) -> LuaResult<()> {
for key_pair in word_map.pairs::<String, Table>() {
let (key, values_list) = key_pair?;
for values_pair in values_list.pairs::<i64, Table>() {
let (_idx, values_table) = values_pair?;
let values: Vec<String> = values_table
.pairs::<i64, String>()
.map(|r| r.map(|(_, v)| v))
.collect::<LuaResult<_>>()?;
out.push(WordCollectionEntry {
key: key.clone(),
values,
is_local,
scene_name: scene_name.clone(),
actor_name: actor_name.clone(),
});
}
}
Ok(())
}
pub fn collect_words(lua: &Lua) -> LuaResult<Vec<WordCollectionEntry>> {
let word_module: Table = lua.load("return require('pasta.word')").eval()?;
let get_all_words: Function = word_module.get("get_all_words")?;
let all_words: Table = get_all_words.call(())?;
let mut entries = Vec::new();
if let Ok(global_words) = all_words.get::<Table>("global") {
collect_word_entries(&global_words, false, None, None, &mut entries)?;
}
if let Ok(local_words) = all_words.get::<Table>("local") {
for scene_pair in local_words.pairs::<String, Table>() {
let (scene_name, scene_words) = scene_pair?;
collect_word_entries(&scene_words, true, Some(scene_name), None, &mut entries)?;
}
}
if let Ok(actor_words) = all_words.get::<Table>("actor") {
for actor_pair in actor_words.pairs::<String, Table>() {
let (actor_name, actor_word_map) = actor_pair?;
collect_word_entries(&actor_word_map, false, None, Some(actor_name), &mut entries)?;
}
}
Ok(entries)
}
fn build_scene_registry(scenes: &[(String, String)]) -> SceneRegistry {
let mut registry = SceneRegistry::new();
let mut grouped: HashMap<String, Vec<String>> = HashMap::new();
for (global_name, local_name) in scenes {
grouped
.entry(global_name.clone())
.or_default()
.push(local_name.clone());
}
for (global_name, local_names) in grouped {
registry.register_global_raw(&global_name, &local_names, HashMap::new());
}
registry
}
fn build_word_registry(entries: &[WordCollectionEntry]) -> WordDefRegistry {
let mut registry = WordDefRegistry::new();
for entry in entries {
if let Some(ref actor_name) = entry.actor_name {
registry.register_actor(actor_name, &entry.key, entry.values.clone());
} else if entry.is_local {
if let Some(ref scene_name) = entry.scene_name {
registry.register_local(scene_name, &entry.key, entry.values.clone());
}
} else {
registry.register_global(&entry.key, entry.values.clone());
}
}
registry
}
pub fn finalize_scene_impl(lua: &Lua) -> LuaResult<bool> {
let scenes = collect_scenes(lua)?;
tracing::debug!(
scene_count = scenes.len(),
"Collected scenes from Lua registry"
);
let word_entries = collect_words(lua)?;
tracing::debug!(
word_count = word_entries.len(),
"Collected words from Lua registry"
);
let scene_registry = build_scene_registry(&scenes);
let word_registry = build_word_registry(&word_entries);
crate::search::register(lua, scene_registry, word_registry)?;
tracing::info!(
scenes = scenes.len(),
words = word_entries.len(),
"SearchContext constructed and registered as @pasta_search"
);
Ok(true)
}
pub fn register_finalize_scene(lua: &Lua) -> LuaResult<()> {
let finalize_fn = lua.create_function(|lua, ()| finalize_scene_impl(lua))?;
let package: Table = lua.globals().get("package")?;
let loaded: Table = package.get("loaded")?;
let pasta_module: Table = if let Ok(module) = loaded.get::<Table>("pasta") {
module
} else {
lua.load("return require('pasta')").eval()?
};
pasta_module.set("finalize_scene", finalize_fn)?;
tracing::debug!("Registered finalize_scene Rust binding");
Ok(())
}