// prowl -- personal knowledge base
// Copyright (c) 2025 zenmaya <maya@zenmaya.xyz>
//
// This file is part of prowl.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use std::{
    path::{Path, PathBuf},
    sync::Arc,
};

use async_lsp::lsp_types::Url;
use futures::StreamExt;
use mlua::{FromLua, FromLuaMulti, Function, IntoLua, Lua, StdLib, Table, UserData, Value};
use tracing::{error, info};

use crate::document_store::DocumentStore;

/// A safe scope wrapper around `require`
/// that intercepts all calls to `require`
/// and forces no-caching
struct ScopedRequire<'a> {
    lua: &'a mlua::Lua,
    previous_require_fn: mlua::Function,
}

impl<'a> ScopedRequire<'a> {
    fn new(lua: &'a mlua::Lua) -> Self {
        let previous_require_fn: mlua::Function = lua.globals().get("require").unwrap();
        let require_fn = lua
            .create_function({
                let previous_require_fn = previous_require_fn.clone();
                move |lua, module_name: String| -> Result<mlua::Value, mlua::Error> {
                    lua.globals()
                        .get::<mlua::Table>("package")
                        .and_then(|package| {
                            package.get::<mlua::Table>("loaded").and_then(|loaded| {
                                loaded.set(module_name.clone(), mlua::Value::Nil)
                            })
                        })
                        .and_then(|_| previous_require_fn.call::<mlua::Value>(module_name.clone()))
                }
            })
            .unwrap();
        lua.globals().set("require", require_fn).unwrap();

        Self {
            lua,
            previous_require_fn,
        }
    }
}

impl<'a> Drop for ScopedRequire<'a> {
    fn drop(&mut self) {
        self.lua
            .globals()
            .set("require", self.previous_require_fn.clone())
            .unwrap();
    }
}

/// A safe wrapper around load functions in a piece of rust+lua code
/// that collects all the performed loads from files.
struct ScopedLoad<'a> {
    lua: &'a mlua::Lua,
    previous_load_fn: mlua::Function,
    loaded_files: std::sync::Arc<std::sync::Mutex<Vec<PathBuf>>>,
}

impl<'a> ScopedLoad<'a> {
    pub fn new(lua: &'a mlua::Lua) -> Self {
        let previous_load_fn: mlua::Function = lua.globals().get("load").unwrap();
        let loaded_files = std::sync::Arc::new(std::sync::Mutex::new(vec![]));

        let load_fn = lua
            .create_function_mut({
                let loaded_files = loaded_files.clone();
                let previous_load_fn = previous_load_fn.clone();
                move |_lua,
                      (chunk, chunkname, mode, env): (
                    mlua::Value,
                    mlua::Value,
                    mlua::Value,
                    mlua::Value,
                )|
                      -> Result<mlua::Value, mlua::Error> {
                    if let Some(chunkname) = chunkname.as_string_lossy() {
                        if chunkname.starts_with("@") {
                            let chunkname = chunkname.trim_start_matches('@');
                            loaded_files
                                .lock()
                                .unwrap()
                                .push(PathBuf::from(chunkname.to_string()));
                        } else {
                            info!(
                                ?chunkname,
                                "Loading chunk without catching it's loaded file"
                            );
                        }
                    }
                    previous_load_fn.call::<mlua::Value>((chunk, chunkname, mode, env))
                }
            })
            .unwrap();
        lua.globals().set("load", load_fn).unwrap();

        Self {
            lua,
            previous_load_fn,
            loaded_files,
        }
    }

    fn loaded_files(self) -> Vec<PathBuf> {
        self.loaded_files.lock().unwrap().clone()
    }
}

impl<'a> Drop for ScopedLoad<'a> {
    fn drop(&mut self) {
        self.lua
            .globals()
            .set("load", self.previous_load_fn.clone())
            .unwrap();
    }
}

pub struct DocumentStoreLua {
    pub store: Arc<DocumentStore>,
}

