pub trait KvStore {
    fn get(&self, key: &str) -> Option<Bytes>;

    /// Set the value associated with a key along with an optional expiration
    /// Duration.
    ///
    /// If a value is already associated with the key, it is removed.
    fn set(&self, key: String, value: Bytes, expire: Option<Duration>);

    /// Returns a `Receiver` for the requested channel.
    ///
    /// The returned `Receiver` is used to receive values broadcast by `PUBLISH`
    /// commands.
    fn subscribe(&self, key: String) -> broadcast::Receiver<Bytes>;

    /// Publish a message to the channel. Returns the number of subscribers
    /// listening on the channel.
    fn publish(&self, key: &str, value: Bytes) -> usize;
}

#[derive(Debug)]
pub(crate) struct Store {
    pub(crate) entries: HashMap<String, Entry>,

    pub(crate) pub_sub: HashMap<String, broadcast::Sender<Bytes>>,

    pub(crate) expirations: BTreeMap<(Instant, u64), String>,

    pub(crate) next_id: u64,

    pub(crate) shutdown: bool,
}

#[derive(Debug)]
pub(crate) struct Entry {
    pub(crate) id: u64,

    pub(crate) data: Bytes,

    pub(crate) expires_at: Option<Instant>,
}

impl Store {
    pub(crate) fn new() -> Store {
        Store {
            entries: HashMap::new(),
            pub_sub: HashMap::new(),
            expirations: BTreeMap::new(),
            next_id: 0,
            shutdown: false,
        }
    }

