use zookeeper::{ZooKeeper, WatchedEvent, WatchedEventType, ZkError, ZkResult, ZooKeeperExt};
use std::sync::{Arc, RwLock};
use serde::{Serialize};
use serde::de::DeserializeOwned;
pub use treediff::{value::Key, diff, tools::ChangeType};
use std::time::{Instant, Duration};
use std::{thread, fmt};
use std::sync::RwLockReadGuard;
use serde_json::Value;
use crossbeam_channel::Receiver;
use std::fmt::Debug;
use rand::Rng;


const MAX_TIMING_DELTA: i64 = 30000; // ms
const LOCK_POLL_INTERVAL: u64 = 5; // ms
const LOCK_POLL_TIMEOUT: u64 = 1000; // ms

#[derive(Debug)]
pub enum ZkStructError {
    /// Timed out when trying to lock the struct for writing
    LockAcquireTimeout,
    StaleRead,
    /// The expected version of the object does not match the remote version.
    StaleWrite,

    ZkError(ZkError),
    Poisoned
}
impl From<ZkError> for ZkStructError {
    fn from(error: ZkError) -> ZkStructError {
        ZkStructError::ZkError(error)
    }
}

#[derive(Debug)]
struct InternalState {
    id: String,

    /// Path of the ZkStruct Dir
    zk_path: String,

    /// Known epoch of the local object. Compared to the remote epoch
    epoch: i32,

    /// Last time the inner object was sync
    timings: chrono::DateTime<chrono::Utc>,
    
    emit_updates: bool,

    chan_rx: crossbeam_channel::Receiver<Change<Key, serde_json::Value>>,
    chan_tx: crossbeam_channel::Sender<Change<Key, serde_json::Value>>,
}

#[derive(Clone)]
pub struct ZkState<T: Serialize + DeserializeOwned + Send + Sync> {
    /// ZooKeeper client
    zk: Option<Arc<ZooKeeper>>,
    id: String,

