mf_collab_client/

utils.rs

use std::time::{SystemTime, UNIX_EPOCH};
use mf_model::mark::Mark;
use mf_model::node_type::NodeEnum;
use mf_model::tree::Tree;
use mf_state::Transaction;
use crate::mapping::{NodeStepConverter, StepConverter};
use crate::AwarenessRef;
use serde_json::Value as JsonValue;
use yrs::{Map, ReadTxn, Transact};
use yrs::{
    types::{array::ArrayRef, map::MapRef, Value},
    Array, ArrayPrelim, MapPrelim, TransactionMut, WriteTxn,
};

use crate::{mapping::Mapper, ClientResult};
use mf_model::{node::Node, attrs::Attrs, types::NodeId};
use std::sync::Arc;
use std::collections::HashMap;

/// 获取当前时间戳（毫秒）
pub fn get_unix_time() -> u64 {
    SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_millis()
        as u64
}

pub struct Utils;
impl Utils {
    
    /// 将 Yrs 文档转换为 ModuForge Tree
    /// 这是从协作状态重建文档树的关键方法
    pub fn apply_yrs_to_tree(
        doc: &yrs::Doc
    ) -> Result<Tree, Box<dyn std::error::Error>> {
        use mf_model::types::NodeId;
        use std::collections::HashMap;

        let root_id = Utils::get_root_id_from_yrs_doc(doc)?;
        let txn = doc.transact();
        let nodes_map =
            txn.get_map("nodes").ok_or("Yrs 文档中没有找到 nodes 映射")?;
        let mut tree_nodes = HashMap::new();
        let mut parent_map = HashMap::new();

        Utils::build_tree_nodes_from_yrs(
            &root_id,
            &nodes_map,
            &txn,
            &mut tree_nodes,
            &mut parent_map,
            None,
        )?;

        let root_node = tree_nodes
            .get(&NodeId::from(root_id))
            .ok_or("根节点不存在")?
            .as_ref()
            .clone();
        let root_enum =
            Utils::build_node_enum_from_map(&root_node, &tree_nodes);
        Ok(Tree::from(root_enum))
    }
    /// 初始化树
    /// 将树同步到 Yrs 文档
    pub async fn apply_tree_to_yrs(
        awareness_ref: AwarenessRef,
        tree: &Tree,
    ) -> ClientResult<()> {
        let mut awareness = awareness_ref.write().await;
        let doc = awareness.doc_mut();
        let mut txn = doc.transact_mut_with(doc.client_id().clone());

        // 清空现有数据（如果有的话）
        let nodes_map = txn.get_or_insert_map("nodes");
        nodes_map.clear(&mut txn);

        // 同步 Tree 中的所有节点到 Yrs 文档
        Utils::sync_tree_to_yrs(tree, &mut txn)?;

        // 添加根节点ID到 meta 区域
        let meta_map = txn.get_or_insert_map("meta");
        meta_map.insert(
            &mut txn,
            "root_id",
            yrs::Any::String(tree.root_id.to_string().into()),
        );

        // 提交事务
        txn.commit();

        tracing::info!(
            "成功初始化树，包含 {} 个节点，根节点ID: {}",
            tree.nodes.len(),
            tree.root_id
        );
        Ok(())
    }

    /// 将树同步到 Yrs 文档
    pub fn sync_tree_to_yrs(
        tree: &Tree,
        txn: &mut yrs::TransactionMut,
    ) -> ClientResult<()> {
        use mf_transform::{step::Step, node_step::AddNodeStep};

        let root_node = tree.get_node(&tree.root_id).unwrap();
        // 获取根节点的所有子树
        if let Some(root_tree) = tree.all_children(&tree.root_id, None) {
            // 创建一个 AddNodeStep 来添加整个子树
            let add_step = AddNodeStep {
                parent_id: tree.root_id.clone(),
                nodes: vec![NodeEnum(
                    root_node.as_ref().clone(),
                    vec![root_tree],
                )],
            };
            let node_step_converter = NodeStepConverter;
            if let Err(e) = node_step_converter
                .apply_to_yrs_txn(&add_step as &dyn Step, txn)
            {
                tracing::error!("🔄 同步树节点到 Yrs 失败: {}", e);
                return Err(anyhow::anyhow!(format!(
                    "Failed to sync tree: {}",
                    e
                ),));
            }
        }

        Ok(())
    }

