use crate::error::RealtimeError;
use crate::indexer::Indexer;
use crate::metrics::CollabRealtimeMetrics;
use anyhow::anyhow;
use app_error::AppError;
use arc_swap::ArcSwap;
use collab::core::collab::DataSource;
use collab::core::origin::CollabOrigin;
use collab::entity::EncodedCollab;
use collab::lock::RwLock;
use collab::preclude::Collab;
use collab_entity::CollabType;
use collab_rt_entity::user::RealtimeUser;
use collab_rt_entity::{
  AckCode, AwarenessSync, BroadcastSync, CollabAck, MessageByObjectId, MsgId,
};
use collab_rt_entity::{ClientCollabMessage, CollabMessage};
use collab_rt_protocol::{Message, MessageReader, RTProtocolError, SyncMessage};
use collab_stream::client::CollabRedisStream;
use collab_stream::collab_update_sink::{AwarenessUpdateSink, CollabUpdateSink};

use collab_stream::error::StreamError;
use collab_stream::model::{AwarenessStreamUpdate, CollabStreamUpdate, MessageId, UpdateFlags};
use dashmap::DashMap;
use database::collab::{CollabStorage, GetCollabOrigin};
use database_entity::dto::{
  AFCollabEmbeddings, CollabParams, InsertSnapshotParams, QueryCollabParams,
};
use futures::{pin_mut, Sink, Stream};
use futures_util::{SinkExt, StreamExt};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use tokio::time::MissedTickBehavior;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, trace};
use yrs::updates::decoder::{Decode, DecoderV1};
use yrs::updates::encoder::{Encode, Encoder, EncoderV1};
use yrs::{ReadTxn, StateVector, Update};

/// A group used to manage a single [Collab] object
pub struct CollabGroup {
  state: Arc<CollabGroupState>,
}

/// Inner state of [CollabGroup] that's private and hidden behind Arc, so that it can be moved into
/// tasks.
struct CollabGroupState {
  workspace_id: String,
  object_id: String,
  collab_type: CollabType,
  /// A list of subscribers to this group. Each subscriber will receive updates from the
  /// broadcast.
  subscribers: DashMap<RealtimeUser, Subscription>,
  persister: CollabPersister,
  metrics: Arc<CollabRealtimeMetrics>,
  /// Cancellation token triggered when current collab group is about to be stopped.
  /// This will also shut down all subsequent [Subscription]s.
  shutdown: CancellationToken,
  last_activity: ArcSwap<Instant>,
  seq_no: AtomicU32,
  /// The most recent state vector from a redis update.
  state_vector: RwLock<StateVector>,
}

impl Drop for CollabGroup {
  fn drop(&mut self) {
    // we're going to use state shutdown to cancel subsequent tasks
    self.state.shutdown.cancel();
  }
}

impl CollabGroup {
  #[allow(clippy::too_many_arguments)]
  pub fn new<S>(
    uid: i64,
    workspace_id: String,
    object_id: String,
    collab_type: CollabType,
    metrics: Arc<CollabRealtimeMetrics>,
    storage: Arc<S>,
    collab_redis_stream: Arc<CollabRedisStream>,
    persistence_interval: Duration,
    prune_grace_period: Duration,
    indexer: Option<Arc<dyn Indexer>>,
    state_vector: StateVector,
  ) -> Result<Self, StreamError>
  where
    S: CollabStorage,
  {
    let is_new_collab = state_vector.is_empty();
    let persister = CollabPersister::new(
      uid,
      workspace_id.clone(),
      object_id.clone(),
      collab_type.clone(),
      storage,
      collab_redis_stream,
      indexer,
      metrics.clone(),
      prune_grace_period,
    );

    let state = Arc::new(CollabGroupState {
      workspace_id,
      object_id,
      collab_type,
      subscribers: DashMap::new(),
      metrics,
      shutdown: CancellationToken::new(),
      persister,
      last_activity: ArcSwap::new(Instant::now().into()),
      seq_no: AtomicU32::new(0),
      state_vector: state_vector.into(),
    });

    /*
     NOTE: we don't want to pass `Weak<CollabGroupState>` to tasks and terminate them when they
     cannot be upgraded since we want to be sure that ie. when collab group is to be removed,
     that we're going to call for a final save of the document state.

     For that we use `CancellationToken` instead, which is racing against internal loops of child
     tasks and triggered when this `CollabGroup` is dropped.
    */

    // setup task used to receive collab updates from Redis
    {
      let state = state.clone();
      tokio::spawn(async move {
        if let Err(err) = Self::inbound_task(state).await {
          tracing::warn!("failed to receive collab update: {}", err);
        }
      });
    }

    // setup task used to receive awareness updates from Redis
    {
      let state = state.clone();
      tokio::spawn(async move {
        if let Err(err) = Self::inbound_awareness_task(state).await {
          tracing::warn!("failed to receive awareness update: {}", err);
        }
      });
    }

    // setup periodic snapshot
    {
      tokio::spawn(Self::snapshot_task(
        state.clone(),
        persistence_interval,
        is_new_collab,
      ));
    }

    Ok(Self { state })
  }