impl UserData for DocumentStoreLua {
    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M) {
        methods.add_async_method(
            "document_uris",
            async |_, this, ()| -> mlua::Result<Vec<String>> {
                let documents = this
                    .store
                    .document_uris()
                    .map(|url| url.path().to_string())
                    .collect::<Vec<_>>()
                    .await;
                Ok(documents)
            },
        );
        methods.add_async_method(
            "exact_str_to_documents",
            async |_, this, str: String| -> mlua::Result<Vec<String>> {
                let documents = this
                    .store
                    .exact_str_to_documents(&str)
                    .await
                    .iter()
                    .map(|url| {
                        let url = url.clone();
                        url.path().to_string()
                    })
                    .collect::<Vec<_>>();
                Ok(documents)
            },
        );
        methods.add_async_method(
            "prefix_to_documents",
            async |lua, this, pattern: String| -> mlua::Result<Vec<Table>> {
                let pattern = nucleo_matcher::pattern::Pattern::parse(
                    &pattern,
                    nucleo_matcher::pattern::CaseMatching::Smart,
                    nucleo_matcher::pattern::Normalization::Smart,
                );
                let documents = this
                    .store
                    .prefix_to_documents(pattern)
                    .await
                    .iter()
                    .map(|(url, rank)| {
                        let table = lua.create_table().unwrap();
                        table.set("url", url.path().to_string()).unwrap();
                        table.set("rank", *rank).unwrap();
                        table
                    })
                    .collect::<Vec<_>>();
                Ok(documents)
            },
        );
        methods.add_async_method(
            "uri_to_document",
            async |lua, this, mut uri: String| -> mlua::Result<Option<mlua::Value>> {
                if uri.starts_with("/") {
                    uri = format!("file://{}", uri);
                }
                let document = this
                    .store
                    .uri_to_document(&match Url::parse(&uri) {
                        Ok(url) => url,
                        Err(_) => {
                            error!(?uri, "Failed to parse uri");
                            return Ok(None);
                        }
                    })
                    .await
                    .map(|meta| meta.into_lua(&lua))
                    .transpose();
                document
            },
        );
        methods.add_async_method(
            "tag_to_documents",
            async |_lua, this, tags: Vec<String>| -> mlua::Result<Vec<String>> {
                let tags = tags.iter().map(|s| s.as_str()).collect::<Vec<_>>();
                let documents = this
                    .store
                    .tags_to_documents(&tags)
                    .await
                    .iter()
                    .map(|url| url.path().to_string())
                    .collect::<Vec<_>>();
                Ok(documents)
            },
        );
        methods.add_async_method(
            "document_out_edges",
            async |_lua, this, uri: String| -> mlua::Result<Vec<String>> {
                let document_urls = this
                    .store
                    .document_out_edges(Url::parse(&uri).unwrap())
                    .await
                    .iter()
                    .map(|url| url.path().to_string())
                    .collect::<Vec<_>>();
                Ok(document_urls)
            },
        );
        methods.add_async_method(
            "document_in_edges",
            async |_lua, this, uri: String| -> mlua::Result<Vec<String>> {
                error!(?uri, "document_in_edges");
                let document_urls = this
                    .store
                    .document_in_edges(Url::parse(&uri).unwrap())
                    .await
                    .iter()
                    .map(|url| {
                        error!(?url, "document_in_edges");
                        url.path().to_string()
                    })
                    .collect::<Vec<_>>();
                error!(?document_urls, "document_in_edges");
                Ok(document_urls)
            },
        );
    }
}

pub async fn init() -> Lua {
    // SAFETY:
    // we need to use unsafe_new() here because we need to load fennel
    // and fennel requires debug library, which is unsafe to use
    let lua = unsafe { Lua::unsafe_new() };
    load_fennel(&lua).await.expect("Failed to load fennel");

    let package: Table = lua.globals().get("package").unwrap();
    let mut path = package.get::<String>("path").unwrap();
    path.push_str(";");
    path.push_str(crate::share_path().to_str().unwrap());
    path.push_str("/lua/?.lua;");
    path.push_str(crate::share_path().to_str().unwrap());
    path.push_str("/lua/?.fnl");
    package.set("path", path).unwrap();

    lua
}

#[derive(Debug, Clone)]
pub struct LuaWriter(Vec<u8>);
impl LuaWriter {
    pub fn new() -> LuaWriter {
        Self(Vec::new())
    }

    pub fn from_lua(value: Value) -> mlua::Result<Self> {
        if let Value::UserData(data) = value {
            return data.take::<LuaWriter>();
        }
        Err(mlua::Error::FromLuaConversionError {
            from: "lua: LuaWriter",
            to: "rust: LuaWriter".to_string(),
            message: Some("Failed to convert Lua value to LuaWriter".to_string()),
        })
    }

    pub fn into_inner(self) -> Vec<u8> {
        self.0
    }
}

impl UserData for LuaWriter {
    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M) {
        methods.add_method_mut("write", |_, this, value: String| {
            this.0.extend(value.as_bytes());
            Ok(())
        });
        methods.add_method_mut("write-byte", |_, this, value: u8| {
            this.0.push(value);
            Ok(())
        });
    }
}

async fn load_fennel(lua: &Lua) -> mlua::Result<()> {
    lua.load_std_libs(StdLib::ALL).map_err(|err| {
        error!(?err, "Failed to load standard libraries");
        err
    })?;
    let fennel_path = crate::share_path().join("lua/fennel.lua");
    let fennel: Table = lua.load(fennel_path).eval_async().await.map_err(|err| {
        error!(?err, "Failed to load fennel");
        err
    })?;
    let fennel_install: Function = fennel.get("install")?;
    let fennel_install_result: Value = fennel_install.call_async(mlua::Nil).await?;
    lua.globals().set("fennel", fennel_install_result)?;
    Ok(())
}

