use std::collections::BTreeMap;
use std::sync::OnceLock;

use arborium_highlight::{Injection, ParseResult, Span};
use wasmi::{
    Caller, CompilationMode, Config, Engine, ExternType, Instance, Linker, Memory, Module, Store,
    StoreLimits, StoreLimitsBuilder, WasmParams, WasmResults,
};

const MAX_MEMORY: usize = 128 << 20;
const WASM_PAGE_SIZE: usize = 1 << 16;
/// wasm-bindgen's heap keeps its first slots permanently populated: 1024..1028
/// are `undefined`, `null`, `true` and `false`, and everything below is padding.
const RESERVED_HANDLES: usize = 1028;

type Result<T> = std::result::Result<T, wasmi::Error>;

fn wasm_error(message: impl Into<String>) -> wasmi::Error {
    wasmi::Error::new(message)
}

/// The JS values serde-wasm-bindgen builds while serializing a parse result.
#[derive(Clone, Debug)]
enum Value {
    Undefined,
    Null,
    Bool(bool),
    Number(f64),
    String(String),
    Array(Vec<Value>),
    Object(BTreeMap<String, Value>),
}

impl Value {
    fn field(&self, name: &str) -> Result<&Value> {
        let Value::Object(fields) = self else {
            return Err(wasm_error("expected object from grammar plugin"));
        };
        fields
            .get(name)
            .ok_or_else(|| wasm_error(format!("grammar result is missing {name:?}")))
    }

    fn array(&self) -> Result<&[Value]> {
        let Value::Array(values) = self else {
            return Err(wasm_error("expected array from grammar plugin"));
        };
        Ok(values)
    }

    fn text(&self) -> Result<String> {
        let Value::String(value) = self else {
            return Err(wasm_error("expected string from grammar plugin"));
        };
        Ok(value.clone())
    }

    fn integer(&self) -> Result<u32> {
        let Value::Number(value) = *self else {
            return Err(wasm_error("expected number from grammar plugin"));
        };
        if !value.is_finite() || value.fract() != 0.0 || !(0.0..=u32::MAX as f64).contains(&value) {
            return Err(wasm_error("invalid integer from grammar plugin"));
        }
        Ok(value as u32)
    }

    fn boolean(&self) -> Result<bool> {
        let Value::Bool(value) = *self else {
            return Err(wasm_error("expected boolean from grammar plugin"));
        };
        Ok(value)
    }

    fn parse_result(&self) -> Result<ParseResult> {
        let spans = self
            .field("spans")?
            .array()?
            .iter()
            .map(|span| {
                Ok(Span {
                    start: span.field("start")?.integer()?,
                    end: span.field("end")?.integer()?,
                    capture: span.field("capture")?.text()?,
                    pattern_index: span.field("pattern_index")?.integer()?,
                })
            })
            .collect::<Result<_>>()?;
        let injections = self
            .field("injections")?
            .array()?
            .iter()
            .map(|injection| {
                Ok(Injection {
                    start: injection.field("start")?.integer()?,
                    end: injection.field("end")?.integer()?,
                    language: injection.field("language")?.text()?,
                    include_children: injection.field("include_children")?.boolean()?,
                })
            })
            .collect::<Result<_>>()?;
        Ok(ParseResult { spans, injections })
    }
}

/// Stands in for wasm-bindgen's JS-side object heap. Serialization builds each
/// value bottom-up and moves it into its parent, so a slot holds its value
/// until the guest drops or consumes the handle. The guest caches field-name
/// strings across calls, so handles must survive between parses.
struct Host {
    slots: Vec<Option<Value>>,
    free: Vec<usize>,
    limits: StoreLimits,
}

impl Host {
    fn new() -> Self {
        let mut slots = vec![Some(Value::Undefined); RESERVED_HANDLES];
        slots[1025] = Some(Value::Null);
        slots[1026] = Some(Value::Bool(true));
        slots[1027] = Some(Value::Bool(false));
        Self {
            slots,
            free: Vec::new(),
            limits: StoreLimitsBuilder::new()
                .memory_size(MAX_MEMORY)
                .instances(1)
                .memories(1)
                .tables(1)
                .trap_on_grow_failure(true)
                .build(),
        }
    }

    fn add(&mut self, value: Value) -> i32 {
        match self.free.pop() {
            Some(index) => {
                self.slots[index] = Some(value);
                index as i32
            }
            None => {
                self.slots.push(Some(value));
                (self.slots.len() - 1) as i32
            }
        }
    }

