aipack 0.8.23

//! Defines the `aip.md` module, used in the Lua engine.
//!
//! This module provides functions for processing Markdown content, particularly useful for
//! extracting structured data like code blocks and metadata from text generated by LLMs.
//!
//! ---
//!
//! ## Lua documentation
//!
//! The `aip.md` module exposes functions to interact with and parse markdown content.
//!
//! ### Functions
//!
//! - `aip.md.extract_blocks(md_content: string, options?: string | {lang?: string, extrude?: "content"}): list<MdBlock> | (list<MdBlock>, string)`
//! - `aip.md.extract_meta(md_content: string): table, string`
//! - `aip.md.outer_block_content_or_raw(md_content: string): string`

use crate::Result;
use crate::runtime::Runtime;
use crate::script::support::into_option_string;
use crate::support::W;
use crate::support::md::{self, MdRefIter};
use crate::types::{Extrude, MdBlock, MdRef};
use mlua::{IntoLua, Lua, LuaSerdeExt, MultiValue, Table, Value};

// region:    --- Module Init

pub fn init_module(lua: &Lua, _runtime: &Runtime) -> Result<Table> {
	let table = lua.create_table()?;

	let extract_blocks_fn = lua.create_function(extract_blocks)?;
	let outer_block_content_or_raw_fn = lua.create_function(outer_block_content_or_raw)?;
	let extract_meta_fn = lua.create_function(extract_meta)?;
	let extract_refs_fn = lua.create_function(extract_refs)?;

	table.set("extract_blocks", extract_blocks_fn)?;
	table.set("extract_meta", extract_meta_fn)?;
	table.set("extract_refs", extract_refs_fn)?;
	table.set("outer_block_content_or_raw", outer_block_content_or_raw_fn)?;

	Ok(table)
}

// endregion: --- Module Init