  #[inline]
  pub fn workspace_id(&self) -> &str {
    &self.state.workspace_id
  }

  #[inline]
  #[allow(dead_code)]
  pub fn object_id(&self) -> &str {
    &self.state.object_id
  }

  pub fn is_cancelled(&self) -> bool {
    self.state.shutdown.is_cancelled()
  }

  /// Task used to receive collab updates from Redis.
  async fn inbound_task(state: Arc<CollabGroupState>) -> Result<(), RealtimeError> {
    let updates = state.persister.collab_redis_stream.live_collab_updates(
      &state.workspace_id,
      &state.object_id,
      None,
    );
    pin_mut!(updates);
    loop {
      tokio::select! {
        _ = state.shutdown.cancelled() => {
          break;
        }
        res = updates.next() => {
          match res {
            Some(Ok((_message_id, update))) => {
              Self::handle_inbound_update(&state, update).await;
            },
            Some(Err(err)) => {
              tracing::warn!("failed to handle incoming update for collab `{}`: {}", state.object_id, err);
              break;
            },
            None => {
              break;
            }
          }
        }
      }
    }
    Ok(())
  }

  async fn handle_inbound_update(state: &CollabGroupState, update: CollabStreamUpdate) {
    // update state vector based on incoming message
    let mut sv = state.state_vector.write().await;
    sv.merge(update.state_vector);
    drop(sv);

    let seq_num = state.seq_no.fetch_add(1, Ordering::SeqCst) + 1;
    tracing::trace!(
      "broadcasting collab update from {} ({} bytes) - seq_num: {}",
      update.sender,
      update.data.len(),
      seq_num
    );
    let payload = Message::Sync(SyncMessage::Update(update.data)).encode_v1();
    let message = BroadcastSync::new(update.sender, state.object_id.clone(), payload, seq_num);
    for mut e in state.subscribers.iter_mut() {
      let subscription = e.value_mut();
      if message.origin == subscription.collab_origin {
        continue; // don't send update to its sender
      }

      if let Err(err) = subscription.sink.send(message.clone().into()).await {
        tracing::debug!(
          "failed to send collab `{}` update to `{}`: {}",
          state.object_id,
          subscription.collab_origin,
          err
        );
      }

      state.last_activity.store(Arc::new(Instant::now()));
    }
  }

  /// Task used to receive awareness updates from Redis.
  async fn inbound_awareness_task(state: Arc<CollabGroupState>) -> Result<(), RealtimeError> {
    let updates = state.persister.collab_redis_stream.awareness_updates(
      &state.workspace_id,
      &state.object_id,
      None,
    );
    pin_mut!(updates);
    loop {
      tokio::select! {
        _ = state.shutdown.cancelled() => {
          break;
        }
        res = updates.next() => {
          match res {
            Some(Ok(awareness_update)) => {
              Self::handle_inbound_awareness(&state, awareness_update).await;
            },
            Some(Err(err)) => {
              tracing::warn!("failed to handle incoming update for collab `{}`: {}", state.object_id, err);
              break;
            },
            None => {
              break;
            }
          }
        }
      }
    }
    Ok(())
  }

  async fn handle_inbound_awareness(state: &CollabGroupState, update: AwarenessStreamUpdate) {
    tracing::trace!(
      "broadcasting awareness update from {} ({} bytes)",
      update.sender,
      update.data.len()
    );
    let sender = update.sender;
    let message = AwarenessSync::new(
      state.object_id.clone(),
      Message::Awareness(update.data).encode_v1(),
      CollabOrigin::Empty,
    );
    for mut e in state.subscribers.iter_mut() {
      let subscription = e.value_mut();
      if sender == subscription.collab_origin {
        continue; // don't send update to its sender
      }

      if let Err(err) = subscription.sink.send(message.clone().into()).await {
        tracing::debug!(
          "failed to send awareness `{}` update to `{}`: {}",
          state.object_id,
          subscription.collab_origin,
          err
        );
      }

      state.last_activity.store(Arc::new(Instant::now()));
    }
  }

  async fn snapshot_task(state: Arc<CollabGroupState>, interval: Duration, is_new_collab: bool) {
    if is_new_collab {
      tracing::trace!("persisting new collab for {}", state.object_id);
      if let Err(err) = state.persister.save().await {
        tracing::warn!(
          "failed to persist new document `{}`: {}",
          state.object_id,
          err
        );
      }
    }

    let mut snapshot_tick = tokio::time::interval(interval);
    // if saving took longer than snapshot_tick, just skip it over and try in the next round
    snapshot_tick.set_missed_tick_behavior(MissedTickBehavior::Skip);

    loop {
      tokio::select! {
        _ = snapshot_tick.tick() => {
          if let Err(err) = state.persister.save().await {
            tracing::warn!("failed to persist collab `{}/{}`: {}", state.workspace_id, state.object_id, err);
          }
        },
        _ = state.shutdown.cancelled() => {
          if let Err(err) = state.persister.save().await {
            tracing::warn!("failed to persist collab on shutdown `{}/{}`: {}", state.workspace_id, state.object_id, err);
          }
          break;
        }
      }
    }
  }

