affine_doc_loader 0.1.9

use super::{
  Any, BTreeMap, BlockFlavour, BlockSpec, DeltaToMdOptions, Doc, HashMap, HashSet, MAX_BLOCKS, Map, MarkdownRenderer,
  MergeError, NOTE_FLAVOUR, PAGE_FLAVOUR, ParseError, ProjectedNode, SourceNode, TextDeltaOp, TextInsert,
  collect_child_ids, failure, get_string, load_doc, parse_markdown_blocks,
};

pub(super) struct State {
  pub doc: Doc,
  pub scope: String,
  pub pool: HashMap<String, Map>,
  pub nodes: Vec<ProjectedNode>,
}

impl State {
  pub fn load(binary: &[u8], doc_id: &str, scope: Option<&str>) -> Result<Self, ParseError> {
    let doc = load_doc(binary, Some(doc_id))?;
    if doc.has_pending_updates() {
      return Err(ParseError::InvalidBinary);
    }
    let blocks = doc.get_map("blocks")?;
    let mut pool = HashMap::new();
    for (id, value) in blocks.iter() {
      if pool.len() >= MAX_BLOCKS {
        return Err(failure(id, "blocks", MergeError::BudgetExceeded));
      }
      let block = value.to_map().ok_or(ParseError::InvalidBinary)?;
      if get_string(&block, "sys:id").as_deref() != Some(id) {
        return Err(ParseError::InvalidBinary);
      }
      if let Some(children) = block.get("sys:children") {
        let array = children.to_array().ok_or(ParseError::InvalidBinary)?;
        if array.iter().any(|v| !matches!(v.to_any(), Some(Any::String(_)))) {
          return Err(ParseError::InvalidBinary);
        }
      }
      pool.insert(id.to_string(), block);
    }
    let pages: Vec<_> = pool
      .iter()
      .filter(|(_, m)| get_string(*m, "sys:flavour").as_deref() == Some(PAGE_FLAVOUR))
      .collect();
    if pages.len() != 1 {
      return Err(ParseError::InvalidBinary);
    }
    let notes: Vec<_> = collect_child_ids(pages[0].1)
      .into_iter()
      .filter(|id| {
        pool
          .get(id)
          .is_some_and(|m| get_string(m, "sys:flavour").as_deref() == Some(NOTE_FLAVOUR))
      })
      .collect();
    let scope = match scope {
      Some(id) if notes.iter().any(|n| n == id) => id.to_owned(),
      None if notes.len() == 1 => notes[0].clone(),
      _ => return Err(ParseError::ParserError("explicit_note_scope_required".into())),
    };
    // Validate children as the authoritative graph, including opaque subtrees.
    let mut parents = HashMap::new();
    for (id, block) in &pool {
      for child in collect_child_ids(block) {
        if !pool.contains_key(&child) || parents.insert(child, id.clone()).is_some() {
          return Err(ParseError::InvalidBinary);
        }
      }
    }
    for id in pool.keys() {
      let mut visited = HashSet::new();
      let mut cursor = Some(id);
      while let Some(id) = cursor {
        if visited.len() > 128 || !visited.insert(id) {
          return Err(ParseError::InvalidBinary);
        }
        cursor = parents.get(id);
      }
    }
    let nodes = project_children(&scope, &pool)?;
    Ok(Self {
      doc,
      scope,
      pool,
      nodes,
    })
  }
}

fn project_children(parent: &str, pool: &HashMap<String, Map>) -> Result<Vec<ProjectedNode>, ParseError> {
  collect_child_ids(&pool[parent])
    .iter()
    .map(|id| {
      let block = &pool[id];
      let kind = get_string(block, "sys:flavour").ok_or(ParseError::InvalidBinary)?;
      let children = collect_child_ids(block);
      let supported = BlockFlavour::from_str(&kind)
        .is_some_and(|flavour| !matches!(flavour, BlockFlavour::Table | BlockFlavour::Callout))
        && (children.is_empty() || kind == "affine:list");
      if !supported {
        return Ok(ProjectedNode {
          id: Some(id.clone()),
          kind,
          properties: BTreeMap::new(),
          text: Vec::new(),
          opaque: Some(String::new()),
          children: Vec::new(),
        });
      }
      let spec = BlockSpec::from_block_map(block)?;
      // Compare only the fields this profile actually exports. The original
      // Y.Text and block map remain the source of metadata and identity.
      let mut own = String::new();
      MarkdownRenderer::new(&DeltaToMdOptions::source_profile()).write_block(&mut own, &spec, 0);
      let source = parse_markdown_blocks(&own)?;
      if source.len() != 1 {
        return Err(failure(id, "projection", MergeError::InvalidProjection));
      }
      let mut node = from_spec(&source[0].spec, Some(id.clone()));
      let mut visible = spec.clone();
      for op in &mut visible.text {
        if let TextDeltaOp::Insert {
          format: Some(format), ..
        } = op
        {
          format.retain(|key, _| super::EDITABLE_MARKS.contains(&key.as_str()));
        }
      }
      let roundtrips = from_spec(&visible, Some(id.clone())).text == node.text;
      node.children = project_children(id, pool)?;
      if !roundtrips || node.children.iter().any(|child| child.opaque.is_some()) {
        node.properties.clear();
        node.text.clear();
        node.children.clear();
        node.opaque = Some(String::new());
      }
      Ok(node)
    })
    .collect()
}

