use std::{
    cell::{RefCell, RefMut},
    collections::{hash_map::Entry, HashMap},
    rc::Rc,
};

use ::log::{debug, trace, log_enabled};
use audi::{Listener, ListenerSet, ListenerWithState};
use futures::{
    channel::mpsc::{Receiver, Sender, channel},
    future::{self, join_all},
    Future, StreamExt,
};
use litl::impl_debug_as_litl;
use mofo::{Mofo};
use ridl::{hashing::HashOf, signing::SignerSecret};
use serde_derive::{Deserialize, Serialize};

mod causal_set;
mod log;
mod object;
mod set;
mod telepathic;

use crate::log::{LogID, LogState, LogWriteAccess};
pub use object::{Diff, ObjectID, ObjectState, WriteAccess};
use object::{ObjectError, StateInfo};
pub use set::SetItem;
use set::{SetDiff, SetID, SetState, SetWriteAccess};
use telepathic::{ApplyDiffErrorFor, Telepathic, TelepathicDiff};

pub type ApplyDiffError = ApplyDiffErrorFor<ObjectError>;

struct NodeInner {
    objects: HashMap<ObjectID, Rc<RefCell<ObjectState>>>,
    local_listeners: HashMap<ObjectID, ListenerSet<Diff>>,
    remotes: ListenerSet<UpdateMessage, RemoteState>,
    background: Mofo
}

impl NodeInner {
    fn load_object(&mut self, id: ObjectID) -> impl Future<Output = Rc<RefCell<ObjectState>>> {
        let object = self.objects.entry(id).or_insert_with(|| match id {
            ObjectID::Log(id) => Rc::new(RefCell::new(ObjectState::Log(LogState::new(id)))),
            ObjectID::Set(id) => Rc::new(RefCell::new(ObjectState::Set(SetState::new(id)))),
        });

        future::ready(Rc::clone(object))
    }

    fn get_object(&self, id: &ObjectID) -> Option<Rc<RefCell<ObjectState>>> {
        self.objects.get(id).cloned()
    }
}

#[derive(Clone)]
pub struct Node(Rc<RefCell<NodeInner>>);


impl Node {
    pub fn new(background: Mofo) -> Self {
        Self(Rc::new(RefCell::new(NodeInner {
            objects: HashMap::new(),
            local_listeners: HashMap::new(),
            remotes: ListenerSet::new(),
            background,
        })))
    }

    pub fn create_log(&self) -> WriteAccess {
        let write_access = LogWriteAccess(SignerSecret::new_random());
        let id = LogID(write_access.0.pub_id());

        match (*self.0).borrow_mut().objects.entry(ObjectID::Log(id)) {
            Entry::Occupied(_) => unreachable!("Log for given ID already exists"),
            Entry::Vacant(entry) => {
                entry.insert(Rc::new(RefCell::new(ObjectState::Log(LogState::new(id)))))
            }
        };

        WriteAccess::Log(write_access)
    }

    pub fn create_set(&self) -> WriteAccess {
        let write_access = SetWriteAccess(SignerSecret::new_random());
        let id = SetID(write_access.0.pub_id());

        match (*self.0).borrow_mut().objects.entry(ObjectID::Set(id)) {
            Entry::Occupied(_) => unreachable!("Set for given ID already exists"),
            Entry::Vacant(entry) => {
                entry.insert(Rc::new(RefCell::new(ObjectState::Set(SetState::new(id)))))
            }
        };

        WriteAccess::Set(write_access)
    }

    pub fn diff_for_log_append(
        &self,
        write_access: &WriteAccess,
        append: &[u8],
    ) -> Diff {
        if append.is_empty() {
            panic!("Empty append")
        }
        match write_access {
            WriteAccess::Log(log_access) => {

                    let self_ref = (*self.0).borrow_mut();

                    let log = self_ref
                        .get_object(&ObjectID::Log(log_access.log_id()))
                        .expect("Expected log to be created or loaded before appending");
                    let log = RefMut::map((*log).borrow_mut(), |object| match object {
                        ObjectState::Log(log) => log,
                        _ => unreachable!("Expected log to be log"),
                    });

                    Diff::Log(log.diff_for_new_append(append, log_access))

            }
            _ => panic!("Tried to append to log with set write access"),
        }
    }
    pub async fn append_to_log(
        &self,
        write_access: &WriteAccess,
        append: &[u8],
    ) {
        let diff = self.diff_for_log_append(write_access, append);
        self.apply_new_diff(diff).await;
    }

    pub async fn insert_into_set<I: IntoIterator<Item = HashOf<SetItem>>>(
        &self,
        write_access: &WriteAccess,
        item_data: Vec<u8>,
        item_prev: I,
    ) -> Result<HashOf<SetItem>, ApplyDiffError> {
        match write_access {
            WriteAccess::Set(set_access) => {
                let set_item = SetItem {
                    data: item_data,
                    prev: item_prev.into_iter().collect(),
                };

                let item_id = HashOf::hash(&set_item);

                let diff = SetDiff {
                    id: set_access.id(),
                    new_items: vec![set_access.0.sign(set_item)],
                };

                self.apply_new_diff(Diff::Set(diff)).await;

                Ok(item_id)
            }
            _ => panic!("Tried to insert into set with log write access"),
        }
    }