  pub async fn calculate_missing_update(
    &self,
    state_vector: StateVector,
  ) -> Result<Vec<u8>, RealtimeError> {
    {
      // first check if we need to send any updates
      let collab_sv = self.state.state_vector.read().await;
      if *collab_sv <= state_vector {
        return Ok(vec![]);
      }
    }

    let encoded_collab = self.encode_collab().await?;
    let collab = Collab::new_with_source(
      CollabOrigin::Server,
      self.object_id(),
      DataSource::DocStateV1(encoded_collab.doc_state.into()),
      vec![],
      false,
    )?;
    let update = collab.transact().encode_state_as_update_v1(&state_vector);
    Ok(update)
  }

  pub async fn encode_collab(&self) -> Result<EncodedCollab, RealtimeError> {
    let snapshot = self.state.persister.load_compact().await?;
    let encode_collab = snapshot.collab.encode_collab_v1(|collab| {
      self
        .state
        .collab_type
        .validate_require_data(collab)
        .map_err(|err| RealtimeError::Internal(err.into()))
    })?;
    Ok(encode_collab)
  }

  pub fn contains_user(&self, user: &RealtimeUser) -> bool {
    self.state.subscribers.contains_key(user)
  }

  pub fn remove_user(&self, user: &RealtimeUser) {
    if self.state.subscribers.remove(user).is_some() {
      trace!(
        "{} remove subscriber from group: {}",
        self.state.object_id,
        user
      );
    }
  }

  pub fn user_count(&self) -> usize {
    self.state.subscribers.len()
  }

  pub fn modified_at(&self) -> Instant {
    *self.state.last_activity.load_full()
  }

  /// Subscribes a new connection to the broadcast group for collaborative activities.
  ///
  pub fn subscribe<Sink, Stream>(
    &self,
    user: &RealtimeUser,
    subscriber_origin: CollabOrigin,
    sink: Sink,
    stream: Stream,
  ) where
    Sink: SubscriptionSink + Clone + 'static,
    Stream: SubscriptionStream + 'static,
  {
    // create new subscription for new subscriber
    let subscriber_shutdown = self.state.shutdown.child_token();

    tokio::spawn(Self::receive_from_client_task(
      self.state.clone(),
      sink.clone(),
      stream,
      subscriber_origin.clone(),
    ));

    let sub = Subscription::new(sink, subscriber_origin, subscriber_shutdown);
    if self
      .state
      .subscribers
      .insert((*user).clone(), sub)
      .is_some()
    {
      tracing::warn!("{}: remove old subscriber: {}", &self.state.object_id, user);
    }

    if cfg!(debug_assertions) {
      trace!(
        "{}: add new subscriber, current group member: {}",
        &self.state.object_id,
        self.user_count(),
      );
    }

    trace!(
      "[realtime]:{} new subscriber:{}, connect at:{}, connected members: {}",
      self.state.object_id,
      user.user_device(),
      user.connect_at,
      self.state.subscribers.len(),
    );
  }

  async fn receive_from_client_task<Sink, Stream>(
    state: Arc<CollabGroupState>,
    mut sink: Sink,
    mut stream: Stream,
    origin: CollabOrigin,
  ) where
    Sink: SubscriptionSink + 'static,
    Stream: SubscriptionStream + 'static,
  {
    loop {
      tokio::select! {
        _ = state.shutdown.cancelled() => {
          break;
        }
        msg = stream.next() => {
          match msg {
            None => break,
            Some(msg) => if let Err(err) =  Self::handle_messages(&state, &mut sink, msg).await {
              tracing::warn!(
                "collab `{}` failed to handle message from `{}`: {}",
                state.object_id,
                origin,
                err
              );

            }
          }
        }
      }
    }
  }

  async fn handle_messages<Sink>(
    state: &CollabGroupState,
    sink: &mut Sink,
    msg: MessageByObjectId,
  ) -> Result<(), RealtimeError>
  where
    Sink: SubscriptionSink + 'static,
  {
    for (message_object_id, messages) in msg.0 {
      if state.object_id != message_object_id {
        error!(
          "Expect object id:{} but got:{}",
          state.object_id, message_object_id
        );
        continue;
      }
      for message in messages {
        match Self::handle_client_message(state, message).await {
          Ok(response) => {
            trace!("[realtime]: sending response: {}", response);
            match sink.send(response.into()).await {
              Ok(()) => {},
              Err(err) => {
                trace!("[realtime]: send failed: {}", err);
                break;
              },
            }
          },
          Err(err) => {
            error!(
              "Error handling collab message for object_id: {}: {}",
              message_object_id, err
            );
            break;
          },
        }
      }
    }
    Ok(())
  }