/// ## Lua Documentation
///
/// Extracts markdown blocks from a string, optionally filtering by language or extracting remaining content.
///
/// ```lua
/// -- API Signatures
///
/// -- Extract all blocks:
/// aip.md.extract_blocks(md_content: string): list<MdBlock>
///
/// -- Extract blocks of a specific language:
/// aip.md.extract_blocks(md_content: string, lang: string): list<MdBlock>
///
/// -- Extract blocks of a specific language and the content outside these blocks:
/// aip.md.extract_blocks(md_content: string, {lang?: string, extrude: "content"}): list<MdBlock>, string
///
/// -- Extract all blocks and the content outside these blocks:
/// aip.md.extract_blocks(md_content: string, {extrude: "content"}): list<MdBlock>, string
/// ```
///
/// Extracts code blocks and other fenced blocks from `md_content`.
///
/// ### Arguments
///
/// - `md_content: string`: The markdown content string to process.
/// - `options?: string | table` (optional):
///   - If a string, it is treated as the `lang` filter. Only blocks with this language identifier will be returned.
///   - If a table, it can contain the following optional fields:
///     - `lang?: string` (optional): Filters blocks by this language identifier. If nil or not present, all blocks are included.
///     - `extrude?: "content"` (optional): If set to the string `"content"`, the function will also return the content of the markdown string that is *outside* of the extracted blocks.
///
/// ### Returns
///
/// The return type depends on the `extrude` option:
///
/// - If `extrude = "content"` is *not* specified:
///   - `list<MdBlock>`: A Lua list (table) of `MdBlock` objects. Each object represents a parsed block:
///     ```ts
///     {
///       block_type: "code" | "block", // Type of block (currently only "code" supported)
///       lang: string | nil,          // Language identifier for code blocks (e.g., "rust", "js"), nil otherwise
///       content: string              // The content inside the block (excluding fence lines)
///     }
///     ```
/// - If `extrude = "content"` is specified:
///   - `list<MdBlock>`: A Lua list (table) of `MdBlock` objects (same structure as above).
///   - `string`: The content of the markdown string that was outside of the extracted blocks.
///
/// The list of blocks is empty if no matching blocks are found.
///
/// ### Example
///
/// ```lua
/// -- Extract all code blocks
/// local all_blocks = aip.md.extract_blocks([[
/// # Example
///
/// ```rust
/// fn main() {}
/// ```
///
/// Some text.
///
/// ```lua
/// print("hello")
/// ```
/// ]])
/// print(#all_blocks) -- Output: 2
///
/// -- Extract only rust blocks
/// local rust_blocks = aip.md.extract_blocks(all_blocks_content, "rust")
/// print(#rust_blocks) -- Output: 1
///
/// -- Extract lua blocks and the remaining content
/// local lua_blocks, remaining = aip.md.extract_blocks(all_blocks_content, {lang = "lua", extrude = "content"})
/// print(#lua_blocks) -- Output: 1
/// print(remaining) -- Output: "# Example\n\n\n\nSome text.\n\n\n\n"
/// ```
///
/// ### Error
///
/// Returns an error if:
/// - The `options` argument is provided but is not a string or a table.
/// - The `options` table contains an `extrude` field with a value other than `"content"`.
/// - An error occurs during parsing.
fn extract_blocks(lua: &Lua, (md_content, options): (String, Option<Value>)) -> mlua::Result<MultiValue> {
	let (lang, extrude): (Option<String>, Option<Extrude>) = match options {
		// if options is of type string, then, just lang name
		Some(Value::String(string)) => (Some(string.to_string_lossy()), None),
		// if it is a table
		Some(Value::Table(table)) => {
			let lang = table.get::<Option<Value>>("lang")?;
			let lang = lang
				.map(|v| {
					v.to_string()
						.map_err(|_err| crate::Error::custom("md_extract_blocks lang options must be of type string"))
				})
				.transpose()?;

			let extrude = Extrude::extract_from_table_value(&table)?;

			(lang, extrude)
		}
		// TODO: Probably need to send error
		_ => (None, None),
	};

	let blocks_it = md::MdBlockIter::new(&md_content, lang.as_deref(), extrude);
	let mut values = MultiValue::new();

	match extrude {
		Some(Extrude::Content) => {
			let (blocks, content) = blocks_it.collect_blocks_and_extruded_content();
			values.push_back(lua.to_value(&blocks)?); // Use lua.to_value for Vec<MdBlock>
			let content = lua.create_string(&content)?;
			values.push_back(Value::String(content));
		}
		_ => {
			let blocks: Vec<MdBlock> = blocks_it.collect();
			values.push_back(lua.to_value(&blocks)?) // Use lua.to_value for Vec<MdBlock>
		}
	}

	Ok(values)
}

/// ## Lua Documentation
///
/// Extracts meta blocks from markdown content and returns a table of the merged meta values and the remaining content.
///
/// ```lua
/// -- API Signature
/// aip.md.extract_meta(md_content: string): table, string
/// ```
///
/// Scans the provided `md_content` for meta blocks (code blocks starting with `#!meta`).
/// Extracts the content of all such meta blocks, attempts to parse them as TOML, and merges
/// the resulting key-value pairs into a single Lua table. The function also returns the
/// original markdown content with the meta blocks removed.
///
/// ### Arguments
///
/// - `md_content: string | nil`: The markdown content string to process.
///
/// ### Returns
///
/// Returns two values:
///
/// - `table`: A Lua table containing the merged data from all meta blocks. If no meta blocks are found, an empty table is returned.
/// - `string`: The remaining content of the markdown string after all meta blocks have been removed.
///
/// If `md_content` is nil, then, return `{nil, nil}`
///
/// ### Example
///
/// ```lua
/// local content = [[
/// Some introductory text.
///
/// ```toml
/// #!meta
/// title = "My Document"
/// version = 1.0
/// ```
///
/// Main content.
///
/// ```yaml
/// #!meta
/// author: John Doe
/// tags: [guide, markdown]
/// ```
/// ]]
///
/// local meta_data, remaining_content = aip.md.extract_meta(content)
///
/// print(meta_data.title) -- Output: "My Document"
/// print(meta_data.version) -- Output: 1.0
/// print(meta_data.author) -- Output: "John Doe"
/// -- Note: Tags will likely be a Lua table/list
///
/// print(remaining_content)
/// ```
///
/// ### Error
///
/// Returns an error if:
/// - A meta block is found but its content cannot be parsed as TOML.
/// - An internal error occurs during processing.
fn extract_meta(lua: &Lua, md_content: Value) -> mlua::Result<MultiValue> {
	// allow md_content to be nil
	let Some(md_content) = into_option_string(md_content, "aip.md.extract_meta")? else {
		return Ok(MultiValue::from_vec(vec![Value::Nil, Value::Nil]));
	};
	// extract value
	let (value, remain) = md::extract_meta(&md_content)?;
	let lua_value = lua.to_value(&value)?;
	// return the tuple of two
	let values = MultiValue::from_vec(vec![lua_value, W(remain).into_lua(lua)?]);
	Ok(values)
}