    inner: Arc<RwLock<T>>,
    state: Arc<RwLock<InternalState>>
}
impl<T: Serialize + DeserializeOwned + Send + Sync + 'static> ZkState<T> {
    pub fn new(zk: Arc<ZooKeeper>, zk_path: String, initial_state: T) -> anyhow::Result<Self> {
        zk.ensure_path(zk_path.as_str())?;

        let id = uuid::Uuid::new_v4().to_string();

        let (chan_tx, chan_rx) = crossbeam_channel::unbounded();
        let r = Self {
            id: id.clone(),
            zk: Some(zk),
            inner: Arc::new(RwLock::new(initial_state)),
            state: Arc::new(RwLock::new(InternalState {
                id,
                zk_path,
                epoch: 0,
                timings: chrono::Utc::now(),
                emit_updates: true,
                chan_rx,
                chan_tx
            }))
        };
        r.initialize()?;
        Ok(r)
    }

    /// Create an Empty, do nothing, ZkState object.
    pub fn empty(initial_state: T) -> Self {
        let (chan_tx, chan_rx) = crossbeam_channel::unbounded();
        Self {
            id: "empty".to_string(),
            zk: None,
            inner: Arc::new(RwLock::new(initial_state)),
            state: Arc::new(RwLock::new(InternalState {
                id: "empty".to_string(),
                zk_path: "empty".to_string(),
                epoch: 0,
                timings: chrono::Utc::now(),
                emit_updates: false,
                chan_rx,
                chan_tx
            }))
        }
    }

    /// Create a new instance of ZkState, but does not attempt to create an empty object if none
    /// already exists. Will wait for data to show up.
    ///
    /// So, if you know that a ZkState instance will eventually appear at the `/foo/bar` znode, this
    /// will block waiting for that to show up.
    ///
    /// TODO describe this better
    pub fn expect(zk: Arc<ZooKeeper>, zk_path: String, timeout: Option<Duration>) -> anyhow::Result<Self> {
        let id = uuid::Uuid::new_v4().to_string();

        // block, waiting for our initial state to show up in zookeeper
        let (l_tx, l_rx) = crossbeam_channel::unbounded();
        let raw_data = zk.get_data_w(format!("{}/payload", &zk_path).as_str(), move |_| {
            let _ = l_tx.send(());
        });

        let data;
        if let Ok(inner) = raw_data {
            data = inner.0;
        } else {
            match timeout {
                None => l_rx.recv()?,
                Some(d) => l_rx.recv_timeout(d)?
            }
            data = zk.get_data(format!("{}/payload", &zk_path).as_str(), false)?.0;
        }

        let (chan_tx, chan_rx) = crossbeam_channel::unbounded();
        let r = Self {
            id: id.clone(),
            zk: Some(zk),
            inner: Arc::new(RwLock::new(serde_json::from_slice(data.as_slice()).unwrap())),
            state: Arc::new(RwLock::new(InternalState {
                id,
                zk_path,
                epoch: 0,
                timings: chrono::Utc::now(),
                emit_updates: true,
                chan_rx,
                chan_tx
            }))
        };
        r.initialize()?;
        Ok(r)
    }

    fn initialize(&self) -> ZkResult<()> {
        if self.zk.is_none() {
            return Ok(());
        }

        let path = format!("{}/payload", &self.state.read().unwrap().zk_path);
        // if the path doesn't exist, let's make it and populate it
        if self.zk.clone().unwrap().exists(path.as_str(), false).unwrap().is_none() {
            log::debug!(target: &*format!("zkstate::{}", &self.id), "{} does not exist, creating", &path);

            let data = self.inner.read().unwrap();
            let inner = serde_json::to_vec(&*data).unwrap();
            self.zk.clone().unwrap().create(path.as_str(), inner, zookeeper::Acl::open_unsafe().clone(), zookeeper::CreateMode::Persistent)?;
        } else {
            log::trace!(target: &*format!("zkstate::{}", &self.id), "{} exists, continuing initialization", &path);
        }

        state_change(self.zk.clone().unwrap(), self.inner.clone(), self.state.clone());

        // Create a thread that will ensure consistency in the background.
        // TODO make the interval configurable
        let zk = self.zk.as_ref().unwrap().clone();
        let state = self.state.clone();
        let inner = self.inner.clone();

        thread::spawn(move || {
            let zk = zk;
            let inner = inner;
            let state = state;
            let mut rng = rand::thread_rng();
            loop {
                thread::sleep(Duration::from_secs_f32(rng.gen_range(3.0..7.0))); // TODO make this configurable

                let handle = state.read().unwrap();
                let path = format!("{}/payload", handle.zk_path);

                if let Some(meta) = zk.exists(path.as_str(), false).unwrap() {
                    let epoch = handle.epoch;
                    let id = handle.id.clone();
                    drop(handle);

                    if epoch != meta.version {
                        // TODO this doesn't seem to actually do anything.
                        // No change notifications are emitted
                        log::warn!(target: &*format!("zkstate::{}", &id), "the remote epoch has drifted. local: {}. remote: {}", epoch, meta.version);
                        state_change(zk.clone(), inner.clone(), state.clone());
                    }

                    state.write().unwrap().timings = chrono::Utc::now();
                }
            }
        });

        Ok(())
    }

    /// Method to be invoked to handle state change notifications
    pub fn update_handler<M: Fn(Change<Key, serde_json::Value>) + Send + 'static>(&self, closure: M) -> Result<(), ZkStructError> {
        let chan_handle = self.state.read().unwrap().chan_rx.clone();
        thread::spawn(move || {
            let rx = chan_handle;
            loop {
                let message = rx.recv().unwrap();
                closure(message);
            }
        });
        Ok(())
    }

    /// Return a reference to the Crossbeam Receiver to get Change notifications
    pub fn get_update_channel(&self) -> Receiver<Change<Key, Value>> {
        self.state.read().unwrap().chan_rx.clone()
    }

    /// Update the shared object using a closure.
    ///
    /// The closure is passed a reference to the contents of the ZkState. Once the closure returns,
    /// the shared state in Zookeeper is committed and the write locks released.
    pub fn update<M: FnOnce(&mut T)>(&self, closure: M) -> Result<(), ZkStructError> {
        if self.zk.is_none() {
            return Ok(());
        }

        let path = format!("{}/payload", &self.state.read().unwrap().zk_path);

        // acquire write lock for the internal object
        let mut inner = self.inner.write().unwrap();
        let mut state = self.state.write().unwrap();

        // get write lock from zookeeper to prevent anyone from modifying this object while we're
        // committing it
        let latch = zookeeper::recipes::leader::LeaderLatch::new(self.zk.as_ref().unwrap().clone(), self.id.clone(),  format!("{}/write_lock", &state.zk_path));
        latch.start()?;

        let mut total_time = 0;
        loop {
            if latch.has_leadership() { break; }

            log::trace!("I do not have write_lock, waiting {}ms before trying again", LOCK_POLL_INTERVAL);

            thread::sleep(Duration::from_millis(LOCK_POLL_INTERVAL));
            if total_time > LOCK_POLL_TIMEOUT {
                return Err(ZkStructError::LockAcquireTimeout)
            } else {
                total_time += LOCK_POLL_INTERVAL;
            }
        }

        // internal state, pre change
        let a = serde_json::to_value(&*inner).unwrap();

        // at this point, we should have an exclusive lock on the object so we execute the closure
        closure(&mut inner);

        // internal state, post change
        let b = serde_json::to_value(&*inner).unwrap();

        let raw_data = serde_json::to_vec(&*inner).unwrap();
        let update_op = self.zk.as_ref().unwrap().set_data(path.as_str(), raw_data, Some(state.epoch));
        match update_op {
            Ok(_) => {}
            Err(err) => {
                if err == ZkError::BadVersion {
                    return Err(ZkStructError::StaleWrite)
                }
                return Err(ZkStructError::ZkError(err))
            }
        }

        // adjust the internal epoch to match the remote
        let result_stats = update_op.unwrap();
        state.epoch = result_stats.version;

        emit_updates(&a, &b, &state);

        drop(inner); // drop the write lock on the inner object
        drop(state); // drop the write lock on the internal state object

        if let Err(inner) = latch.stop() {
            return Err(ZkStructError::ZkError(inner));
        }

        Ok(())
    }

    /// Returns a Result<RwLockReadGuard<T>>
    pub fn read(&self) -> Result<RwLockReadGuard<'_, T>, ZkStructError> {
        let state = self.state.read().unwrap();
        let delta = (chrono::Utc::now() - state.timings).num_milliseconds();
        if delta > MAX_TIMING_DELTA && self.zk.is_some() {
            log::error!(target: &*format!("zkstate::{}", &self.id), "attempted to read stale data. data is {}ms old, limit is {}ms", &delta, MAX_TIMING_DELTA);
            return Err(ZkStructError::StaleRead)
        }
        log::trace!(target: &*format!("zkstate::{}", &self.id), "reading internal data. data is {}ms old, limit is {}ms. epoch is {}", &delta, MAX_TIMING_DELTA, state.epoch);

        match self.inner.read() {
            Ok(inner) => Ok(inner),
            Err(_) => Err(ZkStructError::Poisoned)
        }
    }

    /// Consistent Read.
    ///
    /// Preforms a check to make sure the local version is the same as the remote version before
    /// returning. If there is a mismatch, will preform a sync and return the latest object
    pub fn c_read(&self) { unimplemented!() }

    /// Dirty Read.
    ///
    /// Returns the local data, not failing if the data is too old
    // TODO condense this code, with c_read, and read to remove some code reuse
    pub fn d_read(&self) -> Result<RwLockReadGuard<'_, T>, ZkStructError> {
        let state = self.state.read().unwrap();
        let delta = (chrono::Utc::now() - state.timings).num_milliseconds();
        if delta > MAX_TIMING_DELTA {
            log::warn!(target: &*format!("zkstate::{}", &self.id), "attempted to read stale data. data is {}ms old, limit is {}ms", &delta, MAX_TIMING_DELTA);
        } else {
            log::trace!(target: &*format!("zkstate::{}", &self.id), "dirty reading internal data. data is {}ms old, limit is {}ms. epoch is {}", &delta, MAX_TIMING_DELTA, state.epoch);
        }

        match self.inner.read() {
            Ok(inner) => Ok(inner),
            Err(_) => Err(ZkStructError::Poisoned)
        }
    }

    pub fn metadata(&self) -> (usize, i32) {
        return (self.state.read().unwrap().chan_rx.len(), 0)
    }

    /// Return the ID of this ZkState instance
    pub fn get_id(&self) -> &String {
        &self.id
    }
}
impl<T: Serialize + DeserializeOwned + Send + Sync + Debug + 'static> fmt::Debug for ZkState<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ZkState{{ id: {}, inner: {:?}, state: {:?} }}", self.id, &*self.inner.read().unwrap(), &*self.state.read().unwrap())
    }
}

