//! A sandboxed `require` for loading Lua modules from files.
//!
//! Module names use stock Lua dot notation (`require "db.migrations"`) and are
//! resolved against a fixed list of root directories — the main script's real
//! directory (the CWD in the REPL) followed by the entries of
//! `$A_INCLUDE_PATHS` — through the standard `?.lua` / `?/init.lua` templates.
//! Because a name may only contain identifier-like segments, a module can
//! never name a file outside the roots.
//!
//! Each file is executed once and its return value cached by canonical path,
//! so repeated requires — even under different names or through symlinks —
//! return the same value.

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::rc::Rc;

use mlua::prelude::{LuaResult, LuaTable, LuaValue};
use mlua::Lua;

/// Globals pre-seeded into the module cache, so `require "http"` returns the
/// same table as the global: scripts read more portably and LSPs get an
/// anchor for the stdlib. Covers the project stdlib plus the stock libraries
/// the sandbox loads.
const BUILTINS: &[&str] = &[
    "utils",
    "log",
    "sqlite",
    "http",
    "fs",
    "task",
    "tui",
    "coroutine",
    "math",
    "os",
    "string",
    "table",
    "utf8",
];

/// Strip a leading UTF-8 BOM and a shebang line from Lua source, as stock Lua
/// does. The shebang's newline is kept so error messages still report correct
/// line numbers. Used for the main script and for every required file.
pub(crate) fn prepare_chunk(source: &[u8]) -> &[u8] {
    let source = source.strip_prefix(b"\xEF\xBB\xBF").unwrap_or(source);
    if source.first() == Some(&b'#') {
        let nl = source.iter().position(|&b| b == b'\n').unwrap_or(source.len());
        &source[nl..]
    } else {
        source
    }
}

/// Install the global `require` function, searching `roots` in order.
pub(crate) fn install(lua: &Lua, roots: Vec<PathBuf>) -> LuaResult<()> {
    // Cache of loaded modules, held Lua-side so the values stay GC-anchored.
    // Built-ins are keyed by name, file modules by canonical path; the key
    // spaces cannot collide because canonical paths are absolute.
    let loaded = lua.create_table()?;
    for name in BUILTINS {
        let value: LuaValue = lua.globals().get(*name)?;
        if !value.is_nil() {
            loaded.set(*name, value)?;
        }
    }

    let roots = Rc::new(roots);
    // Module name -> canonical path it resolved to; repeated requires of the
    // same name skip the filesystem probing.
    let resolved: Rc<RefCell<HashMap<String, String>>> = Rc::default();
    // Canonical paths of modules currently executing, to catch require cycles.
    let in_flight: Rc<RefCell<HashSet<String>>> = Rc::default();

    let require = lua.create_async_function(move |lua, name: String| {
        let loaded = loaded.clone();
        let roots = roots.clone();
        let resolved = resolved.clone();
        let in_flight = in_flight.clone();
        async move { require_impl(lua, name, loaded, roots, resolved, in_flight).await }
    })?;
    lua.globals().set("require", require)
}

async fn require_impl(
    lua: Lua,
    name: String,
    loaded: LuaTable,
    roots: Rc<Vec<PathBuf>>,
    resolved: Rc<RefCell<HashMap<String, String>>>,
    in_flight: Rc<RefCell<HashSet<String>>>,
) -> LuaResult<(LuaValue, Option<String>)> {
    // Built-in module?
    let builtin: LuaValue = loaded.get(name.as_str())?;
    if !builtin.is_nil() {
        return Ok((builtin, None));
    }

    // A name that already resolved to a loaded file?
    if let Some(key) = resolved.borrow().get(&name).cloned() {
        let value: LuaValue = loaded.get(key.as_str())?;
        if !value.is_nil() {
            return Ok((value, Some(key)));
        }
    }

    let Some(rel) = module_relpath(&name) else {
        return Err(mlua::Error::runtime(format!(
            "invalid module name '{name}' (expected dot-separated segments of [A-Za-z0-9_-], like 'dir.mod')"
        )));
    };

    // Probe every root so a hit shadowed by an earlier one can be warned about.
    let mut hits: Vec<PathBuf> = Vec::new();
    let mut misses: Vec<PathBuf> = Vec::new();
    for root in roots.iter() {
        let base = root.join(&rel);
        for candidate in [base.with_extension("lua"), base.join("init.lua")] {
            match tokio::fs::metadata(&candidate).await {
                Ok(meta) if meta.is_file() => hits.push(candidate),
                _ => misses.push(candidate),
            }
        }
    }

    let Some(found) = hits.first() else {
        let mut msg = format!("module '{name}' not found:");
        for miss in &misses {
            msg.push_str(&format!("\n\tno file '{}'", miss.display()));
        }
        return Err(mlua::Error::runtime(msg));
    };
    for shadowed in &hits[1..] {
        log::warn!(
            "require '{name}': using '{}', ignoring shadowed '{}'",
            found.display(),
            shadowed.display()
        );
    }

    let real = tokio::fs::canonicalize(found).await.map_err(|e| {
        mlua::Error::runtime(format!(
            "require '{name}': cannot resolve '{}': {e}",
            found.display()
        ))
    })?;
    let key = real.display().to_string();

    // The same file may already be loaded under another name or via a symlink.
    let value: LuaValue = loaded.get(key.as_str())?;
    if !value.is_nil() {
        resolved.borrow_mut().insert(name, key.clone());
        return Ok((value, Some(key)));
    }

    if !in_flight.borrow_mut().insert(key.clone()) {
        return Err(mlua::Error::runtime(format!(
            "circular require of module '{name}' ('{key}')"
        )));
    }
    // Clears the in-flight marker even when the load errors or is cancelled,
    // so a failed module can be required again.
    let _guard = InFlightGuard {
        set: in_flight.clone(),
        key: key.clone(),
    };

    let source = tokio::fs::read(&real)
        .await
        .map_err(|e| mlua::Error::runtime(format!("require '{name}': cannot read '{key}': {e}")))?;

    let value: LuaValue = lua
        .load(prepare_chunk(&source))
        // "@" marks the chunk name as a file path in error messages / tracebacks
        .set_name(format!("@{key}"))
        // like stock Lua, the chunk receives the module name and file path as `...`
        .call_async((name.as_str(), key.as_str()))
        .await?;

    // A module that returns nothing is cached as `true`, as stock require does.
    let value = if value.is_nil() {
        LuaValue::Boolean(true)
    } else {
        value
    };
    loaded.set(key.as_str(), value.clone())?;
    resolved.borrow_mut().insert(name, key.clone());
    Ok((value, Some(key)))
}

/// Convert a dot-separated module name into a relative path, or `None` when
/// the name is malformed. Restricting segments to identifier-like characters
/// is what keeps `require` inside its roots: `..`, `/` and absolute paths are
/// unrepresentable.
fn module_relpath(name: &str) -> Option<PathBuf> {
    if name.is_empty() {
        return None;
    }
    let mut rel = PathBuf::new();
    for segment in name.split('.') {
        let ok = !segment.is_empty()
            && segment
                .bytes()
                .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-');
        if !ok {
            return None;
        }
        rel.push(segment);
    }
    Some(rel)
}

struct InFlightGuard {
    set: Rc<RefCell<HashSet<String>>>,
    key: String,
}

impl Drop for InFlightGuard {
    fn drop(&mut self) {
        self.set.borrow_mut().remove(&self.key);
    }
}