mod registry;

use std::io;
use std::sync::Arc;

use avril_caps::Capabilities;
use avril_core::serde_v8::{JsBuffer, ToJsBuffer};
use avril_core::{Extension, Op, PublicModule};
use avril_kv::{Decision, Store};
use tokio::runtime::Handle;

pub use avril_kv::{MAX_BATCH_BYTES, MAX_BATCH_OPS};
pub use registry::Registry;

struct Shared {
    handle: Handle,
    caps: Arc<Capabilities>,
    registry: Registry,
}

// redb work runs off the isolate thread.
fn blocking<T, W>(shared: &Arc<Shared>, work: W) -> impl Future<Output = io::Result<T>> + use<T, W>
where
    T: Send + 'static,
    W: FnOnce(&Shared) -> io::Result<T> + Send + 'static,
{
    let task = {
        let shared = shared.clone();
        shared.handle.clone().spawn_blocking(move || work(&shared))
    };
    async move {
        task.await
            .unwrap_or_else(|_| Err(io::Error::other("kv background task stopped")))
    }
}

// An op on an open database: `function` splits its arguments into the
// database id and the work to run against the re-authorized store.
fn kv_op<A, R, W>(
    shared: &Arc<Shared>,
    name: &'static str,
    write: bool,
    function: impl Fn(A) -> (u64, W) + 'static,
) -> Op
where
    A: serde::de::DeserializeOwned,
    R: serde::Serialize + Send + 'static,
    W: FnOnce(&Store) -> io::Result<R> + Send + 'static,
{
    let shared = shared.clone();
    Op::asynchronous(name, move |args: A| {
        let (db, work) = function(args);
        blocking(&shared, move |shared| {
            shared.registry.with_entry(db, write, work)
        })
    })
}

#[derive(serde::Deserialize)]
struct ListArgs {
    db: u64,
    keyspace: String,
    start: Option<JsBuffer>,
    end: Option<JsBuffer>,
    prefix: Option<JsBuffer>,
    limit: usize,
    reverse: bool,
}

// Internal tagging buffers each op before picking a variant, which strips the
// typed-array identity that JsBuffer needs; serde_bytes accepts plain bytes.
#[derive(serde::Deserialize)]
#[serde(tag = "t", rename_all = "snake_case")]
enum AtomicOp {
    Check {
        ks: String,
        #[serde(with = "serde_bytes")]
        key: Vec<u8>,
        expected: Option<String>,
    },
    Set {
        ks: String,
        #[serde(with = "serde_bytes")]
        key: Vec<u8>,
        #[serde(with = "serde_bytes")]
        value: Vec<u8>,
    },
    Delete {
        ks: String,
        #[serde(with = "serde_bytes")]
        key: Vec<u8>,
    },
}

impl From<AtomicOp> for avril_kv::BatchOp {
    fn from(op: AtomicOp) -> Self {
        match op {
            AtomicOp::Check { ks, key, expected } => Self::Check { keyspace: ks, key, expected },
            AtomicOp::Set { ks, key, value } => Self::Set { keyspace: ks, key, value },
            AtomicOp::Delete { ks, key } => Self::Delete { keyspace: ks, key },
        }
    }
}

pub fn init(handle: Handle, caps: Arc<Capabilities>) -> Extension {
    let shared = Arc::new(Shared {
        handle,
        caps,
        registry: Registry::new(),
    });
    let open = shared.clone();
    let close = shared.clone();
    Extension {
        name: "ext:kv",
        dependencies: &["ext:file", "ext:url", "ext:web", "ext:encoding"],
        ops: vec![
            Op::asynchronous(
                "op_kv_open",
                move |(path, access, create): (String, String, bool)| {
                    blocking(&open, move |shared| {
                        shared.registry.open(&shared.caps, &path, &access, create)
                    })
                },
            ),
            Op::asynchronous("op_kv_close", move |id: u64| {
                blocking(&close, move |shared| shared.registry.close(id))
            }),
            kv_op(&shared, "op_kv_list_keyspaces", false, |db: u64| {
                (db, |store: &Store| store.read_tx(|tx| tx.list_keyspaces()))
            }),
            kv_op(
                &shared,
                "op_kv_get_many",
                false,
                |(db, keyspace, keys): (u64, String, Vec<JsBuffer>)| {
                    (db, move |store: &Store| {
                        store.read_tx(|tx| {
                            keys.iter()
                                .map(|key| {
                                    let entry = tx.get(&keyspace, key)?;
                                    Ok(entry.map(|(value, stamp)| (ToJsBuffer::from(value), stamp)))
                                })
                                .collect::<io::Result<Vec<_>>>()
                        })
                    })
                },
            ),
            kv_op(
                &shared,
                "op_kv_set",
                true,
                |(db, keyspace, key, value): (u64, String, JsBuffer, JsBuffer)| {
                    (db, move |store: &Store| {
                        store.write_tx(|tx| tx.put(&keyspace, &key, &value).map(Decision::Commit))
                    })
                },
            ),
            kv_op(
                &shared,
                "op_kv_delete",
                true,
                |(db, keyspace, key): (u64, String, JsBuffer)| {
                    (db, move |store: &Store| {
                        store.write_tx(|tx| {
                            tx.delete(&keyspace, &key)?;
                            Ok(Decision::Commit(()))
                        })
                    })
                },
            ),
            kv_op(
                &shared,
                "op_kv_clear",
                true,
                |(db, keyspace): (u64, String)| {
                    (db, move |store: &Store| {
                        store.write_tx(|tx| tx.clear(&keyspace).map(Decision::Commit))
                    })
                },
            ),
            kv_op(&shared, "op_kv_list", false, |args: ListArgs| {
                (args.db, move |store: &Store| {
                    let entries = store.read_tx(|tx| {
                        tx.list(
                            &args.keyspace,
                            args.start.as_deref(),
                            args.end.as_deref(),
                            args.prefix.as_deref(),
                            args.limit,
                            args.reverse,
                        )
                    })?;
                    Ok(entries
                        .into_iter()
                        .map(|(key, value, stamp)| {
                            (ToJsBuffer::from(key), ToJsBuffer::from(value), stamp)
                        })
                        .collect::<Vec<_>>())
                })
            }),
            kv_op(
                &shared,
                "op_kv_atomic",
                true,
                |(db, ops): (u64, Vec<AtomicOp>)| {
                    let ops: Vec<avril_kv::BatchOp> = ops.into_iter().map(Into::into).collect();
                    (db, move |store: &Store| store.atomic_commit(&ops))
                },
            ),
        ],
        modules: const { &[PUBLIC_MODULES[0].module] },
        entry_point: Some("ext:kv/main.js"),
    }
}

pub const PUBLIC_MODULES: &[PublicModule] =
    &[avril_core::public_module!("avril:kv", "ext:kv", "main")];

#[cfg(test)]
mod tests;