mod error;
mod lease;
mod lock;

use std::sync::Arc;

pub use crate::{error::*, lease::*, lock::*};
use etcd_client::{
    Client, Compare, CompareOp, Event, PutOptions, ResponseHeader, Txn, TxnOp, WatchOptions,
    WatchStream,
};
use futures::prelude::*;
use log::*;
use rand::random;

pub async fn try_race(client: &mut Client, key: &str, ttl: i64) -> CrateResult<bool> {
    let value = random::<u128>().to_le_bytes();
    let lease_id = client.lease_grant(ttl, None).await?.id();

    let locked = client
        .txn(
            Txn::new()
                .when(vec![Compare::create_revision(
                    key.as_bytes(),
                    CompareOp::Equal,
                    0,
                )])
                .and_then(vec![TxnOp::put(
                    key.as_bytes(),
                    &value[..],
                    Some(PutOptions::new().with_lease(lease_id)),
                )]),
        )
        .await?
        .succeeded();

    Ok(locked)
}

fn get_lock_txn(name: &str, lock_value: &[u8], lease_id: i64) -> Txn {
    Txn::new()
        .when([Compare::create_revision(name, CompareOp::Equal, 0)])
        .and_then([TxnOp::put(
            name,
            lock_value,
            Some(PutOptions::new().with_lease(lease_id)),
        )])
}

pub(crate) fn get_release_txn(name: &str, lock_value: &[u8], lease_id: i64) -> Txn {
    Txn::new()
        .when([
            Compare::value(name, CompareOp::Equal, lock_value),
            Compare::lease(name, CompareOp::Equal, lease_id),
        ])
        .and_then([TxnOp::delete(name, None)])
}

async fn watch_until_lock_available(
    etcd_client: &mut Client,
    watch_stream: &mut Option<WatchStream>,
    name: &str,
    last_locked_revision: Option<i64>,
) -> CrateResult<()> {
    loop {
        let mut etcd_client = etcd_client.clone();
        let mut inner_watch_stream = match watch_stream.take() {
            Some(watch_stream) => watch_stream,
            None => {
                let mut watch_options = WatchOptions::new();
                if let Some(last_locked_revision) = last_locked_revision {
                    watch_options = watch_options.with_start_revision(last_locked_revision + 1);
                }

                let (_, watch_stream) = etcd_client.watch(name, Some(watch_options)).await?;

                watch_stream
            }
        };

        let watch_response = inner_watch_stream
            .try_next()
            .await?
            .expect("Unexpected end of etcd watch stream");

        *watch_stream = Some(inner_watch_stream);

        let is_available = matches!(watch_response.events().last().and_then(Event::kv), Some(kv) if kv.create_revision() == 0);

        trace!(
            "Got update from watch stream for key: {} (avilable: {})",
            name,
            is_available
        );

        if is_available {
            return Ok(());
        }
    }
}

impl Lease {
    pub(crate) async fn stop_keep_alive(self) -> Result<(), etcd_client::Error> {
        if let Ok(LeaseInner {
            stop_tx,
            keep_alive_handle,
        }) = Arc::try_unwrap(self.inner)
        {
            let _ = stop_tx.send(());
            keep_alive_handle.await.unwrap()?;
        }

        Ok(())
    }

    /// Stop the keep-alive task, propagating any errors that it may have encountered,
    /// and ask the etcd cluster to revoke this lease. Any etcd keys associated with this lease,
    /// including keys used as locks by this library's [`LockGuard`], will be removed.
    ///
    /// See also: [etcd docs]
    ///
    /// [`LockGuard`]: ./struct.LockGuard.html
    /// [etcd docs]: https://etcd.io/docs/v3.3.12/dev-guide/interacting_v3/#revoke-leases
    pub async fn release(self) -> Result<(), etcd_client::Error> {
        let mut client = self.client.clone();
        let lease_id = self.lease_id;

        self.stop_keep_alive().await?;
        client.lease_revoke(lease_id).await?;

        debug!("Released lease: {:#x}", lease_id);

        Ok(())
    }

    async fn _lock<'l, 'n>(
        &'l mut self,
        name: &'n str,
        wait: bool,
    ) -> CrateResult<LockGuard<'l, 'n>> {
        let lock_value = random::<u128>().to_le_bytes();

        let mut watch_stream = None;
        let locked = loop {
            let txn_response = self
                .client
                .txn(get_lock_txn(name, &lock_value, self.lease_id))
                .await?;

            if txn_response.succeeded() {
                break true;
            } else if !wait {
                break false;
            }

            watch_until_lock_available(
                &mut self.client,
                &mut watch_stream,
                name,
                txn_response.header().map(ResponseHeader::revision),
            )
            .await?;
        };

        if !locked {
            return Err(Error::Taken);
        } else {
            debug!("Acquired lock: {} (lease: {:#x})", name, self.lease_id);
        }

        Ok(LockGuard {
            client: self.client.clone(),
            lock_value,
            name,
            lease: &*self,
        })
    }

    pub async fn lock<'l, 'n>(&'l mut self, name: &'n str) -> CrateResult<LockGuard<'l, 'n>> {
        self._lock(name, true).await
    }

    /// Anologous to [`lock`], but if the lock is taken, return [`Error::Taken`] instead of waiting for the lock
    /// to become available.
    ///
    /// [`lock`]: ./struct.Lease.html#method.lock
    /// [`Error::taken`]: ./enum.Error.html#variant.Taken
    pub async fn try_lock<'l, 'n>(&'l mut self, name: &'n str) -> CrateResult<LockGuard<'l, 'n>> {
        self._lock(name, false).await
    }