  /// Handle the message sent from the client
  async fn handle_client_message(
    state: &CollabGroupState,
    collab_msg: ClientCollabMessage,
  ) -> Result<CollabAck, RealtimeError> {
    let msg_id = collab_msg.msg_id();
    let message_origin = collab_msg.origin().clone();

    // If the payload is empty, we don't need to apply any updates .
    // Currently, only the ping message should has an empty payload.
    if collab_msg.payload().is_empty() {
      if !matches!(collab_msg, ClientCollabMessage::ClientCollabStateCheck(_)) {
        error!("receive unexpected empty payload message:{}", collab_msg);
      }
      return Ok(CollabAck::new(
        message_origin,
        state.object_id.to_string(),
        msg_id,
        state.seq_no.load(Ordering::SeqCst),
      ));
    }

    trace!(
      "Applying client updates: {}, origin:{}",
      collab_msg,
      message_origin
    );

    let payload = collab_msg.payload();

    // Spawn a blocking task to handle the message
    let result = Self::handle_message(state, payload, &message_origin, msg_id).await;

    match result {
      Ok(inner_result) => match inner_result {
        Some(response) => Ok(response),
        None => Err(RealtimeError::UnexpectedData("No ack response")),
      },
      Err(err) => Err(RealtimeError::Internal(anyhow!(
        "fail to handle message:{}",
        err
      ))),
    }
  }

  async fn handle_message(
    state: &CollabGroupState,
    payload: &[u8],
    message_origin: &CollabOrigin,
    msg_id: MsgId,
  ) -> Result<Option<CollabAck>, RealtimeError> {
    let mut decoder = DecoderV1::from(payload);
    let reader = MessageReader::new(&mut decoder);
    let mut ack_response = None;
    for msg in reader {
      match msg {
        Ok(msg) => {
          match Self::handle_protocol_message(state, message_origin, msg).await {
            Ok(payload) => {
              // One ClientCollabMessage can have multiple Yrs [Message] in it, but we only need to
              // send one ack back to the client.
              if ack_response.is_none() {
                ack_response = Some(
                  CollabAck::new(
                    message_origin.clone(),
                    state.object_id.to_string(),
                    msg_id,
                    state.seq_no.load(Ordering::SeqCst),
                  )
                  .with_payload(payload.unwrap_or_default()),
                );
              }
            },
            Err(err) => {
              tracing::warn!("[realtime]: failed to handled message: {}", msg_id);
              state.metrics.apply_update_failed_count.inc();

              let code = Self::ack_code_from_error(&err);
              let payload = match err {
                RTProtocolError::MissUpdates {
                  state_vector_v1,
                  reason: _,
                } => state_vector_v1.unwrap_or_default(),
                _ => vec![],
              };

              ack_response = Some(
                CollabAck::new(
                  message_origin.clone(),
                  state.object_id.to_string(),
                  msg_id,
                  state.seq_no.load(Ordering::SeqCst),
                )
                .with_code(code)
                .with_payload(payload),
              );

              break;
            },
          }
        },
        Err(e) => {
          error!("{} => parse sync message failed: {:?}", state.object_id, e);
          break;
        },
      }
    }
    Ok(ack_response)
  }

  async fn handle_protocol_message(
    state: &CollabGroupState,
    origin: &CollabOrigin,
    msg: Message,
  ) -> Result<Option<Vec<u8>>, RTProtocolError> {
    match msg {
      Message::Sync(msg) => match msg {
        SyncMessage::SyncStep1(sv) => Self::handle_sync_step1(state, &sv).await,
        SyncMessage::SyncStep2(update) => Self::handle_sync_step2(state, origin, update).await,
        SyncMessage::Update(update) => Self::handle_update(state, origin, update).await,
      },
      //FIXME: where is the QueryAwareness protocol?
      Message::Awareness(update) => Self::handle_awareness_update(state, origin, update).await,
      Message::Auth(_reason) => Ok(None),
      Message::Custom(_msg) => Ok(None),
    }
  }

  async fn handle_sync_step1(
    state: &CollabGroupState,
    remote_sv: &StateVector,
  ) -> Result<Option<Vec<u8>>, RTProtocolError> {
    if let Ok(sv) = state.state_vector.try_read() {
      // we optimistically try to obtain state vector lock for a fast track:
      // if we remote sv is up-to-date with current one, we don't need to do anything
      match sv.partial_cmp(remote_sv) {
        Some(std::cmp::Ordering::Equal) => return Ok(None), // client and server are in sync
        Some(std::cmp::Ordering::Less) => {
          // server is behind client
          let msg = Message::Sync(SyncMessage::SyncStep1(sv.clone()));
          return Ok(Some(msg.encode_v1()));
        },
        Some(std::cmp::Ordering::Greater) | None => { /* server has some new updates */ },
      }
    }

    // we need to reconstruct document state on the server side
    tracing::debug!("loading collab {}", state.object_id);
    let snapshot = state
      .persister
      .load_compact()
      .await
      .map_err(|err| RTProtocolError::Internal(err.into()))?;

    // prepare document state update and state vector
    let tx = snapshot.collab.transact();
    let doc_state = tx.encode_state_as_update_v1(remote_sv);
    let local_sv = tx.state_vector();
    drop(tx);

    // Retrieve the latest document state from the client after they return online from offline editing.
    tracing::trace!("sending missing data to client ({} bytes)", doc_state.len());
    let mut encoder = EncoderV1::new();
    Message::Sync(SyncMessage::SyncStep2(doc_state)).encode(&mut encoder);
    //FIXME: this should never happen as response to sync step 1 from the client, but rather be
    //  send when a connection is established
    Message::Sync(SyncMessage::SyncStep1(local_sv)).encode(&mut encoder);
    Ok(Some(encoder.to_vec()))
  }