    pub async fn apply_new_diff(&self, diff: Diff) {
        let object_id = diff.id();

        let listeners = {
            let self_ref = (*self.0).borrow_mut();
            let object = self_ref
                .get_object(&diff.id())
                .expect("Expected object to be created or loaded before applying diff");
            (*object)
                .borrow_mut()
                .try_apply_diff(diff.clone())
                .expect("New diffs should always be valid");

            self_ref.local_listeners.get(&diff.id()).cloned()
        };

        if let Some(listeners) = listeners {
            listeners.broadcast(diff).await;
        }

        self.sync_object_with_remotes(object_id).await;
    }

    pub async fn sync_object_with_remotes(&self, id: ObjectID) {
        let (object, remotes) = {
            let self_ref = (*self.0).borrow_mut();

            let object = self_ref
                .get_object(&id)
                .expect("Expected object to be created or loaded before syncing");

            let remotes = self_ref.remotes.clone();

            (object, remotes)
        };

        remotes
            .broadcast_filter_map(|remote| {
                let optimistic_state_info = remote.optimistic_state_info.get(&id).and_then(|s| s.as_ref());

                if optimistic_state_info.is_some() || remote.mode == RemoteMode::SyncEverything {
                    (*object)
                        .borrow()
                        .diff_since(optimistic_state_info)
                        .map(UpdateMessage::NewContent)
                } else {
                    None
                }
            })
            .await
    }

    pub fn get_object(&self, id: &ObjectID) -> Option<Rc<RefCell<ObjectState>>> {
        (*self.0).borrow().get_object(id)
    }

    pub async fn load_object(&self, id: ObjectID) -> Rc<RefCell<ObjectState>> {
        let loaded = (*self.0).borrow_mut().load_object(id);
        loaded.await
    }

    pub async fn add_local_listener(&self, id: ObjectID, listener: Listener<Diff>) {
        let (listener_id, listeners, object) = {
            let listener_id = listener.id().to_owned();
            let mut self_ref = (*self.0).borrow_mut();
            let listeners = self_ref
                .local_listeners
                .entry(id)
                .or_insert_with(ListenerSet::new)
                .clone();
            listeners.add(listener);
            let object = self_ref.load_object(id);

            (listener_id, listeners, object)
        };

        {
            let object = object.await;
            let initial_diff = object.borrow().state_as_initial_diff();
            let is_empty = object.borrow().state_info().is_none();
            if let Some(initial_diff) = initial_diff {
                listeners.send_or_remove(&listener_id, initial_diff).await;
            }

            if is_empty {
                debug!("Listened to object is empty, syncing with remotes");
                let remotes = (*self.0).borrow().remotes.clone();
                remotes.broadcast_filter_map(|remote|
                    if remote.mode == RemoteMode::SyncEverything {
                        Some(UpdateMessage::StateInfo(StateInfoUpdate{id: object.borrow().id(), state_info: None}))
                    } else {
                        None
                    }
                ).await;
            }
        }
    }

    pub async fn add_remote(&self, remote: Remote) {
        (*self.0).borrow_mut().remotes.add(RemoteState {
            id: remote.id.clone(),
            mode: remote.mode,
            optimistic_state_info: HashMap::new(),
            outgoing: remote.outgoing.clone(),
        });

        let self_rc = self.clone();
        let remotes = (*self.0).borrow().remotes.clone();
        let remote_id = remote.id.clone();

        let receiving = remote.incoming.for_each(move |msg| {
            let self_rc = self_rc.clone();
            let remotes = remotes.clone();
            let remote_id = remote_id.clone();

            async move {
                if log_enabled!(::log::Level::Trace) {
                    let mut trunc_msg = format!("{msg:?}");
                    trunc_msg.truncate(500);
                    trace!("Receiving from {remote_id}: {trunc_msg}");
                }
                match msg {
                    UpdateMessage::NewContent(diff) => {
                        self_rc.handle_diff(diff, &remote_id, remotes).await;
                    }
                    UpdateMessage::StateInfo(StateInfoUpdate{id, state_info: new_state_info}) => {
                        self_rc.handle_new_state_info(id, new_state_info, &remote_id, remotes)
                            .await;
                    }
                }
            }
        });

        if remote.mode == RemoteMode::SyncEverything {
            let remotes = (*self.0).borrow().remotes.clone();
            let remote_id = remote.id.clone();
            #[allow(clippy::needless_collect)]
            let all_state_infoes = (*self.0)
                .borrow()
                .objects
                .iter()
                .map(|(id, obj)| (*id, (**obj).borrow().state_info()))
                .collect::<Vec<_>>();

            join_all(all_state_infoes.into_iter().map(|(id, state_info)| {
                remotes.send_or_remove(&remote_id, UpdateMessage::StateInfo(StateInfoUpdate{id, state_info}))
            }))
            .await;

            debug!("Sent all state_infoes to {remote_id}");
        }

        (*self.0).borrow().background.add_background_task(Box::pin(receiving), format!("receiving_from_{}", remote.id));
    }

