use std::{
    collections::BTreeMap,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
};

use mlua::{Function, HookTriggers, Lua, LuaSerdeExt, Table, VmState};
use serde_json::Value;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

use super::{
    CapabilityRequest, ContextDataMode, Contribution, ContributionKind, HostCapabilities,
    InvocationEvent, InvocationMode, InvocationOutcome, PluginError, PluginSource,
};

struct Registered {
    contribution: Contribution,
    callback: Function,
}

pub(super) struct Generation {
    pub source: String,
    pub config: Value,
    pub retirement: CancellationToken,
    lua: Lua,
    callbacks: BTreeMap<(ContributionKind, String), Registered>,
}

// Result::map_err needs an owned-error adapter; no Lua error escapes the module.
#[allow(clippy::needless_pass_by_value)]
fn lua_error(error: mlua::Error) -> PluginError {
    PluginError::Lua(error.to_string())
}

// Unknown fields must not silently imply an unimplemented contribution form.
fn validate_fields(table: &Table, allowed: &[&str]) -> Result<(), PluginError> {
    for key in table.clone().pairs::<String, mlua::Value>() {
        let (key, _) = key.map_err(lua_error)?;
        if !allowed.contains(&key.as_str()) {
            return Err(PluginError::Invalid(format!(
                "unknown declaration field {key}"
            )));
        }
    }
    Ok(())
}

/// One evaluated declaration chunk plus its private VM. The initialization
/// instruction deadline stays installed: callers walk bounded declarations, not
/// arbitrary Lua, and remove the hook once their pass completes.
pub(super) struct EvaluatedChunk {
    pub lua: Lua,
    pub declaration: Table,
}

/// Build the plugin VM and evaluate one declarative chunk. Host capabilities are
/// call-scoped, so this VM holds only the chunk, its owned `config`, and the
/// standard library.
pub(super) fn evaluate(source: &PluginSource) -> Result<EvaluatedChunk, PluginError> {
    if source.source.len() > 1024 * 1024 {
        return Err(PluginError::Invalid(format!(
            "plugin {name}: source exceeds 1 MiB",
            name = source.name
        )));
    }
    // Trusted local plugins, not a hostile-code sandbox. The full safe standard
    // library is present (io, os, load, print) short of the debug library; the
    // limits below are resource bounds, not isolation promises.
    let lua = Lua::new();
    lua.set_memory_limit(16 * 1024 * 1024).map_err(lua_error)?;
    lua.globals().set("null", lua.null()).map_err(lua_error)?;
    // Serialize into this VM before activation. Lua mutations cannot change the
    // desired JSON retained for reconciliation or any other generation's config.
    let config = lua.to_value(&source.config).map_err(lua_error)?;
    lua.globals().set("config", config).map_err(lua_error)?;
    // Initialization is declarative and cannot await. Bound accidental busy loops.
    let instructions = AtomicUsize::new(0);
    lua.set_hook(
        HookTriggers::new().every_nth_instruction(1000),
        move |_, _| {
            if instructions.fetch_add(1000, Ordering::Relaxed) >= 1_000_000 {
                return Err(mlua::Error::RuntimeError(
                    "initialization instruction budget exceeded".into(),
                ));
            }
            Ok(VmState::Continue)
        },
    )
    .map_err(lua_error)?;
    let declaration: Table = lua
        .load(&source.source)
        .set_name(&source.name)
        .eval()
        .map_err(lua_error)?;
    validate_fields(
        &declaration,
        &[
            "exports",
            "commands",
            "hooks",
            "context_strategies",
            "compactors",
            "providers",
        ],
    )?;
    Ok(EvaluatedChunk { lua, declaration })
}

/// Dense declaration arrays share one shape: 1..=128 table entries.
pub(super) fn dense_array(entries: &Table, field: &str) -> Result<Vec<Table>, PluginError> {
    let mut indexed = Vec::new();
    for entry in entries.clone().pairs::<usize, Table>() {
        let (index, entry) = entry.map_err(lua_error)?;
        if index == 0 || index > 128 {
            return Err(PluginError::Invalid(
                "declarations must be dense arrays of at most 128 callbacks".into(),
            ));
        }
        indexed.push((index, entry));
    }
    indexed.sort_by_key(|(index, _)| *index);
    if indexed
        .iter()
        .map(|(index, _)| *index)
        .ne(1..=indexed.len())
    {
        return Err(PluginError::Invalid(format!(
            "{field} declarations must be dense arrays"
        )));
    }
    Ok(indexed.into_iter().map(|(_, entry)| entry).collect())
}

