// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! A Timely Dataflow operator that mirrors a persisted stream.

use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::time::Duration;

use persist_types::Codec;
use timely::dataflow::operators::generic::operator;
use timely::dataflow::operators::{Concat, Map};
use timely::dataflow::{Scope, Stream};
use timely::progress::Antichain;
use timely::Data as TimelyData;

use crate::error::Error;
use crate::indexed::runtime::StreamReadHandle;
use crate::indexed::{ListenEvent, ListenFn};
use crate::operators::replay::Replay;

/// A Timely Dataflow operator that mirrors a persisted stream.
pub trait PersistedSource<G: Scope<Timestamp = u64>, K: TimelyData, V: TimelyData> {
    /// Emits a snapshot of the persisted stream taken as of this call and
    /// listens for any new data added to the persisted stream after that.
    fn persisted_source(
        &mut self,
        read: Result<StreamReadHandle<K, V>, Error>,
    ) -> Stream<G, (Result<(K, V), String>, u64, isize)>;
}

impl<G, K, V> PersistedSource<G, K, V> for G
where
    G: Scope<Timestamp = u64>,
    K: TimelyData + Codec + Send,
    V: TimelyData + Codec + Send,
{
    fn persisted_source(
        &mut self,
        read: Result<StreamReadHandle<K, V>, Error>,
    ) -> Stream<G, (Result<(K, V), String>, u64, isize)> {
        let (listen_tx, listen_rx) = mpsc::channel();
        let listen_fn = ListenFn(Box::new(move |e| {
            // TODO: If send fails, it means the operator is no longer running.
            // We should probably allow the listen to deregister itself.
            let _ = listen_tx.send(e);
        }));

        let snapshot = read.and_then(|read| read.listen(listen_fn));

        let snapshot_seal = snapshot
            .as_ref()
            .map_or(None, |snapshot| Some(snapshot.get_seal()));

        // TODO: Plumb the name of the stream down through the handles and use
        // it for the operator.

        // listen to new data
        let new = listen_source(self, snapshot_seal, listen_rx);

        // Replay the previously persisted data, if any.
        let previous = self.replay(snapshot);

        previous.concat(&new)
    }
}

