affine_doc_loader 0.1.9

use std::{collections::HashSet, ops::Range};

use pulldown_cmark::{Event, Parser, Tag, TagEnd};

use super::{
  BlockFlavour, BlockNode, BlockSpec, MAX_MARKDOWN_CHARS, ParseError, markdown_options,
  parse_markdown_blocks_with_id_hints,
};

pub(super) struct MarkdownTarget {
  pub nodes: Vec<SourceNode>,
}

#[derive(Debug, Clone)]
pub(super) struct SourceNode {
  pub id: Option<String>,
  pub flavour: String,
  pub spec: Option<BlockSpec>,
  pub opaque: Option<String>,
  pub children: Vec<SourceNode>,
  pub range: Range<usize>,
}

struct BlockMarker {
  id: String,
  flavour: String,
  end: bool,
  opaque: bool,
  start: usize,
}

fn marker(line: &str) -> Result<Option<BlockMarker>, ParseError> {
  let line = line.trim();
  let Some(body) = line.strip_prefix("<!--").and_then(|s| s.strip_suffix("-->")) else {
    return Ok(None);
  };
  if !body.trim_start().starts_with("block_id=") {
    return Ok(None);
  }
  let mut id = None;
  let mut flavour = None;
  let mut end = false;
  let mut opaque = false;
  for token in body.split_whitespace() {
    if let Some(value) = token.strip_prefix("block_id=") {
      if value.is_empty() || id.replace(value.to_string()).is_some() {
        return Err(ParseError::ParserError("invalid_block_marker".into()));
      }
    } else if let Some(value) = token.strip_prefix("flavour=") {
      if value.is_empty() || flavour.replace(value.to_string()).is_some() {
        return Err(ParseError::ParserError("invalid_block_marker".into()));
      }
    } else if token == "opaque=true" && !opaque {
      opaque = true;
    } else if matches!(token, "end" | "type=end" | "end=true") && !end {
      end = true;
    } else {
      return Err(ParseError::ParserError("invalid_block_marker".into()));
    }
  }
  Ok(Some(BlockMarker {
    id: id.ok_or_else(|| ParseError::ParserError("invalid_block_marker".into()))?,
    flavour: flavour.ok_or_else(|| ParseError::ParserError("invalid_block_marker".into()))?,
    end,
    opaque,
    start: 0,
  }))
}

pub(super) fn parse_markdown_targets(markdown: &str) -> Result<MarkdownTarget, ParseError> {
  if markdown.len() > MAX_MARKDOWN_CHARS {
    return Err(ParseError::ParserError("markdown_too_large".into()));
  }
  // Only standalone HTML blocks can carry markers. Text inside code, lists and
  // blockquotes is Markdown content, not synchronization metadata.
  let mut markers = Vec::new();
  let mut depth = 0;
  let mut html = false;
  for (event, range) in Parser::new(markdown).into_offset_iter() {
    match event {
      Event::Start(tag) => {
        html = depth == 0 && matches!(tag, Tag::HtmlBlock);
        depth += 1;
        if depth > 128 {
          return Err(ParseError::ParserError("markdown_depth_exceeded".into()));
        }
      }
      Event::End(tag) => {
        depth -= 1;
        if matches!(tag, TagEnd::HtmlBlock) {
          html = false;
        }
      }
      Event::Html(_) if html && depth == 1 => {
        if let Some(mut marker) = marker(&markdown[range.clone()])? {
          marker.start = range.start;
          markers.push((range, marker));
        }
      }
      Event::InlineHtml(html) | Event::Html(html) if html.trim_start().starts_with("<!-- affine:block-id:") => {
        let valid = html
          .trim()
          .strip_prefix("<!-- affine:block-id:merge-v1=")
          .and_then(|s| s.strip_suffix("-->"))
          .map(str::trim)
          .is_some_and(|id| !id.is_empty() && !id.chars().any(char::is_whitespace));
        if !valid {
          return Err(ParseError::ParserError("invalid_block_id_hint".into()));
        }
      }
      _ => {}
    }
  }

  let mut output = MarkdownTarget { nodes: Vec::new() };
  let mut seen = HashSet::new();
  let mut active: Option<BlockMarker> = None;
  let mut start = 0;
  for (range, marker) in markers {
    append_segment(&mut output, active.as_ref(), markdown, start..range.start)?;
    if marker.end {
      if !active
        .as_ref()
        .is_some_and(|open| open.id == marker.id && open.flavour == marker.flavour)
      {
        return Err(ParseError::ParserError("unmatched_block_marker".into()));
      }
      if let Some(node) = output.nodes.last_mut() {
        node.range.end = range.end;
      }
      active = None;
    } else {
      if !seen.insert(marker.id.clone()) {
        return Err(ParseError::ParserError("ambiguous_block_markers".into()));
      }
      active = Some(marker);
    }
    start = range.end;
  }
  append_segment(&mut output, active.as_ref(), markdown, start..markdown.len())?;
  Ok(output)
}