    /// 将事务应用到 Yrs 文档
    pub async fn apply_transaction_to_yrs(
        awareness_ref: AwarenessRef,
        transaction: &Transaction,
    ) -> ClientResult<()> {
        // 使用异步锁获取房间信息

        let mut awareness = awareness_ref.write().await;
        let doc = awareness.doc_mut();
        let mut txn = doc.transact_mut_with(doc.client_id().clone());

        // 使用全局注册表应用所有事务中的步骤
        let registry = Mapper::global_registry();

        
        let steps = &transaction.steps;
        for step in steps {
            if let Some(converter) = registry.find_converter(step.as_ref())
            {
                if let Err(e) =
                    converter.apply_to_yrs_txn(step.as_ref(), &mut txn)
                {
                    tracing::error!("🔄 应用步骤到 Yrs 事务失败: {}", e);
                }
            } else {
                let type_name = std::any::type_name_of_val(step.as_ref());
                tracing::warn!(
                    "🔄 应用步骤到 Yrs 事务失败: 没有找到步骤的转换器: {}",
                    type_name
                );
            }
        }
    
        // 统一提交所有更改
        txn.commit();
        tracing::debug!(
            "🔄 应用 {} 个步骤到文档: {}",
            transaction.steps.len(),
            doc.client_id()
        );

        Ok(())
    }

    /// 将事务应用到 Yrs 文档
    pub async fn apply_transactions_to_yrs(
        awareness_ref: AwarenessRef,
        transactions: &[Transaction],
    ) -> ClientResult<()> {
        // 使用异步锁获取房间信息

        let mut awareness = awareness_ref.write().await;
        let doc = awareness.doc_mut();
        let mut txn = doc.transact_mut_with(doc.client_id().clone());

        // 使用全局注册表应用所有事务中的步骤
        let registry = Mapper::global_registry();

        for tr in transactions {
            let steps = &tr.steps;
            for step in steps {
                if let Some(converter) = registry.find_converter(step.as_ref())
                {
                    if let Err(e) =
                        converter.apply_to_yrs_txn(step.as_ref(), &mut txn)
                    {
                        tracing::error!("🔄 应用步骤到 Yrs 事务失败: {}", e);
                    }
                } else {
                    let type_name = std::any::type_name_of_val(step.as_ref());
                    tracing::warn!(
                        "🔄 应用步骤到 Yrs 事务失败: 没有找到步骤的转换器: {}",
                        type_name
                    );
                }
            }
        }
        // 统一提交所有更改
        txn.commit();
        tracing::debug!(
            "🔄 应用 {} 个事务到文档: {}",
            transactions.len(),
            doc.client_id()
        );

        Ok(())
    }

    // --- 转换器的辅助方法 ---
    /// 将 JSON 值转换为 Yrs 的 Any 类型
    pub fn json_value_to_yrs_any(value: &JsonValue) -> yrs::Any {
        match value {
            JsonValue::Null => yrs::Any::Null,
            JsonValue::Bool(b) => yrs::Any::Bool(*b),
            JsonValue::Number(n) => {
                if let Some(i) = n.as_i64() {
                    yrs::Any::BigInt(i)
                } else if let Some(f) = n.as_f64() {
                    yrs::Any::Number(f)
                } else {
                    yrs::Any::Null
                }
            },
            JsonValue::String(s) => yrs::Any::String(s.clone().into()),
            JsonValue::Array(arr) => {
                let yrs_array: Vec<yrs::Any> =
                    arr.iter().map(Utils::json_value_to_yrs_any).collect();
                yrs::Any::Array(yrs_array.into())
            },
            JsonValue::Object(obj) => {
                let yrs_map: std::collections::HashMap<String, yrs::Any> = obj
                    .iter()
                    .map(|(k, v)| (k.clone(), Utils::json_value_to_yrs_any(v)))
                    .collect();
                yrs::Any::Map(yrs_map.into())
            },
        }
    }