    fn get_mut(&mut self, handle: i32) -> Result<&mut Value> {
        self.slots
            .get_mut(handle as usize)
            .and_then(Option::as_mut)
            .ok_or_else(|| wasm_error("invalid wasm-bindgen object handle"))
    }

    fn take(&mut self, handle: i32) -> Result<Value> {
        let index = handle as usize;
        if index < RESERVED_HANDLES {
            return self.get_mut(handle).cloned();
        }
        let value = self
            .slots
            .get_mut(index)
            .and_then(Option::take)
            .ok_or_else(|| wasm_error("invalid wasm-bindgen object handle"))?;
        self.free.push(index);
        Ok(value)
    }
}

fn engine() -> &'static Engine {
    static ENGINE: OnceLock<Engine> = OnceLock::new();
    ENGINE.get_or_init(|| {
        let mut config = Config::default();
        config
            .ignore_custom_sections(true)
            .compilation_mode(CompilationMode::Lazy);
        Engine::new(&config)
    })
}

fn guest_string(caller: &Caller<'_, Host>, ptr: i32, len: i32) -> Result<String> {
    let memory = caller
        .get_export("memory")
        .and_then(wasmi::Extern::into_memory)
        .ok_or_else(|| wasm_error("grammar plugin has no memory export"))?;
    let mut bytes = vec![0; len as usize];
    memory.read(caller, ptr as usize, &mut bytes)?;
    String::from_utf8(bytes).map_err(|_| wasm_error("grammar plugin produced invalid UTF-8"))
}

/// Provides the imports serde-wasm-bindgen's serializer needs. The plugin also
/// imports a deserializer's worth of JS introspection for `apply_edit`, which
/// we never call, so anything else traps if reached. Import names carry a hash
/// of their JS shim and so are pinned to the wasm-bindgen release that built
/// the grammars.
fn link(linker: &mut Linker<Host>, module: &Module) -> Result<()> {
    for import in module.imports() {
        let ExternType::Func(ty) = import.ty() else {
            continue;
        };
        let (namespace, name) = (import.module(), import.name());
        match name {
            "__wbg_new_ab79df5bd7c26067" => {
                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>| {
                    c.data_mut().add(Value::Object(BTreeMap::new()))
                })?
            }
            "__wbg_new_a70fbab9066b301f" => {
                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>| {
                    c.data_mut().add(Value::Array(Vec::new()))
                })?
            }
            "__wbindgen_cast_0000000000000001" => {
                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>, value: f64| {
                    c.data_mut().add(Value::Number(value))
                })?
            }
            "__wbindgen_cast_0000000000000002" => linker.func_wrap(
                namespace,
                name,
                |mut c: Caller<'_, Host>, ptr: i32, len: i32| -> Result<i32> {
                    let value = guest_string(&c, ptr, len)?;
                    Ok(c.data_mut().add(Value::String(value)))
                },
            )?,
            "__wbindgen_object_clone_ref" => linker.func_wrap(
                namespace,
                name,
                |mut c: Caller<'_, Host>, handle: i32| -> Result<i32> {
                    // Only ever used on immutable strings, so a copy is indistinguishable
                    // from a shared reference.
                    let host = c.data_mut();
                    let value = host.get_mut(handle)?.clone();
                    Ok(host.add(value))
                },
            )?,
            "__wbindgen_object_drop_ref" => {
                linker.func_wrap(namespace, name, |mut c: Caller<'_, Host>, handle: i32| {
                    c.data_mut().take(handle).map(drop)
                })?
            }
            "__wbg_set_282384002438957f" => linker.func_wrap(
                namespace,
                name,
                |mut c: Caller<'_, Host>, target: i32, index: i32, value: i32| -> Result<()> {
                    let host = c.data_mut();
                    let value = host.take(value)?;
                    let Value::Array(values) = host.get_mut(target)? else {
                        return Err(wasm_error("indexed set on non-array"));
                    };
                    let index = index as usize;
                    if values.len() <= index {
                        values.resize(index + 1, Value::Undefined);
                    }
                    values[index] = value;
                    Ok(())
                },
            )?,
            "__wbg_set_6be42768c690e380" => linker.func_wrap(
                namespace,
                name,
                |mut c: Caller<'_, Host>, target: i32, key: i32, value: i32| -> Result<()> {
                    let host = c.data_mut();
                    let key = host.take(key)?.text()?;
                    let value = host.take(value)?;
                    let Value::Object(fields) = host.get_mut(target)? else {
                        return Err(wasm_error("property set on non-object"));
                    };
                    fields.insert(key, value);
                    Ok(())
                },
            )?,
            "__wbg___wbindgen_throw_6ddd609b62940d55" => linker.func_wrap(
                namespace,
                name,
                |c: Caller<'_, Host>, ptr: i32, len: i32| -> Result<()> {
                    Err(wasm_error(guest_string(&c, ptr, len)?))
                },
            )?,
            _ => {
                let message = format!("grammar plugin called unsupported import {name}");
                linker.func_new(namespace, name, ty.clone(), move |_, _, _| {
                    Err(wasm_error(message.clone()))
                })?
            }
        };
    }
    Ok(())
}