/// ## Lua Documentation
///
/// Extracts the content within the outermost code block, or returns the raw content if no such block exists.
///
/// ```lua
/// -- API Signature
/// aip.md.outer_block_content_or_raw(md_content: string): string
/// ```
///
/// Checks if the `md_content` starts and ends with a markdown fenced code block (e.g., lines starting with ```` ``` ```` `).
///
/// - If it does, it extracts and returns the content that is *inside* the outermost pair of code block fences.
/// - If it does not start and end with a code block, or if the structure is malformed (e.g., only starts with a block but doesn't end with one), the original `md_content` is returned as is.
///
/// This function is particularly useful for processing responses from LLMs that often wrap their primary output (like code or structured text) within a markdown code block. Using this function allows you to easily get the "raw" content intended by the LLM, stripping away the markdown fences.
///
/// ### Arguments
///
/// - `md_content: string`: The markdown content string to process.
///
/// ### Returns
///
/// - `string`: The content inside the outermost code block fences if present, otherwise the original `md_content`.
///
/// ### Example
///
/// ```lua
/// -- Example with a code block
/// local content_with_block = [[
/// ```rust
/// fn main() {
///   println!("Hello!");
/// }
/// ```
/// ]]
/// local extracted = aip.md.outer_block_content_or_raw(content_with_block)
/// print(extracted)
/// -- Output:
/// -- [[
/// -- fn main() {
/// --   println!("Hello!");
/// -- }
/// -- ]]
///
/// -- Example without a code block
/// local content_raw = [[Just some plain text
/// without any blocks.]]
/// local extracted_raw = aip.md.outer_block_content_or_raw(content_raw)
/// print(extracted_raw)
/// -- Output:
/// -- [[Just some plain text
/// -- without any blocks.]]
/// ```
fn outer_block_content_or_raw(_lua: &Lua, md_content: String) -> mlua::Result<String> {
	let res = md::outer_block_content_or_raw(&md_content);
	Ok(res.into_owned())
}

/// ## Lua Documentation
///
/// Extracts all markdown references (links and images) from markdown content.
///
/// ```lua
/// -- API Signature
/// aip.md.extract_refs(md_content: string): list<MdRef>
/// ```
///
/// Scans the provided `md_content` for markdown references in the forms:
/// - Links: `[text](target)`
/// - Images: `![alt text](target)`
///
/// References inside code blocks (fenced with ``` or ````) and inline code (backticks) are skipped.
///
/// ### Arguments
///
/// - `md_content: string | nil`: The markdown content string to process.
///
/// ### Returns
///
/// - `list<MdRef>`: A Lua list (table) of `MdRef` objects. Each object represents a parsed reference:
///   ```ts
///   {
///     _type: "MdRef",       // Type identifier
///     target: string,       // URL, file path, or in-document anchor
///     text: string | nil,   // Content inside the brackets (nil if empty)
///     inline: boolean,      // True if prefixed with '![' (image)
///     kind: string          // "Anchor" | "File" | "Url"
///   }
///   ```
///
/// If `md_content` is nil, returns an empty list.
///
/// ### Example
///
/// ```lua
/// local content = [[
/// Check out [this link](https://example.com) and [docs](docs/page.md).
///
/// Also see ![image](assets/photo.jpg) for reference.
///
/// ```
/// [not a link](https://fake.com)
/// ```
/// ]]
///
/// local refs = aip.md.extract_refs(content)
/// print(#refs) -- Output: 3
///
/// for _, ref in ipairs(refs) do
///   print(ref.target, ref.kind, ref.inline)
/// end
/// -- Output:
/// -- https://example.com    Url    false
/// -- docs/page.md           File   false
/// -- assets/photo.jpg       File   true
/// ```
///
/// ### Error
///
/// Returns an error if an internal error occurs during processing.
fn extract_refs(lua: &Lua, md_content: Value) -> mlua::Result<Value> {
	// allow md_content to be nil
	let Some(md_content) = into_option_string(md_content, "aip.md.extract_refs")? else {
		// Return empty table for nil input
		return Ok(Value::Table(lua.create_table()?));
	};

	let refs: Vec<MdRef> = MdRefIter::new(&md_content).collect();
	let lua_value = lua.to_value(&refs)?;

	Ok(lua_value)
}