async fn fennel_eval<T: FromLuaMulti>(
    lua: &Lua,
    code: &str,
    path: Option<&Path>,
) -> mlua::Result<T> {
    let fennel: Table = lua.globals().get("fennel")?;
    let eval: Function = fennel.get("eval")?;

    let opts = lua.create_table()?;
    if let Some(path) = path {
        opts.set("filename", path.to_string_lossy())?;
    }

    eval.call_async((code, opts)).await
}

async fn lua_eval<T: FromLuaMulti>(lua: &Lua, code: &str, path: &Path) -> mlua::Result<T> {
    lua.load(code)
        .set_name(format!("@{path}", path = path.to_string_lossy()))
        .eval_async()
        .await
}

pub async fn eval_file<T: FromLuaMulti>(lua: &Lua, path: &Path) -> mlua::Result<T> {
    let file = tokio::fs::read_to_string(path).await?;
    if path.extension().map(|s| s == "fnl").unwrap_or(false) {
        fennel_eval(lua, file.as_str(), Some(path)).await
    } else if path.extension().map(|s| s == "lua").unwrap_or(false) {
        lua_eval(lua, file.as_str(), path).await
    } else {
        error!("File must be either .fnl or .lua");
        panic!();
    }
}

pub async fn load_config_collect_loads(lua: &Lua) -> Result<(Table, Vec<PathBuf>), mlua::Error> {
    let scope = ScopedLoad::new(lua);
    Ok((load_config(lua).await?, scope.loaded_files()))
}

pub fn config_file_path() -> PathBuf {
    let fennel_conf = PathBuf::from("prowl.fnl");
    let lua_conf = PathBuf::from("project.lua");

    if fennel_conf.exists() {
        return fennel_conf;
    } else if lua_conf.exists() {
        return lua_conf;
    }
    fennel_conf
}

async fn access_config_entry_inner<V: FromLua>(path: &[&str], config: Table) -> mlua::Result<V> {
    let mut config_subtable = config;
    let mut path_iter = path.iter().peekable();
    while let Some(entry) = path_iter.next() {
        if path_iter.peek() == None {
            return config_subtable.get(*entry);
        }
        config_subtable = config_subtable.get::<Table>(*entry)?;
    }
    unreachable!()
}

pub async fn access_config_entry<V: FromLua>(lua: &Lua, path: &[&str]) -> V {
    let config = load_config_safe(lua).await;
    match access_config_entry_inner::<V>(path, config).await {
        Ok(v) => v,
        Err(err) => {
            error!(
                ?err,
                ?path,
                "Failed to find path in config table, loading default config value"
            );
            access_config_entry_inner::<V>(path, load_default_config(lua).await)
                .await
                .expect("The entry should always appear in the default config")
        }
    }
}

const DEFAULT_CONFIG: &'static str = include_str!("../../prowl.fnl");

pub async fn load_default_config(lua: &Lua) -> Table {
    fennel_eval(lua, DEFAULT_CONFIG, None).await.unwrap()
}

/// Caches last working config
static LAST_CONFIG: tokio::sync::Mutex<Option<Table>> = tokio::sync::Mutex::const_new(None);

/// Load the configuration file
/// and return the loaded table.
/// If the file does not exist, it will be created
/// with a default configuration.
/// If the file is not a valid fennel file, it will return the error and the last
/// loaded config.
pub async fn load_config(lua: &Lua) -> Result<Table, mlua::Error> {
    let _scope = ScopedRequire::new(lua);

    let config_file_path = config_file_path();

    if !config_file_path.exists() {
        tokio::fs::create_dir_all(config_file_path.parent().unwrap()).await?;
        if config_file_path
            .extension()
            .map(|s| s == "fnl")
            .unwrap_or(false)
        {
            tokio::fs::write(&config_file_path, DEFAULT_CONFIG).await?;
        } else if config_file_path
            .extension()
            .map(|s| s == "lua")
            .unwrap_or(false)
        {
            todo!()
        } else {
            unreachable!()
        }
    }
    let config = eval_file::<Table>(lua, &config_file_path).await?;
    let mut last_config = LAST_CONFIG.lock().await;
    *last_config = Some(config.clone());
    Ok(config)
}

pub async fn load_config_safe(lua: &Lua) -> Table {
    match load_config(lua).await {
        Ok(config) => config,
        Err(err) => {
            if let Some(config) = LAST_CONFIG.lock().await.clone() {
                error!(?err, "Failed to load config, loading last working config");
                config
            } else {
                error!(
                    ?err,
                    "Failed to load config and no last working config found, loading default config"
                );
                load_default_config(lua).await
            }
        }
    }
}