    /// 将标记添加到 Yrs 数组中
    pub fn add_mark_to_array(
        marks_array: &ArrayRef,
        txn: &mut TransactionMut,
        mark: &Mark,
    ) {
        let mark_map = MapPrelim::<yrs::Any>::from([
            ("type".to_string(), yrs::Any::String(mark.r#type.clone().into())),
            ("attrs".to_string(), {
                let attrs_map: std::collections::HashMap<String, yrs::Any> =
                    mark.attrs
                        .iter()
                        .map(|(k, v)| {
                            (k.clone(), Utils::json_value_to_yrs_any(v))
                        })
                        .collect();
                yrs::Any::Map(attrs_map.into())
            }),
        ]);
        marks_array.push_back(txn, mark_map);
    }

    /// 获取或创建节点数据映射
    pub fn get_or_create_node_data_map(
        nodes_map: &MapRef,
        txn: &mut TransactionMut,
        node_id: &str,
    ) -> MapRef {
        if let Some(Value::YMap(map)) = nodes_map.get(txn, node_id) {
            map
        } else {
            nodes_map.insert(
                txn,
                node_id.to_string(),
                MapPrelim::<yrs::Any>::new(),
            )
        }
    }

    /// 获取或创建节点属性映射
    pub fn get_or_create_node_attrs_map(
        node_data_map: &MapRef,
        txn: &mut TransactionMut,
    ) -> MapRef {
        if let Some(Value::YMap(map)) = node_data_map.get(txn, "attrs") {
            map
        } else {
            node_data_map.insert(txn, "attrs", MapPrelim::<yrs::Any>::new())
        }
    }

    /// 获取或创建标记数组
    pub fn get_or_create_content_array(
        node_data_map: &MapRef,
        txn: &mut TransactionMut,
    ) -> ArrayRef {
        if let Some(Value::YArray(array)) = node_data_map.get(txn, "content") {
            array
        } else {
            node_data_map.insert(
                txn,
                "content",
                ArrayPrelim::from(Vec::<yrs::Any>::new()),
            )
        }
    }

    /// 将 Yrs 的 Any 类型转换为 JSON 值
    pub fn yrs_any_to_json_value(value: &yrs::Any) -> Option<JsonValue> {
        match value {
            yrs::Any::Null => Some(JsonValue::Null),
            yrs::Any::Bool(b) => Some(JsonValue::Bool(*b)),
            yrs::Any::Number(n) => {
                Some(JsonValue::Number(serde_json::Number::from_f64(*n)?))
            },
            yrs::Any::BigInt(i) => {
                Some(JsonValue::Number(serde_json::Number::from(*i)))
            },
            yrs::Any::String(s) => Some(JsonValue::String(s.to_string())),
            yrs::Any::Array(arr) => {
                let json_array: Vec<JsonValue> = arr
                    .iter()
                    .filter_map(Utils::yrs_any_to_json_value)
                    .collect();
                Some(JsonValue::Array(json_array))
            },
            yrs::Any::Map(map) => {
                let json_map: std::collections::HashMap<String, JsonValue> =
                    map.iter()
                        .filter_map(|(k, v)| {
                            Utils::yrs_any_to_json_value(v)
                                .map(|json_v| (k.to_string(), json_v))
                        })
                        .collect();
                Some(JsonValue::Object(serde_json::Map::from_iter(json_map)))
            },
            _ => None, // 处理其他类型，如 YText, YMap 等
        }
    }

    /// 获取或创建标记数组
    pub fn get_or_create_marks_array(
        node_data_map: &MapRef,
        txn: &mut TransactionMut,
    ) -> ArrayRef {
        if let Some(Value::YArray(array)) = node_data_map.get(txn, "marks") {
            array
        } else {
            node_data_map.insert(
                txn,
                "marks",
                ArrayPrelim::from(Vec::<yrs::Any>::new()),
            )
        }
    }

    /// 从 Yrs 文档中获取根节点ID
    pub fn get_root_id_from_yrs_doc(
        doc: &yrs::Doc
    ) -> Result<String, Box<dyn std::error::Error>> {
        let txn = doc.transact();
        // 优先从 meta 区域读取
        if let Some(meta_map) = txn.get_map("meta") {
            if let Some(yrs::types::Value::Any(any)) =
                meta_map.get(&txn, "root_id")
            {
                return Ok(any.to_string());
            }
        }
        // fallback: 兼容老数据，取 nodes_map 第一个节点
        let nodes_map =
            txn.get_map("nodes").ok_or("Yrs 文档中没有找到 nodes 映射")?;
        for (key, _) in nodes_map.iter(&txn) {
            return Ok(key.to_string());
        }
        Err("Yrs 文档中没有找到根节点".into())
    }

    /// 从 Yrs 文档的 nodes_map 递归构建所有节点和 parent_map
    pub fn build_tree_nodes_from_yrs(
        node_id: &str,
        nodes_map: &yrs::types::map::MapRef,
        txn: &yrs::Transaction,
        tree_nodes: &mut HashMap<NodeId, Arc<Node>>,
        parent_map: &mut HashMap<NodeId, NodeId>,
        parent_id: Option<&NodeId>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let node_data = nodes_map.get(txn, node_id);
        if node_data.is_none() {
            return Err(format!("节点 {} 在 Yrs 文档中不存在", node_id).into());
        }
        let node_data = node_data.unwrap();
        if let yrs::types::Value::YMap(node_map) = node_data {
            // 提取节点类型
            let node_type = node_map
                .get(txn, "type")
                .and_then(|v| match v {
                    yrs::types::Value::Any(any) => Some(any.to_string()),
                    _ => None,
                })
                .unwrap_or_else(|| "unknown".to_string());

            // 提取属性
            let mut attrs = Attrs::default();
            if let Some(attrs_map) = node_map.get(txn, "attrs") {
                if let yrs::types::Value::YMap(attrs_yrs_map) = attrs_map {
                    for (key, value) in attrs_yrs_map.iter(txn) {
                        if let yrs::types::Value::Any(any_value) = value {
                            if let Some(json_value) =
                                Utils::yrs_any_to_json_value(&any_value)
                            {
                                attrs.insert(key.to_string(), json_value);
                            }
                        }
                    }
                }
            }

            // 提取内容（子节点ID列表）
            let mut content = im::Vector::new();
            if let Some(content_array) = node_map.get(txn, "content") {
                if let yrs::types::Value::YArray(content_yrs_array) =
                    content_array
                {
                    for item in content_yrs_array.iter(txn) {
                        if let yrs::types::Value::Any(any) = item {
                            content.push_back(NodeId::from(any.to_string()));
                        }
                    }
                }
            }

            // 提取标记
            let mut marks = im::Vector::new();
            if let Some(marks_array) = node_map.get(txn, "marks") {
                if let yrs::types::Value::YArray(marks_yrs_array) = marks_array
                {
                    for item in marks_yrs_array.iter(txn) {
                        if let yrs::types::Value::YMap(mark_map) = item {
                            let mark_type = mark_map
                                .get(txn, "type")
                                .and_then(|v| match v {
                                    yrs::types::Value::Any(any) => {
                                        Some(any.to_string())
                                    },
                                    _ => None,
                                })
                                .unwrap_or_else(|| "unknown".to_string());

                            let mut mark_attrs = Attrs::default();
                            if let Some(mark_attrs_map) =
                                mark_map.get(txn, "attrs")
                            {
                                if let yrs::types::Value::YMap(attrs_yrs_map) =
                                    mark_attrs_map
                                {
                                    for (key, value) in attrs_yrs_map.iter(txn)
                                    {
                                        if let yrs::types::Value::Any(
                                            any_value,
                                        ) = value
                                        {
                                            if let Some(json_value) =
                                                Utils::yrs_any_to_json_value(
                                                    &any_value,
                                                )
                                            {
                                                mark_attrs.insert(
                                                    key.to_string(),
                                                    json_value,
                                                );
                                            }
                                        }
                                    }
                                }
                            }

                            marks.push_back(Mark {
                                r#type: mark_type,
                                attrs: mark_attrs,
                            });
                        }
                    }
                }
            }

            // 创建节点
            let content_vec: Vec<String> =
                content.clone().into_iter().collect();
            let marks_vec: Vec<Mark> = marks.clone().into_iter().collect();
            let node =
                Node::new(node_id, node_type, attrs, content_vec, marks_vec);

            let node_id_typed = NodeId::from(node_id);
            tree_nodes.insert(node_id_typed.clone(), Arc::new(node));

            // 记录父子关系
            if let Some(parent) = parent_id {
                parent_map.insert(node_id_typed.clone(), parent.clone());
            }

            // 递归处理子节点
            for child_id in content {
                Utils::build_tree_nodes_from_yrs(
                    &child_id,
                    nodes_map,
                    txn,
                    tree_nodes,
                    parent_map,
                    Some(&node_id_typed),
                )?;
            }
        }
        Ok(())
    }

    /// 递归构建 NodeEnum
    pub fn build_node_enum_from_map(
        node: &Node,
        tree_nodes: &HashMap<NodeId, Arc<Node>>,
    ) -> mf_model::node_type::NodeEnum {
        let mut children = Vec::new();
        for child_id in &node.content {
            if let Some(child_node) = tree_nodes.get(child_id) {
                children.push(Utils::build_node_enum_from_map(
                    child_node, tree_nodes,
                ));
            }
        }
        mf_model::node_type::NodeEnum(node.clone(), children)
    }
}