// region:    --- Tests

#[cfg(test)]
mod tests {
	type Result<T> = core::result::Result<T, Box<dyn std::error::Error>>; // For tests.

	use crate::_test_support::{assert_contains, assert_not_contains, eval_lua, run_reflective_agent, setup_lua};
	use serde_json::Value;
	use value_ext::JsonValueExt;

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_blocks_simple() -> Result<()> {
		// -- Setup & Fixtures
		// NOTE: the [[ ]] for multi line in lua breaks with the ``` for code block, so, reading files.
		let fx_script = r#"
local file = aip.file.load("agent-script/agent-before-all-inputs-gen.aip")
return aip.md.extract_blocks(file.content, {lang = "lua"})
		"#;

		// -- Exec
		let res = run_reflective_agent(fx_script, None).await?;

		// -- Check
		assert!(res.is_array());
		let blocks = res.as_array().ok_or("Res should be array")?;
		assert_eq!(blocks.len(), 4, "Should have found 4 lua blocks");

		// Check first block
		let first_block = &blocks[0];
		assert_eq!(first_block.x_get_str("lang")?, "lua");
		assert!(first_block.x_get_str("content")?.contains("before_all_response"));

		// Check second block
		let second_block = &blocks[1];
		assert_eq!(second_block.x_get_str("lang")?, "lua");
		assert!(second_block.x_get_str("content")?.contains("Data with input"));

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_blocks_with_lang_and_extruded_content() -> Result<()> {
		// -- Setup & Fixtures
		// NOTE: the [[ ]] for multi line in lua breaks when line starts with ```, so work around
		let fx_script = r#"
local content = "This is some content\n"
content = content .. "\n```lua\n--some lua \n```\n"
content = content .. "and other block\n\n```rust\n//! some rust block \n```\n"
content = content .. "The end"

local blocks, extruded_content = aip.md.extract_blocks(content, {lang = "lua", extrude = "content"})
return {
		blocks = blocks,
		extruded_content = extruded_content
}
		"#;

		// -- Exec
		let res = run_reflective_agent(fx_script, None).await?;

		// -- Check Blocks
		let blocks = res.pointer("/blocks").ok_or("Should have blocks")?;
		assert!(blocks.is_array());
		let blocks = blocks.as_array().ok_or("Should be arr")?;
		assert_eq!(blocks.len(), 1, "Should have found 1 lua blocks");

		// Check first and only blockblock
		let first_block = &blocks[0];
		assert_eq!(first_block.x_get_str("lang")?, "lua");
		assert!(first_block.x_get_str("content")?.contains("some lua"));

		// -- Check Extruded Content
		let content = res.x_get_str("extruded_content")?;
		assert_contains(content, "This is some content");
		assert_contains(content, "and other block");
		assert_contains(content, "```rust"); // The rust block should remain in extruded content
		assert_contains(content, "```\n");
		assert_contains(content, "The end");

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_blocks_with_all_lang_and_extruded_content() -> Result<()> {
		// -- Setup & Fixtures
		// NOTE: the [[ ]] for multi line in lua breaks when line starts with ```, so work around
		let fx_script = r#"
local content = "This is some content\n"
content = content .. "\n```lua\n--some lua \n```\n"
content = content .. "and other block\n\n```rust\n//! some rust block \n```\n"
content = content .. "The end"

local blocks, extruded_content = aip.md.extract_blocks(content, {extrude = "content"})
return {
		blocks = blocks,
		extruded_content = extruded_content
}
		"#;

		// -- Exec
		let res = run_reflective_agent(fx_script, None).await?;

		// -- Check Blocks
		let blocks = res.pointer("/blocks").ok_or("Should have blocks")?;
		assert!(blocks.is_array());
		let blocks = blocks.as_array().ok_or("Should be arr")?;
		assert_eq!(blocks.len(), 2, "Should have found 2 blocks, lua and rust");

		// Check first block (order might vary, so check both)
		let block1 = &blocks[0];
		let block2 = &blocks[1];
		if block1.x_get_str("lang")? == "lua" {
			assert_eq!(block1.x_get_str("lang")?, "lua");
			assert!(block1.x_get_str("content")?.contains("some lua"));
			assert_eq!(block2.x_get_str("lang")?, "rust");
			assert!(block2.x_get_str("content")?.contains("some rust"));
		} else {
			assert_eq!(block1.x_get_str("lang")?, "rust");
			assert!(block1.x_get_str("content")?.contains("some rust"));
			assert_eq!(block2.x_get_str("lang")?, "lua");
			assert!(block2.x_get_str("content")?.contains("some lua"));
		}

		// -- Check Extruded Content
		let content = res.x_get_str("extruded_content")?;
		assert_contains(content, "This is some content");
		assert_contains(content, "and other block");
		assert_not_contains(content, "```lua");
		assert_not_contains(content, "```rust");
		assert_not_contains(content, "```"); // Ensure fence markers are removed from extruded content
		assert_contains(content, "The end");

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_meta() -> Result<()> {
		// -- Setup & Fixtures
		let lua = setup_lua(super::init_module, "md").await?;
		let lua_code = r#"
local content = [[
Some content

```toml
#!meta
some = "stuff"
```

some more content

```toml
#!meta
# Another meta block
num = 123
```

And this is the end
]]
local meta, remain = aip.md.extract_meta(content)
return {
   meta   = meta,
	 remain = remain
}
		"#;

		// -- Exec
		let res: Value = eval_lua(&lua, lua_code)?;

		// -- Check meta
		let meta = res.get("meta").ok_or("Should have meta")?;
		assert_eq!(meta.x_get_str("some")?, "stuff");
		assert_eq!(meta.x_get_i64("num")?, 123);

		// -- Check remain
		let remain = res.x_get_str("remain")?;
		assert_contains(remain, "Some content");
		assert_contains(remain, "some more content");
		assert_contains(remain, "And this is the end");
		assert_not_contains(remain, "Another meta block");
		assert_not_contains(remain, "num = 123");
		assert_not_contains(remain, "#!meta");
		// Should also remove code fences and blank lines resulting from block removal
		assert_not_contains(remain, "```toml");
		assert_not_contains(remain, "```");

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_refs_simple() -> Result<()> {
		// -- Setup & Fixtures
		let lua = setup_lua(super::init_module, "md").await?;
		let lua_code = r#"
local content = [[
Check out [this link](https://example.com) and [docs](docs/page.md).

Also see ![image](assets/photo.jpg) for reference.
]]
return aip.md.extract_refs(content)
		"#;

		// -- Exec
		let res: Value = eval_lua(&lua, lua_code)?;

		// -- Check
		assert!(res.is_array());
		let refs = res.as_array().ok_or("Res should be array")?;
		assert_eq!(refs.len(), 3, "Should have found 3 refs");

		// Check first ref (link)
		let first = &refs[0];
		assert_eq!(first.x_get_str("_type")?, "MdRef");
		assert_eq!(first.x_get_str("target")?, "https://example.com");
		assert_eq!(first.x_get_str("text")?, "this link");
		assert_eq!(first.get("inline").and_then(|v| v.as_bool()), Some(false));
		assert_eq!(first.x_get_str("kind")?, "Url");

		// Check second ref (file link)
		let second = &refs[1];
		assert_eq!(second.x_get_str("target")?, "docs/page.md");
		assert_eq!(second.x_get_str("text")?, "docs");
		assert_eq!(second.get("inline").and_then(|v| v.as_bool()), Some(false));
		assert_eq!(second.x_get_str("kind")?, "File");

		// Check third ref (image)
		let third = &refs[2];
		assert_eq!(third.x_get_str("target")?, "assets/photo.jpg");
		assert_eq!(third.x_get_str("text")?, "image");
		assert_eq!(third.get("inline").and_then(|v| v.as_bool()), Some(true));
		assert_eq!(third.x_get_str("kind")?, "File");

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_refs_skip_code_blocks() -> Result<()> {
		// -- Setup & Fixtures
		let lua = setup_lua(super::init_module, "md").await?;
		// NOTE: the [[ ]] for multi line in lua breaks when line starts with ```, so work around
		let lua_code = r#"
local content = "Here is a [real link](https://real.com).\n"
content = content .. "\n```\n[not a link](https://fake.com)\n```\n"
content = content .. "\nAnd [another real](page.md)."

return aip.md.extract_refs(content)
		"#;

		// -- Exec
		let res: Value = eval_lua(&lua, lua_code)?;

		// -- Check
		assert!(res.is_array());
		let refs = res.as_array().ok_or("Res should be array")?;
		assert_eq!(refs.len(), 2, "Should have found 2 refs (skipping code block)");
		assert_eq!(refs[0].x_get_str("target")?, "https://real.com");
		assert_eq!(refs[1].x_get_str("target")?, "page.md");

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_refs_nil_input() -> Result<()> {
		// -- Setup & Fixtures
		let lua = setup_lua(super::init_module, "md").await?;
		let lua_code = r#"
return aip.md.extract_refs(nil)
		"#;

		// -- Exec
		let res: Value = eval_lua(&lua, lua_code)?;

		// -- Check
		assert!(res.is_array());
		let refs = res.as_array().ok_or("Res should be array")?;
		assert_eq!(refs.len(), 0, "Should return empty list for nil input");

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_extract_refs_anchor() -> Result<()> {
		// -- Setup & Fixtures
		let lua = setup_lua(super::init_module, "md").await?;
		let lua_code = r#"
local content = [[
[go to section](#my-section)
]]
return aip.md.extract_refs(content)
		"#;

		// -- Exec
		let res: Value = eval_lua(&lua, lua_code)?;

		// -- Check
		assert!(res.is_array());
		let refs = res.as_array().ok_or("Res should be array")?;
		assert_eq!(refs.len(), 1, "Should have found 1 ref");
		assert_eq!(refs[0].x_get_str("target")?, "#my-section");
		assert_eq!(refs[0].x_get_str("kind")?, "Anchor");

		Ok(())
	}

	#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
	async fn test_lua_md_outer_block_content_or_raw() -> Result<()> {
		// -- Setup & Fixtures
		// NOTE: Here we put the ``` out of the multiline [[ ]]
		let fx_script = r#"
local content = "```" .. [[rust
fn main() {
    // Some nested blocks
    let example = ```typescript
    const x = 42;
    ```;
    println!("Hello!");
}
]] .. "```"

return aip.md.outer_block_content_or_raw(content)
		"#;

		// -- Exec
		let res = run_reflective_agent(fx_script, None).await?;

		// -- Check
		let content = res.as_str().ok_or("Should have res")?;
		assert!(content.contains("fn main()"));
		assert!(content.contains("const x = 42"));
		assert!(!content.contains("```rust")); // Should not contain the outer markers
		assert!(content.contains("```typescript")); // Should still contain inner markers

		// Test with raw content (no blocks)
		let fx_script_raw = r#"
local content = [[Just some plain
text without any code blocks]]

return aip.md.outer_block_content_or_raw(content)
		"#;

		let res_raw = run_reflective_agent(fx_script_raw, None).await?;
		let content_raw = res_raw.as_str().ok_or("Should have res")?;
		assert_eq!(content_raw, "Just some plain\ntext without any code blocks");

		Ok(())
	}
}

// endregion: --- Tests