use bincode::{self, deserialize, serialize};
use log::*;
use serde::{Deserialize, Serialize};
use sled::{self, Db};
use std::ops::Deref;

use super::cryptograph::{Edge, Vertex};

#[derive(Serialize, Deserialize, PartialEq, Debug)]
struct PersistentVertex<V: AsRef<[u8]>> {
    vertex: Option<V>,
    edges: Vec<[u8; 32]>,
}

impl<V: AsRef<[u8]>> PersistentVertex<V> {
    fn borrowed(&self) -> PersistentVertex<&[u8]> {
        PersistentVertex {
            vertex: self.vertex.as_ref().map(AsRef::as_ref),
            edges: self.edges.clone(),
        }
    }

    fn cloned(&self) -> PersistentVertex<Vec<u8>> {
        PersistentVertex {
            vertex: self.vertex.as_ref().map(AsRef::as_ref).map(Vec::from),
            edges: self.edges.clone(),
        }
    }
}

#[derive(Clone)]
pub struct PersistentGraph {
    storage: Db,
}

#[derive(Debug, thiserror::Error)]
pub enum PersistError {
    #[error("Error serializing {0}")]
    SerializationError(#[from] bincode::Error),
    #[error("Error during persist {0}")]
    SledError(#[from] sled::Error),
}

fn merge_vertices<V: AsRef<[u8]>>(
    _key: &[u8],
    old: Option<PersistentVertex<V>>,
    new: PersistentVertex<V>,
) -> PersistentVertex<Vec<u8>> {
    let old = old.as_ref().map(PersistentVertex::borrowed);
    let new = new.borrowed();

    if let Some(old) = old {
        // Real merge
        let PersistentVertex {
            vertex: old_vertex,
            edges: mut old_edges,
        } = old;

        let PersistentVertex {
            vertex: new_vertex,
            mut edges,
        } = new;

        edges.append(&mut old_edges);

        let old_vertex = old_vertex.map(Vertex::from_bytes).map(Result::unwrap);
        let new_vertex = new_vertex.map(Vertex::from_bytes).map(Result::unwrap);

        let vertex = match (old_vertex, new_vertex) {
            (Some(a), Some(b)) => {
                if a.clock > b.clock {
                    Some(a)
                } else {
                    Some(b)
                }
            }
            (a, b) => a.or(b),
        }
        .as_ref()
        .map(Vertex::as_bytes);

        PersistentVertex { vertex, edges }
    } else {
        new.cloned()
    }
}

fn merge_encoded_vertices(key: &[u8], old: Option<&[u8]>, new: &[u8]) -> Vec<u8> {
    let old: Option<PersistentVertex<&[u8]>> = old.map(deserialize).map(Result::unwrap);
    let new = deserialize(new).unwrap();
    let merged = merge_vertices(key, old, new);
    serialize(&merged).unwrap()
}

impl PersistentGraph {
    pub fn temp() -> PersistentGraph {
        Self::open(sled::Config::new().temporary(true)).expect("temp graph always openable")
    }

    pub fn open(config: sled::Config) -> Result<PersistentGraph, PersistError> {
        let config = config.open().unwrap();
        config.set_merge_operator(|x: &[u8], y: Option<&[u8]>, z: &[u8]| {
            Some(merge_encoded_vertices(x, y, z))
        });
        Ok(PersistentGraph { storage: config })
    }

    pub fn persist_vertex(&self, v: &Vertex) -> Result<(), PersistError> {
        // TODO: check clock, only store most recent
        let persist = PersistentVertex {
            vertex: Some(v.as_bytes()),
            edges: vec![],
        };
        let encoded = serialize(&persist)?;
        let key = v.identifier();
        debug!("Persisting vertex {:?}", v);
        self.storage.merge(key.to_vec(), encoded)?;

        Ok(())
    }

    pub fn vertices(&self) -> impl Iterator<Item = (Vertex, impl Iterator<Item = Edge> + '_)> + '_ {
        // XXX some Result::ok calls should actually yield an error here.
        self.storage
            .iter()
            .values()
            .filter_map(Result::ok)
            .filter_map(move |x| {
                let v = deserialize::<PersistentVertex<Vec<u8>>>(&x).ok()?;
                let vertex = Vertex::from_bytes(v.vertex.as_deref()?).ok()?;
                let edges = self.find_persisted_edges(v).filter_map(Result::ok);
                Some((vertex, edges))
            })
    }

    pub fn persist_edge(&self, e: &Edge) -> Result<(), PersistError> {
        // XXX: make this a transaction when `sled' implements transactions:
        // https://github.com/spacejam/sled/issues/209

        let e_key = e.key();
        self.storage.insert(e_key, e.as_bytes())?;

        debug_assert!(e.vertex_keys().count() > 0);
        for v_key in e.vertex_keys() {
            let persist = PersistentVertex::<&[u8]> {
                vertex: None,
                edges: vec![e_key],
            };
            // XXX: don't unwrap this thing, wrong bytes can be fed apparently.
            debug!("Persisting vertex {:?} for edge", z85::encode(&v_key));
            self.storage.merge(v_key.to_vec(), serialize(&persist)?)?;
            self.storage.flush()?;
        }
        Ok(())
    }

    fn find_persistent_vertex<K: AsRef<[u8]>>(
        &self,
        key: K,
    ) -> Result<Option<PersistentVertex<Vec<u8>>>, PersistError> {
        Ok(self
            .storage
            .get(key.as_ref())?
            .map(|v| deserialize::<PersistentVertex<_>>(&v).expect("Database corruption FIXME")))
    }

    fn find_persisted_edges<'a, 's: 'a, V: AsRef<[u8]> + 'a>(
        &'s self,
        persisted: PersistentVertex<V>,
    ) -> impl Iterator<Item = Result<Edge, PersistError>> + 'a {
        persisted.edges.into_iter().map(move |edge_key| {
            let bytes = self
                .storage
                .get(&edge_key)?
                .expect("Database corruption 3 FIXME");
            Ok(Edge::from_bytes(bytes.deref()).expect("Database corruption 4 FIXME"))
        })
    }

    pub fn find_vertex<K: AsRef<[u8]>>(
        &self,
        key: K,
    ) -> Result<
        (
            Option<Vertex>,
            Option<impl Iterator<Item = Result<Edge, PersistError>> + '_>,
        ),
        PersistError,
    > {
        if let Some(persisted) = self.find_persistent_vertex(key)? {
            let vertex = persisted.vertex.as_ref().map(|persisted| {
                Vertex::from_bytes(&persisted).expect("Database corruption 2 FIXME")
            });
            trace!("Retrieving vertex with {} edges", persisted.edges.len());

            Ok((vertex, Some(self.find_persisted_edges(persisted))))
        } else {
            Ok((None, None))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::crypto::ec::*;
    use crate::graph::cryptograph::*;
    use ring::rand;

    #[test]
    fn test_persist_and_retrieve() {
        let config = sled::Config::new().temporary(true);
        let pg = PersistentGraph::open(config).unwrap();

        let rng = rand::SystemRandom::new();

        let alice_sk = PrivateKey::generate(&rng).unwrap();
        let alice_pub = alice_sk.compute_public_key().unwrap();

        let (alice, alice_signer) = VertexBuilder::new(&rng, &alice_pub)
            .unwrap()
            .build(&rng, &alice_sk)
            .unwrap();

        let (bob, _bob_signer) = VertexBuilder::new(&rng, &alice_pub)
            .unwrap()
            .build(&rng, &alice_sk)
            .unwrap();

        pg.persist_vertex(&alice).unwrap();
        pg.persist_vertex(&bob).unwrap();

        let edge = EdgeBuilder::default()
            .from(&alice)
            .to(&bob)
            .label("isFriend")
            .build(&rng, &alice_signer)
            .unwrap();

        pg.persist_edge(&edge).unwrap();

        let (retrieved, retrieved_edges) = pg.find_vertex(alice.identifier()).unwrap();
        let retrieved = retrieved.unwrap();
        assert!(retrieved.verify_signature());

        let retrieved_edges = retrieved_edges.unwrap().map(|e| {
            let e = e.unwrap();
            let edge = e.decrypt(&alice_signer).unwrap();
            assert_eq!(edge.predicate, "isFriend");
            edge
        });
        assert_eq!(retrieved_edges.count(), 1);

        let collected: Vec<_> = pg.vertices().collect();
        assert_eq!(collected.len(), 2);
    }

    #[test]
    fn test_update() {
        let config = sled::Config::new().temporary(true);
        let pg = PersistentGraph::open(config).unwrap();

        let rng = rand::SystemRandom::new();

        let alice_sk = PrivateKey::generate(&rng).unwrap();
        let alice_pub = alice_sk.compute_public_key().unwrap();

        let (alice, _alice_signer) = VertexBuilder::new(&rng, &alice_pub)
            .unwrap()
            .build(&rng, &alice_sk)
            .unwrap();

        assert_eq!(alice.clock, 0);

        pg.persist_vertex(&alice).unwrap();

        let (retrieved, retrieved_edges) = pg.find_vertex(alice.identifier()).unwrap();
        let retrieved = retrieved.unwrap();
        let retrieved_edges = retrieved_edges.unwrap();
        assert!(retrieved.verify_signature());

        assert_eq!(retrieved_edges.count(), 0);

        let (alice_new, _alice_signer) = VertexBuilder::update_vertex(&rng, &alice, &alice_sk)
            .unwrap()
            .build(&rng, &alice_sk)
            .unwrap();

        assert_eq!(alice.identifier(), alice_new.identifier());
        assert_eq!(retrieved.identifier(), alice_new.identifier());
        assert_eq!(alice_new.clock, 1);

        pg.persist_vertex(&alice_new).unwrap();

        let (retrieved, retrieved_edges) = pg.find_vertex(alice.identifier()).unwrap();
        let _retrieved_edges = retrieved_edges.unwrap();
        let retrieved = retrieved.unwrap();
        assert!(retrieved.verify_signature());

        assert_eq!(
            retrieved.clock, alice_new.clock,
            "Persisting newer vertex still returns old"
        );

        let collected: Vec<_> = pg.vertices().collect();
        assert_eq!(collected.len(), 1);
    }
}