pub(super) fn from_source(node: &SourceNode) -> ProjectedNode {
  let mut view = if let Some(spec) = &node.spec {
    from_spec(spec, node.id.clone())
  } else {
    ProjectedNode {
      id: node.id.clone(),
      kind: node.flavour.clone(),
      properties: BTreeMap::new(),
      text: Vec::new(),
      opaque: node.opaque.clone(),
      children: Vec::new(),
    }
  };
  view.children = node.children.iter().map(from_source).collect();
  view
}

fn from_spec(spec: &BlockSpec, id: Option<String>) -> ProjectedNode {
  let mut properties = BTreeMap::new();
  if let Some(value) = spec.block_type_str() {
    properties.insert("prop:type".into(), Any::String(value.into()));
  }
  if let Some(value) = spec.checked {
    properties.insert("prop:checked".into(), Any::from(value));
  }
  if let Some(value) = &spec.language {
    properties.insert("prop:language".into(), Any::String(value.clone()));
  }
  if let Some(value) = spec.order {
    properties.insert("prop:order".into(), Any::from(value));
  }
  if let Some(image) = &spec.image {
    properties.insert("prop:sourceId".into(), Any::String(image.source_id.clone()));
    if let Some(value) = &image.caption {
      properties.insert("prop:caption".into(), Any::String(value.clone()));
    }
    if let Some(value) = image.width {
      properties.insert("prop:width".into(), Any::from(value));
    }
    if let Some(value) = image.height {
      properties.insert("prop:height".into(), Any::from(value));
    }
  }
  if let Some(bookmark) = &spec.bookmark {
    properties.insert("prop:url".into(), Any::String(bookmark.url.clone()));
  }
  if let Some(embed) = &spec.embed_iframe {
    properties.insert("prop:url".into(), Any::String(embed.url.clone()));
  }
  if let Some(embed) = &spec.embed_youtube {
    properties.insert("prop:videoId".into(), Any::String(embed.video_id.clone()));
  }
  let mut text: Vec<TextDeltaOp> = Vec::new();
  for op in &spec.text {
    if let TextDeltaOp::Insert {
      insert: TextInsert::Text(value),
      format,
    } = op
    {
      if value.is_empty() {
        continue;
      }
      let format = format.clone().filter(|f| !f.is_empty());
      if let Some(TextDeltaOp::Insert {
        insert: TextInsert::Text(previous),
        format: old,
      }) = text.last_mut()
        && *old == format
      {
        previous.push_str(value);
      } else {
        text.push(TextDeltaOp::Insert {
          insert: TextInsert::Text(value.clone()),
          format,
        });
      }
    }
  }
  ProjectedNode {
    id,
    kind: spec.flavour.as_str().into(),
    properties,
    text,
    opaque: None,
    children: Vec::new(),
  }
}

pub(super) fn pair_baseline(actual: &[ProjectedNode], source: &mut [ProjectedNode]) -> Result<(), ParseError> {
  if actual.len() != source.len() {
    return Err(failure("", "baseline", MergeError::InvalidProjection));
  }
  for (a, b) in actual.iter().zip(source) {
    if b.id.as_ref().is_some_and(|id| Some(id) != a.id.as_ref())
      || a.kind != b.kind
      || a.properties != b.properties
      || a.text != b.text
      || a.opaque.is_some() != b.opaque.is_some()
    {
      return Err(failure(
        a.id.as_deref().unwrap_or(""),
        "baseline",
        MergeError::InvalidProjection,
      ));
    }
    b.id = a.id.clone();
    pair_baseline(&a.children, &mut b.children)?;
  }
  Ok(())
}

pub(super) fn flatten<'a>(nodes: &'a [ProjectedNode], out: &mut Vec<&'a ProjectedNode>) {
  for node in nodes {
    out.push(node);
    flatten(&node.children, out);
  }
}
pub(super) fn flatten_source<'a>(nodes: &'a [SourceNode], out: &mut Vec<&'a SourceNode>) {
  for node in nodes {
    out.push(node);
    flatten_source(&node.children, out);
  }
}