    pub(crate) fn next_expiration(&self) -> Option<Instant> {
        self.expirations.keys().next().map(|expire| expire.0)
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Db {
    shared: Arc<SharedDb>,
}

impl Db {
    pub(crate) fn new() -> Db {
        let shared = Arc::new(SharedDb::new());

        tokio::spawn(Db::purge_expired_tasks(shared.clone()));

        Db { shared }
    }

    async fn purge_expired_tasks(shared: Arc<SharedDb>) {
        while !shared.is_shutdown() {
            if let Some(when) = shared.purge_expired_keys() {
                tokio::select! {
                    _ = time::sleep_until(when) => {}
                    _ = shared.background_task.notified() => {}
                }
            } else {
                shared.background_task.notified().await;
            }
        }

        info!("Purge background task shut down")
    }

    fn shutdown_purge_task(&self) {
        let mut store = self.shared.store.lock().unwrap();
        store.shutdown = true;

        drop(store);
        self.shared.background_task.notify_one();
    }
}

#[derive(Debug)]
struct SharedDb {
    store: Mutex<Store>,

    background_task: Notify,
}

impl SharedDb {
    fn new() -> Self {
        SharedDb {
            store: Mutex::new(Store::new()),
            background_task: Notify::new(),
        }
    }

    fn purge_expired_keys(&self) -> Option<Instant> {
        let mut store = self.store.lock().unwrap();

        if store.shutdown {
            return None;
        }

        let store = &mut *store;

        let now = Instant::now();
        while let Some((&(when, id), key)) = store.expirations.iter().next() {
            if when > now {
                return Some(when);
            }

            store.entries.remove(key);
            store.expirations.remove(&(when, id));
        }

        None
    }

    fn is_shutdown(&self) -> bool {
        self.store.lock().unwrap().shutdown
    }
}

impl KvStore for Db {

    fn get(&self, key: &str) -> Option<Bytes> {
        let store = self.shared.store.lock().unwrap();
        store.entries.get(key).map(|entry| entry.data.clone())
    }

    fn set(&self, key: String, value: Bytes, expire: Option<Duration>) {
        let mut store = self.shared.store.lock().unwrap();

        let id = store.next_id;
        store.next_id += 1;

        let mut notify = false;

        let expires_at = expire.map(|duration| {
            let when = Instant::now() + duration;

            notify = store
                .next_expiration()
                .map(|expiration| expiration > when)
                .unwrap_or(true);

            store.expirations.insert((when, id), key.clone());
            when
        });

        let prev = store.entries.insert(
            key,
            Entry {
                id,
                data: value,
                expires_at,
            },
        );

        if let Some(prev) = prev {
            if let Some(when) = prev.expires_at {
                // clear expiration
                store.expirations.remove(&(when, prev.id));
            }
        }

        drop(store);

        if notify {
            self.shared.background_task.notify_one();
        }
    }

    fn subscribe(&self, key: String) -> broadcast::Receiver<Bytes> {
        use std::collections::hash_map::Entry;

        let mut store = self.shared.store.lock().unwrap();

        match store.pub_sub.entry(key) {
            Entry::Occupied(e) => e.get().subscribe(),
            Entry::Vacant(e) => {
                let (tx, rx) = broadcast::channel(1024);
                e.insert(tx);
                rx
            }
        }
    }

    fn publish(&self, key: &str, value: Bytes) -> usize {
        debug!("publish: (key={}, len(value)={})", key, value.len());

        let state = self.shared.store.lock().unwrap();

        state
            .pub_sub
            .get(key)
            .map(|tx| tx.send(value).unwrap_or(0))
            .unwrap_or(0)
    }
}

fn get(&self, key: &str) -> Option<Bytes> {
  // Acquire the lock, get the entry and clone the value.
  //
  // Because data is stored using `Bytes`, a clone here is a shallow
  // clone. Data is not copied.
  let store = self.shared.store.lock().unwrap();
  store.entries.get(key).map(|entry| entry.data.clone())
}

fn set(&self, key: String, value: Bytes, expire: Option<Duration>) {
  let mut store = self.shared.store.lock().unwrap();

  let id = store.next_id;
  store.next_id += 1;

  let mut notify = false;

  let expires_at = expire.map(|duration| {
    let when = Instant::now() + duration;

    notify = store
    .next_expiration()
    .map(|expiration| expiration > when)
    .unwrap_or(true);

    store.expirations.insert((when, id), key.clone());
    when
  });

  let prev = store.entries.insert(
    key,
    Entry {
      id,
      data: value,
      expires_at,
    },
  );

  if let Some(prev) = prev {
    if let Some(when) = prev.expires_at {
      store.expirations.remove(&(when, prev.id));
    }
  }

  drop(store);

  if notify {
    self.shared.background_task.notify_one();
  }
}

fn subscribe(&self, key: String) -> broadcast::Receiver<Bytes> {
  use std::collections::hash_map::Entry;

  let mut store = self.shared.store.lock().unwrap();

  match store.pub_sub.entry(key) {
    Entry::Occupied(e) => e.get().subscribe(),
    Entry::Vacant(e) => {
      let (tx, rx) = broadcast::channel(1024);
      e.insert(tx);
      rx
    }
  }
}

fn publish(&self, key: &str, value: Bytes) -> usize {
  let state = self.shared.store.lock().unwrap();

  state
  .pub_sub
  .get(key)
  .map(|tx| tx.send(value).unwrap_or(0))
  .unwrap_or(0)
}

/// A wrapper around a `Db` instance. This exists to allow orderly cleanup
/// of the `Db` by signalling the background purge task to shut down when
/// this struct is dropped.
#[derive(Debug)]
pub(crate) struct DbDropGuard {
    /// The `Db` instance that will be shut down when this `DbHolder` struct
    /// is dropped.
    db: Db,
}

impl DbDropGuard {
    /// Create a new `DbHolder`, wrapping a `Db` instance. When this is dropped
    /// the `Db`'s purge task will be shut down.
    pub(crate) fn new() -> DbDropGuard {
        DbDropGuard { db: Db::new() }
    }

    /// Get the shared database. Internally, this is an `Arc`,
    /// so a clone only increments the ref count.
    pub(crate) fn db(&self) -> Db {
        self.db.clone()
    }
}

impl Drop for DbDropGuard {
    fn drop(&mut self) {
        // Signal the 'Db' instance to shut down the task that purges expired keys
        self.db.shutdown_purge_task();
    }
}