impl Generation {
    pub fn load(source: &PluginSource) -> Result<Self, PluginError> {
        let EvaluatedChunk { lua, declaration } = evaluate(source)?;

        let id = Uuid::new_v4();
        let mut callbacks = BTreeMap::new();
        for (field, kind) in [
            ("exports", ContributionKind::Export),
            ("commands", ContributionKind::Command),
            ("hooks", ContributionKind::Hook),
            ("context_strategies", ContributionKind::ContextStrategy),
            ("compactors", ContributionKind::Compactor),
            ("providers", ContributionKind::Provider),
        ] {
            let Some(entries) = declaration.get::<Option<Table>>(field).map_err(lua_error)? else {
                continue;
            };
            for entry in dense_array(&entries, field)? {
                validate_fields(
                    &entry,
                    &["name", "description", "mode", "context_data", "run"],
                )?;
                let context_data = entry
                    .get::<Option<String>>("context_data")
                    .map_err(lua_error)?;
                if context_data.is_some()
                    && !matches!(
                        kind,
                        ContributionKind::ContextStrategy | ContributionKind::Compactor
                    )
                {
                    return Err(PluginError::Invalid(format!(
                        "{field} callbacks cannot declare context_data"
                    )));
                }
                let context_data = match context_data.as_deref() {
                    None | Some("json") => ContextDataMode::Json,
                    Some("refs_v1") => ContextDataMode::RefsV1,
                    Some(other) => {
                        return Err(PluginError::Invalid(format!(
                            "unknown context data mode {other}"
                        )))
                    }
                };
                let name: String = entry.get("name").map_err(lua_error)?;
                super::names::validate_contribution_name(kind, &name)?;
                let mode = match entry
                    .get::<Option<String>>("mode")
                    .map_err(lua_error)?
                    .as_deref()
                {
                    None if kind == ContributionKind::Provider => InvocationMode::Streaming,
                    None | Some("unary") => InvocationMode::Unary,
                    Some("streaming") => InvocationMode::Streaming,
                    Some(other) => {
                        return Err(PluginError::Invalid(format!(
                            "unknown invocation mode {other}"
                        )))
                    }
                };
                if matches!(
                    kind,
                    ContributionKind::ContextStrategy | ContributionKind::Compactor
                ) && mode != InvocationMode::Unary
                {
                    return Err(PluginError::Invalid(format!(
                        "{field} callbacks must be unary"
                    )));
                }
                let registered = Registered {
                    contribution: Contribution {
                        plugin: source.name.clone(),
                        generation: id,
                        kind,
                        name: name.clone(),
                        description: entry
                            .get::<Option<String>>("description")
                            .map_err(lua_error)?
                            .unwrap_or_default(),
                        mode,
                        context_data,
                    },
                    callback: entry.get("run").map_err(lua_error)?,
                };
                if callbacks.insert((kind, name.clone()), registered).is_some() {
                    return Err(PluginError::Invalid(format!(
                        "duplicate {field} callback {name}"
                    )));
                }
            }
        }
        lua.remove_hook();
        // Every running coroutine cooperatively yields, even pure Lua busy loops.
        // mlua releases its internal VM lock on Pending; no registry/session lock
        // crosses an async capability wait. The caller/retirement select then wins.
        lua.set_global_hook(HookTriggers::new().every_nth_instruction(10_000), |_, _| {
            Ok(VmState::Yield)
        })
        .map_err(lua_error)?;
        Ok(Self {
            source: source.source.clone(),
            config: source.config.clone(),
            retirement: CancellationToken::new(),
            lua,
            callbacks,
        })
    }

    pub fn contributions(&self) -> impl Iterator<Item = &Contribution> {
        self.callbacks
            .values()
            .map(|registered| &registered.contribution)
    }
}

#[derive(Clone)]
struct CallScope {
    invocation: Uuid,
    contribution: Contribution,
    capabilities: Arc<dyn HostCapabilities>,
    cancellation: CancellationToken,
    retirement: CancellationToken,
    closed: CancellationToken,
}

