//! Async Arborium WASM highlighting powered by Wasmi.
//!
//! Callers own grammar retrieval and integrity policy through [`ModuleLoader`].
//! Arbrisseau caches initialized grammars in [`Options`] and resolves injections recursively.

mod runtime;

use std::collections::HashMap;
use std::error::Error as StdError;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};

use arborium_highlight::{Injection, ParseResult};
use async_lock::OnceCell;

pub use arborium_highlight::Span;

const MAX_INJECTION_DEPTH: u32 = 3;

pub type LoaderError = Box<dyn StdError + Send + Sync + 'static>;
pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("unsupported language {language:?}")]
    UnsupportedLanguage { language: String },

    #[error("failed to load {language} grammar")]
    LoadModule {
        language: String,
        #[source]
        source: LoaderError,
    },

    #[error("failed to initialize {language} grammar")]
    Initialize {
        language: String,
        #[source]
        source: wasmi::Error,
    },

    #[error("failed to parse {language} source")]
    Parse {
        language: String,
        #[source]
        source: wasmi::Error,
    },
}

/// The asynchronous result of loading a grammar's raw WebAssembly bytes.
/// `None` means that the requested language is unsupported.
pub type ModuleFuture<'a> =
    Pin<Box<dyn Future<Output = std::result::Result<Option<Vec<u8>>, LoaderError>> + Send + 'a>>;

/// Loads a grammar from a caller-defined source.
pub type ModuleLoader = fn(language: &str) -> ModuleFuture;

/// Shared configuration and initialized-grammar cache for highlighting calls.
pub struct Options {
    /// Callback used when a language's module is first needed.
    pub load_module: ModuleLoader,
    /// A cached `None` marks a language the loader reported as unsupported.
    grammars: Mutex<HashMap<String, Arc<OnceCell<Option<runtime::Grammar>>>>>,
}

impl Options {
    /// Creates options with an initially empty in-memory grammar cache.
    pub fn new(load_module: ModuleLoader) -> Self {
        Self {
            load_module,
            grammars: Mutex::new(HashMap::new()),
        }
    }

    async fn runtime(&self, language: &str) -> Result<Option<runtime::Runtime>> {
        // Cloning the cell out releases the map lock, so unrelated languages
        // stay loadable while this one initializes.
        let cell = self
            .grammars
            .lock()
            .unwrap_or_else(|error| error.into_inner())
            .entry(language.to_owned())
            .or_default()
            .clone();
        let grammar = cell
            .get_or_try_init(|| async {
                let bytes =
                    (self.load_module)(language)
                        .await
                        .map_err(|source| Error::LoadModule {
                            language: language.to_owned(),
                            source,
                        })?;
                bytes
                    .as_deref()
                    .map(runtime::Grammar::new)
                    .transpose()
                    .map_err(|source| Error::Initialize {
                        language: language.to_owned(),
                        source,
                    })
            })
            .await?;
        grammar
            .as_ref()
            .map(runtime::Grammar::runtime)
            .transpose()
            .map_err(|source| Error::Initialize {
                language: language.to_owned(),
                source,
            })
    }
}

/// Reuses Wasmi instances and parser sessions across highlighting calls.
#[derive(Default)]
pub struct Highlighter {
    runtimes: HashMap<String, runtime::Runtime>,
}

impl Highlighter {
    pub fn new() -> Self {
        Self::default()
    }

    async fn parse(
        &mut self,
        language: &str,
        source: &str,
        options: &Options,
    ) -> Result<Option<ParseResult>> {
        if !self.runtimes.contains_key(language) {
            let Some(runtime) = options.runtime(language).await? else {
                return Ok(None);
            };
            self.runtimes.insert(language.to_owned(), runtime);
        }
        match self.runtimes.get_mut(language).unwrap().parse(source) {
            Ok(result) => Ok(Some(result)),
            Err(source) => {
                // A trap leaves guest state arbitrary, so discard the session.
                self.runtimes.remove(language);
                Err(Error::Parse {
                    language: language.to_owned(),
                    source,
                })
            }
        }
    }

    /// Highlights a source block and recursively includes spans from injected languages.
    pub async fn highlight(
        &mut self,
        language: &str,
        source: &str,
        options: &Options,
    ) -> Result<Vec<Span>> {
        let result = self
            .parse(language, source, options)
            .await?
            .ok_or_else(|| Error::UnsupportedLanguage {
                language: language.to_owned(),
            })?;
        let mut spans = result.spans;
        let mut pending = Vec::new();
        push_injections(
            &mut pending,
            source,
            result.injections,
            0,
            MAX_INJECTION_DEPTH,
        );

        while let Some(injection) = pending.pop() {
            let Some(result) = self
                .parse(&injection.language, injection.text, options)
                .await?
            else {
                continue;
            };
            spans.extend(result.spans.into_iter().map(|mut span| {
                span.start += injection.offset;
                span.end += injection.offset;
                span
            }));
            push_injections(
                &mut pending,
                injection.text,
                result.injections,
                injection.offset,
                injection.depth,
            );
        }

        Ok(spans)
    }
}

struct PendingInjection<'a> {
    language: String,
    text: &'a str,
    offset: u32,
    depth: u32,
}

/// Queues injections in reverse so that popping the stack visits them in source order.
fn push_injections<'a>(
    pending: &mut Vec<PendingInjection<'a>>,
    source: &'a str,
    injections: Vec<Injection>,
    offset: u32,
    depth: u32,
) {
    if depth == 0 {
        return;
    }
    pending.extend(injections.into_iter().rev().filter_map(|injection| {
        let text = source.get(injection.start as usize..injection.end as usize)?;
        (!text.is_empty()).then_some(PendingInjection {
            language: injection.language,
            text,
            offset: offset + injection.start,
            depth: depth - 1,
        })
    }));
}