// TODO this seems to fire twice
fn handle_zk_watcher<T: Serialize + DeserializeOwned + Send + Sync + 'static>(ev: WatchedEvent, zk: Arc<ZooKeeper>, inner: Arc<RwLock<T>>, state: Arc<RwLock<InternalState>>) {
    if let WatchedEventType::NodeDataChanged = ev.event_type {
        log::trace!(target: &*format!("zkstate::{}", &state.read().unwrap().id), "handling state change. got WatchedEventType::NodeDataChanged");
        state_change(zk, inner, state)
    }
}

#[derive(PartialEq, Debug)]
pub enum Change<K, V> {
    /// The Value was removed
    Removed(Vec<K>, V),
    /// The Value was added
    Added(Vec<K>, V),
    /// No change was performed to the Value
    Unchanged(),
    /// The first Value was modified and became the second Value
    Modified(Vec<K>, V, V),
}

/// Pull the full state from ZooKeeper, compare it to the current inner, Enqueue Changes, and then
/// Update the inner field with the state
fn state_change<T: Serialize + DeserializeOwned + Send + Sync + 'static>(zk: Arc<ZooKeeper>, inner: Arc<RwLock<T>>, state: Arc<RwLock<InternalState>>) {
    let start = Instant::now();

    let path = format!("{}/payload", &state.read().unwrap().zk_path);
    let movable = (zk.clone(), inner.clone(), state.clone());

    let raw_obj = zk.get_data_w(path.as_str(), move |ev| {
        let movable = movable.clone();
        handle_zk_watcher(ev, movable.0, movable.1, movable.2);
    }).unwrap();

    // explicitly hold on to the handle while we compare the diff. it might take a bit for big objects
    // and we do not want to allow for something else to update the object while we're in the process
    // of updating.
    let mut a_handle = inner.write().unwrap();
    let mut state = state.write().unwrap();

    let id = state.id.clone();
    log::trace!(target: &*format!("zkstate::{}", &id), "reading {}, version: {}. last modified: {}", &path, &raw_obj.1.version, &raw_obj.1.mtime);

    let b: serde_json::Value = serde_json::from_slice(&*raw_obj.0).unwrap();

    // only do a delta if we want to emit updates
    if state.emit_updates {
        let a: serde_json::Value = serde_json::to_value(&*a_handle).unwrap();
        emit_updates(&a, &b, &state);
    }

    *a_handle = serde_json::from_value(b).unwrap();
    state.epoch = raw_obj.1.version;
    state.timings = chrono::Utc::now();

    drop(a_handle); // drop the write handle for the internal object
    drop(state);    // drop the write handle for the state object

    let dur = start.elapsed().as_millis();
    if dur > 225 {
        log::warn!(target: &*format!("zkstate::{}", &id), "took {}ms to handle state change", dur);
    } else {
        log::trace!(target: &*format!("zkstate::{}", &id), "took {}ms to handle state change", dur);
    }
}

fn emit_updates(a: &serde_json::Value, b: &serde_json::Value, state: &InternalState) {
    let mut delta = treediff::tools::Recorder::default();
    diff(a, b, &mut delta);

    let mut ops = (0, 0, 0, 0);
    for change in delta.calls {
        let op = match change {
            ChangeType::Added(k, v) => {
                ops.0 += 1;
                Change::Added(k.clone(), v.clone())
            },
            ChangeType::Removed(k, v) => {
                ops.1 += 1;
                Change::Removed(k.clone(), v.clone())
            },
            ChangeType::Modified(k, a, v) => {
                ops.2 += 1;
                Change::Modified(k.clone(), a.clone(), v.clone())
            },
            ChangeType::Unchanged(_, _) => {
                ops.3 += 1;
                Change::Unchanged()
            },
        };
        if op != Change::Unchanged() {
            log::trace!(target: &*format!("zkstate::{}", &state.id), "enqueuing operation {:?}", &op);
            let _insert = state.chan_tx.send(op);
        }
    }
    log::debug!(target: &*format!("zkstate::{}", &state.id), "resulting operations: {} added, {} removed, {} modified, {} noop", ops.0, ops.1, ops.2, ops.3);

}