fn instantiate(module: &Module) -> Result<(Store<Host>, Instance, Memory)> {
    let mut store = Store::new(engine(), Host::new());
    store.limiter(|host| &mut host.limits);
    let mut linker = Linker::new(engine());
    link(&mut linker, module)?;
    let instance = linker.instantiate_and_start(&mut store, module)?;
    let memory = instance
        .get_memory(&store, "memory")
        .ok_or_else(|| wasm_error("grammar plugin has no memory export"))?;
    Ok((store, instance, memory))
}

/// A compiled grammar plus a snapshot of guest memory taken right after its
/// parser session was created, since compiling the highlight queries in
/// `create_session` is expensive and every runtime would otherwise repeat it.
pub(crate) struct Grammar {
    module: Module,
    memory: Vec<u8>,
    session: i32,
}

impl Grammar {
    pub(crate) fn new(bytes: &[u8]) -> Result<Self> {
        let module = Module::new(engine(), bytes)?;
        let (mut store, instance, memory) = instantiate(&module)?;
        let session = instance
            .get_typed_func::<(), i32>(&store, "create_session")?
            .call(&mut store, ())?;
        Ok(Self {
            memory: memory.data(&store).to_vec(),
            module,
            session,
        })
    }

    pub(crate) fn runtime(&self) -> Result<Runtime> {
        let (mut store, instance, memory) = instantiate(&self.module)?;
        let growth = self.memory.len().saturating_sub(memory.data_size(&store));
        memory.grow(&mut store, growth.div_ceil(WASM_PAGE_SIZE) as u64)?;
        memory.write(&mut store, 0, &self.memory)?;
        Ok(Runtime {
            store,
            instance,
            memory,
            session: self.session,
        })
    }
}

pub(crate) struct Runtime {
    store: Store<Host>,
    instance: Instance,
    memory: Memory,
    session: i32,
}

impl Runtime {
    fn call<P: WasmParams, R: WasmResults>(&mut self, name: &str, params: P) -> Result<R> {
        self.instance
            .get_typed_func::<P, R>(&self.store, name)?
            .call(&mut self.store, params)
    }

    /// Parses `source`. Errors may be traps that leave guest state arbitrary,
    /// so the runtime should be discarded after one.
    pub(crate) fn parse(&mut self, source: &str) -> Result<ParseResult> {
        let source_len = i32::try_from(source.len())
            .map_err(|_| wasm_error("source is too large to highlight"))?;
        let ptr: i32 = self.call("__wbindgen_export", (source_len, 1))?;
        self.memory
            .write(&mut self.store, ptr as usize, source.as_bytes())?;
        let () = self.call("set_text", (self.session, ptr, source_len))?;

        let retptr: i32 = self.call("__wbindgen_add_to_stack_pointer", -16)?;
        let () = self.call("parse", (retptr, self.session))?;
        let mut returned = [0; 12];
        self.memory
            .read(&self.store, retptr as usize, &mut returned)?;
        let _: i32 = self.call("__wbindgen_add_to_stack_pointer", 16)?;
        let [value, error, is_error] =
            std::array::from_fn(|i| i32::from_le_bytes(returned[4 * i..][..4].try_into().unwrap()));

        let host = self.store.data_mut();
        if is_error != 0 {
            return Err(wasm_error(match host.take(error)? {
                Value::String(message) => message,
                other => format!("{other:?}"),
            }));
        }
        host.take(value)?.parse_result()
    }
}