/// Creates a source that listens on the given `listen_rx`.
fn listen_source<G, K, V>(
    scope: &G,
    // This uses an `Option<Antichain<_>>` and not an empty `Antichain` to signal that there is no
    // initial frontier because an empty frontier would signal that we are at "the end of time",
    // meaning that there will be no more events in the future.
    initial_frontier: Option<Antichain<u64>>,
    listen_rx: Receiver<ListenEvent<Vec<u8>, Vec<u8>>>,
) -> Stream<G, (Result<(K, V), String>, u64, isize)>
where
    G: Scope<Timestamp = u64>,
    K: TimelyData + Codec + Send,
    V: TimelyData + Codec + Send,
{
    let source_stream = operator::source(scope, "PersistedSource", |mut capability, info| {
        let worker_index = scope.index();
        let activator = scope.activator_for(&info.address[..]);

        // For Antichain<u64> there can only ever be one element, because
        // u64 has a total order. If we ever change this to an Antichain<T>
        // where we only have T: PartialOrder, this would change. If we have
        // that, though, I don't think we can downgrade this capability because
        // downgrading can only happen when `old_ts` <= `new_ts`, which is not
        // the case when an Antichain has multiple elements.
        if let Some(initial_frontier) = initial_frontier {
            for element in initial_frontier.iter() {
                capability.downgrade(element);
            }
        }

        let mut cap = Some(capability);
        move |output| {
            let mut done = false;
            if let Some(cap) = cap.as_mut() {
                let mut session = output.session(cap);
                match listen_rx.try_recv() {
                    Ok(e) => match e {
                        ListenEvent::Records(mut records) => {
                            // TODO: This currently works by only emitting the persisted data on worker
                            // 0 because that was the simplest thing to do initially. Instead, we should
                            // shard up the responsibility between all the workers.
                            if worker_index == 0 {
                                for record in records.drain(..) {
                                    let ((k, v), ts, diff) = record;
                                    let k = K::decode(&k);
                                    let v = V::decode(&v);
                                    session.give(((k, v), ts, diff));
                                }
                            }
                            activator.activate();
                        }
                        ListenEvent::Sealed(ts) => {
                            cap.downgrade(&ts);
                            activator.activate();
                        }
                    },
                    Err(TryRecvError::Empty) => {
                        // TODO: Hook the activator up to the callback instead of
                        // TryRecvError::Empty.
                        activator.activate_after(Duration::from_millis(100));
                    }
                    Err(TryRecvError::Disconnected) => {
                        done = true;
                    }
                }
            }
            if done {
                cap = None;
            }
        }
    });

    source_stream.map(|u| match u {
        ((Ok(k), Ok(v)), ts, diff) => (Ok((k, v)), ts, diff),
        ((Err(err), _), ts, diff) => (Err(err), ts, diff),
        ((_, Err(err)), ts, diff) => (Err(err), ts, diff),
    })
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use timely::dataflow::operators::capture::Extract;
    use timely::dataflow::operators::{Capture, OkErr, Probe};
    use timely::dataflow::ProbeHandle;
    use timely::Config;

    use crate::error::Error;
    use crate::mem::MemRegistry;
    use crate::operators::split_ok_err;

    use super::*;

    #[test]
    fn persisted_source() -> Result<(), Error> {
        let mut registry = MemRegistry::new();
        let p = registry.runtime_no_reentrance()?;

        let (oks, errs) = timely::execute_directly(move |worker| {
            let (oks, errs) = worker.dataflow(|scope| {
                let read = p
                    .create_or_load::<String, ()>("1")
                    .map(|(_write, read)| read);
                let (ok_stream, err_stream) = scope.persisted_source(read).ok_err(split_ok_err);
                (ok_stream.capture(), err_stream.capture())
            });

            let (write, _) = p.create_or_load("1").unwrap();
            for i in 1..=5 {
                write
                    .write(&[((i.to_string(), ()), i, 1)])
                    .recv()
                    .expect("write was successful");
            }
            write.seal(6).recv().expect("seal was successful");
            (oks, errs)
        });

        assert_eq!(
            errs.extract()
                .into_iter()
                .flat_map(|(_time, data)| data.into_iter().map(|(err, _ts, _diff)| err))
                .collect::<Vec<_>>(),
            Vec::<String>::new()
        );

        let mut actual = oks
            .extract()
            .into_iter()
            .flat_map(|(_, xs)| xs.into_iter().map(|((k, _), _, _)| k))
            .collect::<Vec<_>>();
        actual.sort();

        let expected = (1usize..=5usize).map(|x| x.to_string()).collect::<Vec<_>>();
        assert_eq!(actual, expected);

        Ok(())
    }

    /// Verifies that `persisted_source()` correctly downgrades to the upper frontier
    /// of already existing/sealed data when it is re-created based on an existing collection.
    ///
    /// There is no easy way to assert on a probe. Instead, we step the worker until the probe is
    /// no longer `less_than` the expected timestamp. This test would hang if the source were not
    /// working as expected.
    #[test]
    fn eager_capability_downgrade() -> Result<(), Error> {
        let mut registry = MemRegistry::new();

        let seal_ts = {
            let p = registry.runtime_no_reentrance()?;

            let (write, _) = p.create_or_load("1").unwrap();
            for i in 1..=5 {
                write
                    .write(&[((i.to_string(), ()), i, 1)])
                    .recv()
                    .expect("write was successful");
            }
            write.seal(6).recv().expect("seal was successful");
            6
        };

        let registry = Arc::new(Mutex::new(registry));

        let result = timely::execute_directly(move |worker| {
            let mut registry = registry.lock().expect("poisoned lock");
            let p = registry.runtime_no_reentrance().expect("missing registry");
            let mut probe = ProbeHandle::new();

            worker.dataflow(|scope| {
                let read = p
                    .create_or_load::<String, ()>("1")
                    .map(|(_write, read)| read);
                let (oks, _rrs) = scope.persisted_source(read).ok_err(split_ok_err);

                oks.probe_with(&mut probe);
            });

            while probe.less_than(&seal_ts) {
                worker.step();
            }

            true
        });

        assert!(result);

        Ok(())
    }

    #[test]
    #[ignore]
    fn multiple_workers() -> Result<(), Error> {
        let mut registry = MemRegistry::new();
        let p = registry.runtime_no_reentrance()?;

        // Write some data using 3 workers.
        timely::execute(Config::process(3), move |worker| {
            worker.dataflow::<u64, _, _>(|scope| {
                let (write, _) = p.create_or_load("1").unwrap();
                // Write one thing from each worker.
                write
                    .write(&[((format!("worker-{}", scope.index()), ()), 1, 1)])
                    .recv()
                    .expect("write was successful")
            });
        })?;

        // Execute a second dataflow with a different number of workers (2).
        // This is mainly testing that we only get one copy of the original
        // persisted data in the stream (as opposed to one per worker).
        let p = registry.runtime_no_reentrance()?;
        let (tx, rx) = mpsc::channel();
        let capture_tx = Arc::new(Mutex::new(tx));
        timely::execute(Config::process(2), move |worker| {
            worker.dataflow(|scope| {
                let (write, read) = p.create_or_load("1").unwrap();
                let (ok_stream, _) = scope.persisted_source(Ok(read)).ok_err(split_ok_err);

                // Write one thing from each worker again. This time at timestamp 2.
                write
                    .write(&[((format!("worker-{}", scope.index()), ()), 2, 1)])
                    .recv()
                    .expect("write was successful");

                // Now seal time 3 so we can probe the output.
                if scope.index() == 0 {
                    write.seal(3).recv().expect("seal was successful");
                }

                // Send the data to be captured by a channel so that we can replay
                // its contents outside of the dataflow and verify they are correct
                let tx = capture_tx.lock().expect("lock is not poisoned").clone();
                ok_stream.capture_into(tx)
            });
            // TODO: This hangs because the PersistedSource operator is still
            // hanging on to its capability. We need to wait until we've
            // captured through ts=3 and then unregister the listener/close the
            // persister to unblock things.
        })?;

        let mut actual = rx
            .extract()
            .into_iter()
            .flat_map(|(_, xs)| xs.into_iter())
            .collect::<Vec<_>>();
        actual.sort();

        let expected = vec![
            (("worker-0".to_owned(), ()), 1, 1),
            (("worker-0".to_owned(), ()), 2, 1),
            (("worker-1".to_owned(), ()), 1, 1),
            (("worker-1".to_owned(), ()), 2, 1),
            (("worker-2".to_owned(), ()), 1, 1),
        ];
        assert_eq!(expected, actual);

        Ok(())
    }

    #[test]
    fn error_stream() -> Result<(), Error> {
        let mut p = MemRegistry::new().runtime_no_reentrance()?;
        let read = p.create_or_load::<(), ()>("1").map(|(_write, read)| read);
        p.stop()?;

        let recv = timely::execute_directly(move |worker| {
            let recv = worker.dataflow(|scope| {
                let (_, err_stream) = scope.persisted_source(read).ok_err(split_ok_err);
                err_stream.capture()
            });
            recv
        });

        let actual = recv
            .extract()
            .into_iter()
            .flat_map(|(_, xs)| xs.into_iter())
            .collect::<Vec<_>>();

        let expected = vec![(
            "replaying persisted data: runtime shutdown".to_string(),
            0,
            1,
        )];
        assert_eq!(actual, expected);

        Ok(())
    }

    #[test]
    fn initial_error_handling() -> Result<(), Error> {
        let read: Result<StreamReadHandle<(), ()>, _> = Err(Error::String("ciao".to_owned()));

        let recv = timely::execute_directly(move |worker| {
            let recv = worker.dataflow(|scope| {
                let stream = scope.persisted_source(read);
                stream.capture()
            });

            recv
        });

        let actual = recv
            .extract()
            .into_iter()
            .flat_map(|(_, xs)| xs.into_iter())
            .collect::<Vec<_>>();

        let expected = vec![(Err("replaying persisted data: ciao".to_string()), 0, 1)];
        assert_eq!(actual, expected);

        Ok(())
    }

    // Previously, creating a persistent source would create differently shaped operator graphs,
    // depending on whether an internal call to persistence was successful or not. Timely dataflow
    // doesn't like that, which would manifest in hanging worker threads when trying to shut down.
    //
    // This test concurrently shuts down the persist runtime (which will make the internal
    // `listen()` call fail) and creates a bunch of timely workers. Some workers are expected to
    // call `listen` before the runtime is shut down and some after. We "help" the test a bit by
    // inserting `thread::sleep()` calls in strategic places, this way it is almost guaranteed to
    // fail without the fix.
    #[test]
    fn regression_8687_deterministic_operator_construction() -> Result<(), Error> {
        let p = MemRegistry::new().runtime_no_reentrance()?;
        let (write, _read) = p.create_or_load::<String, ()>("1").unwrap();

        // Write some data.
        let data = vec![(("ciao".into(), ()), 1, 1), (("bello".into(), ()), 1, 1)];
        write.write(&data).recv().expect("write failed");
        write.seal(1).recv().expect("seal failed");

        // Concurrently shut down the runtime and create multiple timely workers/threads from a
        // second thread that create a persistent source.
        let mut cloned_p = p.clone();
        let runtime_shutdown = std::thread::spawn(move || {
            std::thread::sleep(Duration::from_millis(1));
            cloned_p.stop().expect("shutdown failed");
        });

        let timely = std::thread::spawn(move || {
            let guards = timely::execute(Config::process(3), move |worker| {
                let mut probe = ProbeHandle::new();
                if worker.index() == 0 {
                    std::thread::sleep(Duration::from_millis(2));
                }
                worker.dataflow(|scope| {
                    let read = p
                        .create_or_load::<String, ()>("1")
                        .map(|(_write, read)| read);
                    let data = scope.persisted_source(read);
                    data.probe_with(&mut probe);
                });

                while probe.less_than(&1) {
                    worker.step();
                }
            })?;

            let result: Result<Vec<_>, _> = guards.join().into_iter().collect();

            result
        });

        // Assert that we can shut down cleanly. Without the fix, this test would not fail but
        // would hang indefinitely while trying to shut down the timely workers.
        runtime_shutdown
            .join()
            .expect("joining runtime thread failed");

        timely
            .join()
            .expect("joining timely launcher thread failed")
            .expect("timely workers failed");

        Ok(())
    }
}