fn append_segment(
  output: &mut MarkdownTarget,
  marker: Option<&BlockMarker>,
  markdown: &str,
  range: Range<usize>,
) -> Result<(), ParseError> {
  let body = &markdown[range.clone()];
  if let Some(marker) = marker {
    if marker.opaque
      || matches!(
        marker.flavour.as_str(),
        "affine:database" | "affine:callout" | "affine:table"
      )
      || BlockFlavour::from_str(&marker.flavour).is_none()
    {
      output.nodes.push(SourceNode {
        id: Some(marker.id.clone()),
        flavour: marker.flavour.clone(),
        spec: None,
        opaque: Some(body.trim().to_string()),
        children: Vec::new(),
        range: marker.start..range.end,
      });
      return Ok(());
    }
  }
  let mut spans = Vec::new();
  let mut list_spans = Vec::new();
  let mut depth = 0;
  for (event, span) in Parser::new_ext(body, markdown_options()).into_offset_iter() {
    match event {
      Event::Start(tag) => {
        if matches!(tag, Tag::Item) {
          let line_start = body[..span.start].rfind('\n').map_or(0, |i| i + 1);
          let item_start = if body[line_start..span.start].trim().is_empty() {
            line_start
          } else {
            span.start
          };
          list_spans.push((range.start + item_start)..(range.start + span.end));
        }
        if depth == 0 && !matches!(tag, Tag::List(_)) || depth == 1 && matches!(tag, Tag::Item) {
          spans.push(span);
        }
        depth += 1;
        if depth > 128 {
          return Err(ParseError::ParserError("markdown_depth_exceeded".into()));
        }
      }
      Event::End(_) => depth -= 1,
      Event::Rule if depth == 0 => spans.push(span),
      _ => {}
    }
  }
  let nodes = parse_markdown_blocks_with_id_hints(body, "merge-v1")?;
  if marker.is_some() && nodes.is_empty() {
    return Err(ParseError::ParserError("empty_block_marker".into()));
  }
  // A source unit must have a reliable span before it can be rewritten.
  if spans.len() != nodes.len() && !nodes.is_empty() {
    if nodes.len() == 1 {
      spans = vec![0..body.len()];
    } else {
      return Err(ParseError::ParserError("unsupported_source_ranges".into()));
    }
  }
  fn assign_ranges(node: &mut SourceNode, spans: &mut impl Iterator<Item = Range<usize>>) {
    if node.flavour == "affine:list" {
      node.range = spans.next().unwrap_or(0..0);
    }
    for child in &mut node.children {
      assign_ranges(child, spans);
    }
  }
  let mut list_spans = list_spans.into_iter();
  for (index, node) in nodes.into_iter().enumerate() {
    let mut node = from_block(node);
    assign_ranges(&mut node, &mut list_spans);
    let start = if index == 0 {
      marker.map_or(range.start, |m| m.start)
    } else {
      range.start + spans[index].start
    };
    let end = spans.get(index + 1).map_or(range.end, |s| range.start + s.start);
    node.range = start..end;
    if index == 0
      && let Some(marker) = marker
    {
      if node.flavour != marker.flavour {
        return Err(ParseError::ParserError("marker_flavour_mismatch".into()));
      }
      if node.id.as_ref().is_some_and(|id| id != &marker.id) {
        return Err(ParseError::ParserError("conflicting_block_markers".into()));
      }
      node.id = Some(marker.id.clone());
    }
    output.nodes.push(node);
  }
  Ok(())
}

fn from_block(node: BlockNode) -> SourceNode {
  SourceNode {
    id: node.id,
    flavour: node.spec.flavour.as_str().into(),
    spec: Some(node.spec),
    opaque: None,
    children: node.children.into_iter().map(from_block).collect(),
    range: 0..0,
  }
}