impl CallScope {
    fn context(&self, lua: &Lua, sender: mpsc::Sender<InvocationEvent>) -> mlua::Result<Table> {
        let context = lua.create_table()?;
        let call_scope = self.clone();
        context.set("call", lua.create_async_function(move |lua, (operation, input): (String, mlua::Value)| {
            let scope = call_scope.clone();
            async move {
                let request = CapabilityRequest {
                    invocation: scope.invocation,
                    operation,
                    input: lua.from_value(input)?,
                };
                let output = tokio::select! {
                    biased;
                    () = scope.retirement.cancelled() => return Err(mlua::Error::RuntimeError("retired".into())),
                    () = scope.cancellation.cancelled() => return Err(mlua::Error::RuntimeError("cancelled".into())),
                    () = scope.closed.cancelled() => return Err(mlua::Error::RuntimeError("invocation closed".into())),
                    // Construct the adapter future only after the cancellation fences
                    // have been polled. select! eagerly evaluates branch expressions.
                    output = async { scope.capabilities.call(request, scope.cancellation.child_token()).await } => output.map_err(mlua::Error::RuntimeError)?,
                };
                lua.to_value(&output)
            }
        })?)?;
        let emit_scope = self.clone();
        context.set("emit", lua.create_async_function(move |lua, value: mlua::Value| {
            let scope = emit_scope.clone();
            let sender = sender.clone();
            async move {
                if scope.contribution.mode != InvocationMode::Streaming {
                    return Err(mlua::Error::RuntimeError("unary callbacks cannot emit".into()));
                }
                let value = lua.from_value(value)?;
                tokio::select! {
                    biased;
                    () = scope.retirement.cancelled() => Err(mlua::Error::RuntimeError("retired".into())),
                    () = scope.cancellation.cancelled() => Err(mlua::Error::RuntimeError("cancelled".into())),
                    () = scope.closed.cancelled() => Err(mlua::Error::RuntimeError("invocation closed".into())),
                    result = sender.send(InvocationEvent::Update(value)) => result.map_err(mlua::Error::external),
                }
            }
        })?)?;
        Ok(context)
    }
}

#[tracing::instrument(skip_all, fields(plugin = %contribution.plugin, generation = %contribution.generation))]
pub(super) async fn run(
    generation: Arc<Generation>,
    contribution: Contribution,
    input: Value,
    capabilities: Arc<dyn HostCapabilities>,
    cancellation: CancellationToken,
    sender: mpsc::Sender<InvocationEvent>,
) {
    let closed = CancellationToken::new();
    let _close_on_exit = closed.clone().drop_guard();
    // Retained Lua contexts must not keep working after their originating invocation.
    // Cancel this adapter token on both success and failure, independently of delivery.
    let work = cancellation.child_token();
    let _cancel_work_on_exit = work.clone().drop_guard();
    let scope = CallScope {
        invocation: Uuid::new_v4(),
        contribution: contribution.clone(),
        capabilities,
        cancellation: work,
        retirement: generation.retirement.clone(),
        closed,
    };
    let execution = async {
        let context = scope
            .context(&generation.lua, sender.clone())
            .map_err(lua_error)?;
        let input = generation.lua.to_value(&input).map_err(lua_error)?;
        let registered = generation
            .callbacks
            .get(&(contribution.kind, contribution.name.clone()))
            .ok_or(PluginError::Missing)?;
        // call_async owns a recyclable coroutine. Its drop resets the thread and
        // drops pending Rust capability futures on cancellation, not just their result.
        let output: mlua::Value = registered
            .callback
            .call_async((input, context))
            .await
            .map_err(lua_error)?;
        let output = generation.lua.from_value(output).map_err(lua_error)?;
        Ok(match contribution.mode {
            InvocationMode::Unary => InvocationOutcome::Unary(output),
            InvocationMode::Streaming => InvocationOutcome::Streamed(output),
        })
    };
    let result = tokio::select! {
        biased;
        () = generation.retirement.cancelled() => Err(PluginError::Retired),
        () = cancellation.cancelled() => Err(PluginError::Cancelled),
        result = execution => result,
    };
    scope.closed.cancel();
    scope.cancellation.cancel();
    tokio::select! {
        biased;
        () = generation.retirement.cancelled() => {},
        () = cancellation.cancelled() => {},
        result = sender.send(InvocationEvent::Finished(result)) => {
            if result.is_err() {
                tracing::debug!("plugin invocation consumer dropped");
            }
        }
    }
}