    async fn handle_diff(
        &self,
        diff: Diff,
        remote_id: &str,
        remotes: ListenerSet<UpdateMessage, RemoteState>,
    ) {
        let (apply_result, id, current_state_info) = {
            let object_promise = (*self.0).borrow_mut().load_object(diff.id());
            let object_rc = object_promise.await;
            let mut object = (*object_rc).borrow_mut();

            (
                object.try_apply_diff(diff.clone()),
                object.id(),
                object.state_info(),
            )
        };
        match apply_result {
            Ok(new_state_info) => {
                if let Some(remote) = remotes.state_of(remote_id) {
                    (*remote)
                        .borrow_mut()
                        .optimistic_state_info
                        .insert(id, Some(new_state_info));
                }

                let listeners = (*self.0).borrow_mut().local_listeners.get(&id).cloned();

                if let Some(listeners) = listeners {
                    listeners.broadcast(diff).await;
                }

                remotes
                    .broadcast_filter_map(|remote| {
                        let object = self.get_object(&id).expect("Should have object now");
                        remote
                            .optimistic_state_info
                            .get(&id)
                            .and_then(|optimistic_state_info| {
                                (*object).borrow().diff_since(optimistic_state_info.as_ref())
                            })
                            .map(UpdateMessage::NewContent)
                    })
                    .await;
            }
            Err(ApplyDiffError::InvalidKnownStateAssumption) => {
                // remote must be mistaken about our state, let them know
                remotes
                    .send_or_remove(
                        remote_id,
                        UpdateMessage::StateInfo(StateInfoUpdate{id, state_info: current_state_info}),
                    )
                    .await;
            }
            Err(err) => Err(err).unwrap(),
        }
    }

    async fn handle_new_state_info(
        &self,
        id: ObjectID,
        new_state_info: Option<StateInfo>,
        remote_id: &str,
        remotes: ListenerSet<UpdateMessage, RemoteState>,
    ) {
        let update_to_respond_with = if let Some(remote) = remotes.state_of(remote_id) {
            (*remote)
                .borrow_mut()
                .optimistic_state_info
                .insert(id, new_state_info.clone());

            if let Some(object) = (*self.0).borrow_mut().objects.get(&id).cloned() {
                if let Some(update) = (*object).borrow().diff_since(new_state_info.as_ref()) {
                    Some(UpdateMessage::NewContent(update))
                } else {
                    let own_state_info = (*object).borrow().state_info();

                    if own_state_info != new_state_info {
                        // TODO: we can still get stuck in an endless loop here if we can't reconcile
                        // compare if we made progress since the last remembered optimistic state_info, only then send again
                        Some(UpdateMessage::StateInfo(StateInfoUpdate{id, state_info:own_state_info}))
                    } else {
                        None
                    }
                }
            } else {
                Some(UpdateMessage::StateInfo(StateInfoUpdate{id, state_info: None}))
            }
        } else {
            None
        };
        if let Some(update) = update_to_respond_with {
            let remotes = remotes.clone();
            remotes.send_or_remove(remote_id, update).await;
        }
    }
}

#[derive(Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum UpdateMessage {
    NewContent(Diff),
    StateInfo(StateInfoUpdate),
}

#[derive(Clone, Serialize, Deserialize)]
pub struct StateInfoUpdate{
    id: ObjectID,
    state_info: Option<StateInfo>
}

impl_debug_as_litl!(UpdateMessage);

pub struct Remote {
    pub id: String,
    pub mode: RemoteMode,
    pub incoming: Receiver<UpdateMessage>,
    pub outgoing: Sender<UpdateMessage>,
}

impl Remote {
    pub fn new_connected_test_pair(node1_name: &str, node2_name: &str) -> (Remote, Remote) {
        let (to_node1, node1_rx) = channel(100);
        let (to_node2, node2_rx) = channel(100);

        (
            Remote {
                id: node1_name.to_string(),
                mode: RemoteMode::ProvideWhatRemoteWants,
                incoming: node2_rx,
                outgoing: to_node1,
            },
            Remote {
                id: node2_name.to_string(),
                mode: RemoteMode::SyncEverything,
                incoming: node1_rx,
                outgoing: to_node2,
            },
        )
    }
}

#[derive(Copy, Clone, PartialEq, Eq)]
pub enum RemoteMode {
    SyncEverything,
    ProvideWhatRemoteWants,
}

struct RemoteState {
    id: String,
    mode: RemoteMode,
    optimistic_state_info: HashMap<ObjectID, Option<StateInfo>>,
    outgoing: Sender<UpdateMessage>,
}

impl ListenerWithState<UpdateMessage> for RemoteState {
    fn id(&self) -> &str {
        &self.id
    }

    fn sender(&self) -> &Sender<UpdateMessage> {
        &self.outgoing
    }
}