    /// Run and await future-returning function `f`, but only in the context of an acquired lock.
    ///
    /// Attempt to acquire lock with `name`, waiting for it to become available if needed, and run `f` after
    /// it is successfuly acquired. After `f` is finished and the future it returned is awaited, the lock is released
    /// automatically.
    ///
    /// # Panic handling
    /// This function does not attempt to catch panics from `f`, and does not guarantee that the lock is released in such
    /// cases. If the program exits abruptly, the lock will be not be available to other clients until this lease's TTL runs out.
    ///
    /// # Example
    /// ```no_run
    /// # async fn test(mut lease: tokio_etcd_lock::Lease) -> Result<(), tokio_etcd_lock::Error> {
    ///
    /// lease.with_lock("lock-key", || async {
    ///     println!("Holding mutually exclusive access to the lock!");
    /// }).await?;
    ///
    /// // at this point, the lock has been released and is available for other clients
    /// # Ok(())
    /// # }
    /// ```
    pub async fn with_lock<F, Fut, T>(&mut self, name: &str, f: F) -> CrateResult<T>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = T>,
    {
        let guard = self.lock(name).await?;
        let result = f().await;
        guard.release().await?;

        Ok(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::{random, Rng};
    use std::{
        ops::AddAssign,
        sync::{
            atomic::{AtomicUsize, Ordering},
            Arc,
        },
        time::Duration,
    };
    use tokio::{
        sync::{Barrier, Mutex},
        time::delay_for,
    };

    async fn create_client() -> CrateResult<Client> {
        Ok(Client::connect(vec!["http://localhost:2379"], None).await?)
    }

    fn create_lock_name() -> String {
        format!("{:x}", random::<u128>())
    }

    #[tokio::test]
    async fn test_mutex() -> anyhow::Result<()> {
        const WORKERS: usize = 10;

        type Result<T> = anyhow::Result<T, anyhow::Error>;

        let lock_name = create_lock_name();

        let client = create_client().await?;
        let locked_count = Arc::new(Mutex::new(0_usize));

        let mut rng = rand::thread_rng();

        stream::iter(0..WORKERS)
            .map(Result::Ok)
            .try_for_each_concurrent(None, |_| {
                let client = client.clone();
                let lock_name = lock_name.clone();
                let locked_count = locked_count.clone();

                async move {
                    let mut lease = acquire_lease(&client, 10).await?;

                    lease
                        .with_lock(&lock_name, || async move {
                            {
                                let mut locked_count = locked_count.try_lock()?;
                                delay_for(Duration::from_millis(rng.gen_range(0, 2_000))).await;
                                locked_count.add_assign(1);
                            }

                            Result::Ok(())
                        })
                        .await??;

                    Result::Ok(())
                }
            })
            .await?;

        assert_eq!(*locked_count.try_lock().unwrap(), WORKERS);

        Ok(())
    }

    #[tokio::test]
    async fn test_lease_expire() -> CrateResult<()> {
        let lock_name = create_lock_name();

        let client = create_client().await?;

        let lease1 = acquire_lease(&client, 10).await?;
        let mut lease2 = acquire_lease(&client, 10).await?;

        let barrier1 = Arc::new(Barrier::new(2));
        let barrier2 = Arc::new(Barrier::new(2));

        future::try_join(
            {
                let mut client = client.clone();
                let lock_name = lock_name.clone();
                let barrier1 = barrier1.clone();
                let barrier2 = barrier2.clone();
                async move {
                    client
                        .txn(get_lock_txn(&lock_name, b"", lease1.lease_id))
                        .await?;
                    barrier1.wait().await;
                    barrier2.wait().await;
                    lease1.stop_keep_alive().await?;
                    Ok(())
                }
            },
            {
                async move {
                    barrier1.wait().await;

                    // make sure the lock is still taken at this point
                    assert!(matches!(
                        lease2.try_lock(&lock_name).await.unwrap_err(),
                        Error::Taken
                    ));

                    // allow first worker to stop the lease keep-alive
                    barrier2.wait().await;

                    lease2
                        .with_lock(&lock_name, || async {
                            info!("Inside lock.");
                        })
                        .await
                }
            },
        )
        .await?;

        Ok(())
    }

    #[tokio::test]
    async fn test_lock_lost() -> CrateResult<()> {
        let lock_name = create_lock_name();

        let mut client = create_client().await?;
        let mut lease = acquire_lease(&client, 10).await?;

        let lease_id = lease.lease_id;

        let res = lease
            .with_lock(&lock_name, || async {
                client.lease_revoke(lease_id).await.unwrap();
            })
            .await;

        assert!(matches!(res, Err(Error::Lost)));

        Ok(())
    }

    #[tokio::test]
    async fn test_try_race() -> CrateResult<()> {
        const TTL: u64 = 5;

        let lock_name = create_lock_name();

        let mut client = create_client().await?;
        let locked_count = Arc::new(AtomicUsize::new(0));

        let mut rng = rand::thread_rng();

        let wait_for_ttl = delay_for(Duration::from_secs(TTL + 1));

        stream::iter(0..20)
            .map(|_| {
                let mut client = client.clone();
                let lock_name = lock_name.clone();
                let locked_count = locked_count.clone();

                async move {
                    delay_for(Duration::from_secs(rng.gen_range(0, TTL - 1))).await;

                    if try_race(&mut client, &lock_name, TTL as i64).await? {
                        locked_count.fetch_add(1, Ordering::SeqCst);
                    }

                    CrateResult::Ok(())
                }
            })
            .buffer_unordered(20)
            .try_collect::<Vec<_>>()
            .await?;

        assert_eq!(locked_count.load(Ordering::SeqCst), 1);

        wait_for_ttl.await;
        assert!(try_race(&mut client, &lock_name, 10).await?);

        Ok(())
    }
}