  async fn handle_sync_step2(
    state: &CollabGroupState,
    origin: &CollabOrigin,
    update: Vec<u8>,
  ) -> Result<Option<Vec<u8>>, RTProtocolError> {
    state.metrics.collab_size.observe(update.len() as f64);

    let start = tokio::time::Instant::now();
    // we try to decode update to make sure it's not malformed and to extract state vector
    let (update, decoded_update) = if update.len() <= collab_rt_protocol::LARGE_UPDATE_THRESHOLD {
      let decoded_update = Update::decode_v1(&update)?;
      (update, decoded_update)
    } else {
      tokio::task::spawn_blocking(move || {
        let decoded_update = Update::decode_v1(&update)?;
        Ok::<(Vec<u8>, yrs::Update), yrs::encoding::read::Error>((update, decoded_update))
      })
      .await
      .map_err(|err| RTProtocolError::Internal(err.into()))??
    };
    let missing_updates = {
      let state_vector = state.state_vector.read().await;
      match state_vector.partial_cmp(&decoded_update.state_vector_lower()) {
        None | Some(std::cmp::Ordering::Less) => Some(state_vector.clone()),
        _ => None,
      }
    };

    if let Some(missing_updates) = missing_updates {
      let msg = Message::Sync(SyncMessage::SyncStep1(missing_updates));
      tracing::debug!("subscriber {} send update with missing data", origin);
      Ok(Some(msg.encode_v1()))
    } else {
      let upper_state_vector = decoded_update.state_vector();
      state
        .persister
        .send_update(origin.clone(), update, upper_state_vector)
        .await
        .map_err(|err| RTProtocolError::Internal(err.into()))?;
      let elapsed = start.elapsed();

      state
        .metrics
        .load_collab_time
        .observe(elapsed.as_millis() as f64);

      Ok(None)
    }
  }

  async fn handle_update(
    state: &CollabGroupState,
    origin: &CollabOrigin,
    update: Vec<u8>,
  ) -> Result<Option<Vec<u8>>, RTProtocolError> {
    Self::handle_sync_step2(state, origin, update).await
  }

  async fn handle_awareness_update(
    state: &CollabGroupState,
    origin: &CollabOrigin,
    update: Vec<u8>,
  ) -> Result<Option<Vec<u8>>, RTProtocolError> {
    state
      .persister
      .send_awareness(origin, update)
      .await
      .map_err(|err| RTProtocolError::Internal(err.into()))?;
    Ok(None)
  }

  #[inline]
  fn ack_code_from_error(error: &RTProtocolError) -> AckCode {
    match error {
      RTProtocolError::YrsTransaction(_) => AckCode::Retry,
      RTProtocolError::YrsApplyUpdate(_) => AckCode::CannotApplyUpdate,
      RTProtocolError::YrsEncodeState(_) => AckCode::EncodeStateAsUpdateFail,
      RTProtocolError::MissUpdates { .. } => AckCode::MissUpdate,
      _ => AckCode::Internal,
    }
  }

  /// Check if the group is active. A group is considered active if it has at least one
  /// subscriber
  pub fn is_inactive(&self) -> bool {
    let modified_at = self.modified_at();

    // In debug mode, we set the timeout to 60 seconds
    if cfg!(debug_assertions) {
      trace!(
        "Group:{}:{} is inactive for {} seconds, subscribers: {}",
        self.state.object_id,
        self.state.collab_type,
        modified_at.elapsed().as_secs(),
        self.state.subscribers.len()
      );
      modified_at.elapsed().as_secs() > 60 * 3
    } else {
      let elapsed_secs = modified_at.elapsed().as_secs();
      if elapsed_secs > self.timeout_secs() {
        // Mark the group as inactive if it has been inactive for more than 3 hours, regardless of the number of subscribers.
        // Otherwise, return `true` only if there are no subscribers remaining in the group.
        // If a client modifies a group that has already been marked as inactive (removed),
        // the client will automatically send an initialization sync to reinitialize the group.
        const MAXIMUM_SECS: u64 = 3 * 60 * 60;
        if elapsed_secs > MAXIMUM_SECS {
          info!(
            "Group:{}:{} is inactive for {} seconds, subscribers: {}",
            self.state.object_id,
            self.state.collab_type,
            modified_at.elapsed().as_secs(),
            self.state.subscribers.len()
          );
          true
        } else {
          self.state.subscribers.is_empty()
        }
      } else {
        false
      }
    }
  }

  /// Returns the timeout duration in seconds for different collaboration types.
  ///
  /// Collaborative entities vary in their activity and interaction patterns, necessitating
  /// different timeout durations to balance efficient resource management with a positive
  /// user experience. This function assigns a timeout duration to each collaboration type,
  /// ensuring that resources are utilized judiciously without compromising user engagement.
  ///
  /// # Returns
  /// A `u64` representing the timeout duration in seconds for the collaboration type in question.
  #[inline]
  fn timeout_secs(&self) -> u64 {
    match self.state.collab_type {
      CollabType::Document => 30 * 60, // 30 minutes
      CollabType::Database | CollabType::DatabaseRow => 30 * 60, // 30 minutes
      CollabType::WorkspaceDatabase | CollabType::Folder | CollabType::UserAwareness => 6 * 60 * 60, // 6 hours,
      CollabType::Unknown => {
        10 * 60 // 10 minutes
      },
    }
  }
}

pub trait SubscriptionSink:
  Sink<CollabMessage, Error = RealtimeError> + Send + Sync + Unpin
{
}
impl<T> SubscriptionSink for T where
  T: Sink<CollabMessage, Error = RealtimeError> + Send + Sync + Unpin
{
}

pub trait SubscriptionStream: Stream<Item = MessageByObjectId> + Send + Sync + Unpin {}
impl<T> SubscriptionStream for T where T: Stream<Item = MessageByObjectId> + Send + Sync + Unpin {}

struct Subscription {
  collab_origin: CollabOrigin,
  sink: Box<dyn SubscriptionSink>,
  shutdown: CancellationToken,
}

impl Subscription {
  fn new<S>(sink: S, collab_origin: CollabOrigin, shutdown: CancellationToken) -> Self
  where
    S: SubscriptionSink + 'static,
  {
    Subscription {
      sink: Box::new(sink),
      collab_origin,
      shutdown,
    }
  }
}

impl Drop for Subscription {
  fn drop(&mut self) {
    tracing::trace!("closing subscription: {}", self.collab_origin);
    self.shutdown.cancel();
  }
}

struct CollabPersister {
  uid: i64,
  workspace_id: String,
  object_id: String,
  collab_type: CollabType,
  storage: Arc<dyn CollabStorage>,
  collab_redis_stream: Arc<CollabRedisStream>,
  indexer: Option<Arc<dyn Indexer>>,
  metrics: Arc<CollabRealtimeMetrics>,
  update_sink: CollabUpdateSink,
  awareness_sink: AwarenessUpdateSink,
  /// A grace period for prunning Redis collab updates. Instead of deleting all messages we
  /// read right away, we give 1min for other potential client to catch up.
  prune_grace_period: Duration,
}

impl CollabPersister {
  #[allow(clippy::too_many_arguments)]
  pub fn new(
    uid: i64,
    workspace_id: String,
    object_id: String,
    collab_type: CollabType,
    storage: Arc<dyn CollabStorage>,
    collab_redis_stream: Arc<CollabRedisStream>,
    indexer: Option<Arc<dyn Indexer>>,
    metrics: Arc<CollabRealtimeMetrics>,
    prune_grace_period: Duration,
  ) -> Self {
    let update_sink = collab_redis_stream.collab_update_sink(&workspace_id, &object_id);
    let awareness_sink = collab_redis_stream.awareness_update_sink(&workspace_id, &object_id);
    Self {
      uid,
      workspace_id,
      object_id,
      collab_type,
      storage,
      collab_redis_stream,
      indexer,
      metrics,
      update_sink,
      awareness_sink,
      prune_grace_period,
    }
  }

  async fn send_update(
    &self,
    sender: CollabOrigin,
    update: Vec<u8>,
    state_vector: StateVector,
  ) -> Result<MessageId, StreamError> {
    let len = update.len();
    // send updates to redis queue
    let update = CollabStreamUpdate::new(update, state_vector, sender, UpdateFlags::default());
    let msg_id = self.update_sink.send(&update).await?;
    tracing::trace!(
      "persisted update from {} ({} bytes) - msg id: {}",
      update.sender,
      len,
      msg_id
    );
    Ok(msg_id)
  }

  async fn send_awareness(
    &self,
    sender_session: &CollabOrigin,
    awareness_update: Vec<u8>,
  ) -> Result<MessageId, StreamError> {
    // send awareness updates to redis queue:
    // QUESTION: is it needed? Maybe we could reuse update_sink?
    let len = awareness_update.len();
    let update = AwarenessStreamUpdate {
      data: awareness_update,
      sender: sender_session.clone(),
    };
    let msg_id = self.awareness_sink.send(&update).await?;
    tracing::trace!(
      "persisted awareness from {} ({} bytes) - msg id: {}",
      update.sender,
      len,
      msg_id
    );
    Ok(msg_id)
  }

  /// Loads collab without its history. Used for handling y-sync protocol messages.
  async fn load_compact(&self) -> Result<CollabSnapshot, RealtimeError> {
    // 1. Try to load the latest snapshot from storage
    let start = Instant::now();
    let mut collab = match self.load_collab_full(false).await? {
      Some(collab) => collab,
      None => Collab::new_with_origin(CollabOrigin::Server, self.object_id.clone(), vec![], false),
    };
    self.metrics.load_collab_count.inc();

    // 2. consume all Redis updates on top of it (keep redis msg id)
    let mut last_message_id = None;
    let mut tx = collab.transact_mut();
    let updates = self
      .collab_redis_stream
      .current_collab_updates(
        &self.workspace_id,
        &self.object_id,
        None, //TODO: store Redis last msg id somewhere in doc state snapshot and replay from there
      )
      .await?;
    let mut i = 0;
    for (message_id, update) in updates {
      i += 1;
      let update: Update = update.into_update()?;
      tx.apply_update(update)
        .map_err(|err| RTProtocolError::YrsApplyUpdate(err.to_string()))?;
      last_message_id = Some(message_id); //TODO: shouldn't this happen before decoding?
      self.metrics.apply_update_count.inc();
    }
    drop(tx);
    tracing::trace!(
      "loaded collab compact state: {} replaying {} updates",
      self.object_id,
      i
    );
    self
      .metrics
      .load_collab_time
      .observe(start.elapsed().as_millis() as f64);

    // now we have the most recent version of the document
    let snapshot = CollabSnapshot {
      collab,
      last_message_id,
    };
    Ok(snapshot)
  }

  /// Returns a collab state (with GC turned off), but only if there were any pending updates
  /// waiting to be merged into main document state.
  async fn load_if_changed(&self) -> Result<Option<CollabSnapshot>, RealtimeError> {
    // 1. load pending Redis updates
    let updates = self
      .collab_redis_stream
      .current_collab_updates(&self.workspace_id, &self.object_id, None)
      .await?;

    let start = Instant::now();
    let mut i = 0;
    let mut collab = None;
    let mut last_message_id = None;
    for (message_id, update) in updates {
      i += 1;
      let update: Update = update.into_update()?;
      if collab.is_none() {
        collab = Some(match self.load_collab_full(true).await? {
          Some(collab) => collab,
          None => {
            Collab::new_with_origin(CollabOrigin::Server, self.object_id.clone(), vec![], true)
          },
        })
      };
      let collab = collab.as_mut().unwrap();
      collab
        .transact_mut()
        .apply_update(update)
        .map_err(|err| RTProtocolError::YrsApplyUpdate(err.to_string()))?;
      last_message_id = Some(message_id); //TODO: shouldn't this happen before decoding?
      self.metrics.apply_update_count.inc();
    }

    // if there were no Redis updates, collab is still not initialized
    match collab {
      Some(collab) => {
        self.metrics.load_full_collab_count.inc();
        let elapsed = start.elapsed();
        self
          .metrics
          .load_collab_time
          .observe(elapsed.as_millis() as f64);
        tracing::trace!(
          "loaded collab full state: {} replaying {} updates in {:?}",
          self.object_id,
          i,
          elapsed
        );
        {
          let tx = collab.transact();
          if tx.store().pending_update().is_some() || tx.store().pending_ds().is_some() {
            tracing::trace!(
              "loaded collab {} is incomplete: has pending data",
              self.object_id
            );
          }
        }
        Ok(Some(CollabSnapshot {
          collab,
          last_message_id,
        }))
      },
      None => Ok(None),
    }
  }

  async fn save(&self) -> Result<(), RealtimeError> {
    // load collab but only if there were pending updates in Redis
    if let Some(mut snapshot) = self.load_if_changed().await? {
      tracing::debug!("requesting save for collab {}", self.object_id);
      if let Some(message_id) = snapshot.last_message_id {
        // non-nil message_id means that we had to update the most recent collab state snapshot
        // with new updates from Redis. This means that our snapshot state is newer than the last
        // persisted one in the database
        self.save_attempt(&mut snapshot.collab, message_id).await?;
      }
    } else {
      tracing::trace!("collab {} state has not changed", self.object_id);
    }
    Ok(())
  }

  /// Tries to save provided `snapshot`. This snapshot is expected to have **GC turned off**, as
  /// first it will try to save it as a historical snapshot (will all updates available), then it
  /// will generate another (compact) snapshot variant that will be used as main one for loading
  /// for the sake of y-sync protocol.
  async fn save_attempt(
    &self,
    collab: &mut Collab,
    message_id: MessageId,
  ) -> Result<(), RealtimeError> {
    if !collab.get_awareness().doc().skip_gc() {
      return Err(RealtimeError::UnexpectedData(
        "tried to save history for snapshot with GC turned on",
      ));
    }
    // try to acquire snapshot lease - it's possible that multiple web services will try to
    // perform snapshot at the same time, so we'll use lease to let only one of them atm.
    if let Some(mut lease) = self
      .collab_redis_stream
      .lease(&self.workspace_id, &self.object_id)
      .await?
    {
      // 1. Save full historic document state
      let mut tx = collab.transact_mut();
      let sv = tx.state_vector().encode_v1();
      let doc_state_full = tx.encode_state_as_update_v1(&StateVector::default());
      let full_len = doc_state_full.len();
      let encoded_collab = EncodedCollab::new_v1(sv.clone(), doc_state_full)
        .encode_to_bytes()
        .map_err(|err| RealtimeError::Internal(err.into()))?;
      self
        .metrics
        .full_collab_size
        .observe(encoded_collab.len() as f64);
      let params = InsertSnapshotParams {
        object_id: self.object_id.clone(),
        data: encoded_collab.into(),
        workspace_id: self.workspace_id.clone(),
        collab_type: self.collab_type.clone(),
      };
      self
        .storage
        .create_snapshot(params)
        .await
        .map_err(|err| RealtimeError::Internal(err.into()))?;

      // 2. Generate document state with GC turned on and save it.
      tx.force_gc();
      drop(tx);

      let doc_state_light = collab
        .transact()
        .encode_state_as_update_v1(&StateVector::default());
      let light_len = doc_state_light.len();
      let encoded_collab = EncodedCollab::new_v1(sv, doc_state_light)
        .encode_to_bytes()
        .map_err(|err| RealtimeError::Internal(err.into()))?;
      self
        .metrics
        .collab_size
        .observe(encoded_collab.len() as f64);
      let mut params = CollabParams::new(&self.object_id, self.collab_type.clone(), encoded_collab);

      match self.embeddings(collab).await {
        Ok(embeddings) => params.embeddings = embeddings,
        Err(err) => tracing::warn!("failed to fetch embeddings `{}`: {}", self.object_id, err),
      }

      self
        .storage
        .queue_insert_or_update_collab(&self.workspace_id, &self.uid, params, true)
        .await
        .map_err(|err| RealtimeError::Internal(err.into()))?;

      tracing::debug!(
        "persisted collab {} snapshot at {}: {} and {} bytes",
        self.object_id,
        message_id,
        full_len,
        light_len
      );

      // 3. finally we can drop Redis messages
      let now = SystemTime::UNIX_EPOCH.elapsed().unwrap().as_millis();
      let msg_id = MessageId {
        timestamp_ms: (now - self.prune_grace_period.as_millis()) as u64,
        sequence_number: 0,
      };
      let stream_key = CollabStreamUpdate::stream_key(&self.workspace_id, &self.object_id);
      self
        .collab_redis_stream
        .prune_stream(&stream_key, msg_id)
        .await?;

      let _ = lease.release().await;
    }

    Ok(())
  }

  async fn embeddings(&self, collab: &Collab) -> Result<Option<AFCollabEmbeddings>, AppError> {
    if let Some(indexer) = self.indexer.clone() {
      let params = indexer.embedding_params(collab).await?;
      let embeddings = indexer.embeddings(params).await?;
      Ok(embeddings)
    } else {
      Ok(None)
    }
  }

  async fn load_collab_full(&self, keep_history: bool) -> Result<Option<Collab>, RealtimeError> {
    let doc_state = if keep_history {
      // if we want history-keeping variant, we need to get a snapshot
      let snapshot = self
        .storage
        .get_latest_snapshot(
          &self.workspace_id,
          &self.object_id,
          self.collab_type.clone(),
        )
        .await
        .map_err(|err| RealtimeError::Internal(err.into()))?;
      match snapshot {
        None => None,
        Some(snapshot) => {
          let encoded_collab = EncodedCollab::decode_from_bytes(&snapshot.encoded_collab_v1)
            .map_err(|err| RealtimeError::Internal(err.into()))?;
          Some(encoded_collab.doc_state)
        },
      }
    } else {
      None // if we want a lightweight variant, we'll fallback to default
    };
    let doc_state = match doc_state {
      Some(doc_state) => doc_state,
      None => {
        // we didn't find a snapshot, or we want a lightweight collab version
        let params = QueryCollabParams::new(
          self.object_id.clone(),
          self.collab_type.clone(),
          self.workspace_id.clone(),
        );
        let result = self
          .storage
          .get_encode_collab(GetCollabOrigin::Server, params)
          .await;
        match result {
          Ok(encoded_collab) => encoded_collab.doc_state,
          Err(AppError::RecordNotFound(_)) => return Ok(None),
          Err(err) => return Err(RealtimeError::Internal(err.into())),
        }
      },
    };

    let collab: Collab = Collab::new_with_source(
      CollabOrigin::Server,
      &self.object_id,
      DataSource::DocStateV1(doc_state.into()),
      vec![],
      keep_history, // should we use history-remembering version (true) or lightweight one (false)?
    )?;
    Ok(Some(collab))
  }
}

pub struct CollabSnapshot {
  pub collab: Collab,
  pub last_message_id: Option<MessageId>,
}