affine_importer 0.1.2

use std::collections::{BTreeMap, BTreeSet};

use nanoid::nanoid;

use super::{
  ImportAssetRef, ImportDocRef, IndexedArchive, add_asset_batch, doc_snapshot, empty_batch, entry_from_indexed,
  file_name, folders_for_doc_path, frontmatter_meta, has_emitted_assets, hash_bytes, push_skipped_doc_warning,
  read_archive, register_page_path, strip_extension, strip_frontmatter,
};
use crate::{
  ImportBatch, ImportBatchLimits, ImportCursor, ImportError, ImportOptions, ImportProgress, ImportProvider,
  ImportResult, ImportSource, ImportedDocMeta, ImportedDocSnapshot,
};

pub(super) struct MarkdownZipProvider;

impl ImportProvider for MarkdownZipProvider {
  fn format(&self) -> &'static str {
    "markdownZip"
  }

  fn create_cursor(
    &self,
    source: ImportSource,
    options: ImportOptions,
    limits: ImportBatchLimits,
  ) -> ImportResult<Box<dyn ImportCursor>> {
    let archive = read_archive(source, false, &options)?;
    Ok(Box::new(MarkdownImportCursor::new(archive, options, limits)?))
  }
}

struct MarkdownImportCursor {
  archive: IndexedArchive,
  docs: Vec<ImportDocRef>,
  assets: Vec<ImportAssetRef>,
  page_ids_by_path: BTreeMap<String, String>,
  blob_ids_by_path: BTreeMap<String, String>,
  folders: Vec<crate::FolderHierarchyDelta>,
  options: ImportOptions,
  limits: ImportBatchLimits,
  next_doc: usize,
  emitted_assets: BTreeSet<usize>,
  emitted_final: bool,
  completed: usize,
  total: usize,
  entry_id: Option<String>,
}

impl MarkdownImportCursor {
  fn new(archive: IndexedArchive, options: ImportOptions, limits: ImportBatchLimits) -> ImportResult<Self> {
    let mut docs = Vec::new();
    let mut assets = Vec::new();
    let mut page_ids_by_path = BTreeMap::new();
    let mut blob_ids_by_path = BTreeMap::new();
    for (entry_index, entry) in archive.entries.iter().enumerate() {
      if entry.meta.path.to_lowercase().ends_with(".md") {
        let doc_id = nanoid!();
        register_page_path(&mut page_ids_by_path, &entry.meta.path, &doc_id);
        docs.push(ImportDocRef {
          entry_index,
          doc_id,
          title: None,
          icon: None,
        });
      }
    }
    for (entry_index, entry) in archive.entries.iter().enumerate() {
      if !entry.meta.path.to_lowercase().ends_with(".md") {
        let blob_id = hash_bytes(&archive.read_entry(entry)?);
        blob_ids_by_path.insert(entry.meta.path.clone(), blob_id.clone());
        assets.push(ImportAssetRef { entry_index, blob_id });
      }
    }
    if docs.is_empty() {
      return Err(ImportError::InvalidSource(
        "No Markdown pages found in the archive".to_string(),
      ));
    }
    let folders = markdown_folders(&docs, &archive);
    let total = docs.len();
    let entry_id = docs.first().map(|doc| doc.doc_id.clone());
    Ok(Self {
      archive,
      docs,
      assets,
      page_ids_by_path,
      blob_ids_by_path,
      folders,
      options,
      limits,
      next_doc: 0,
      emitted_assets: BTreeSet::new(),
      emitted_final: false,
      completed: 0,
      total,
      entry_id,
    })
  }
}

impl ImportCursor for MarkdownImportCursor {
  fn total(&self) -> usize {
    self.total
  }

  fn next_batch(&mut self) -> ImportResult<Option<ImportBatch>> {
    if self.completed >= self.total && self.emitted_final {
      return Ok(None);
    }
    if self.options.cancel_after_entries == Some(self.completed) {
      return Err(ImportError::Cancelled);
    }

    let mut batch = empty_batch(self.total);
    batch.entry_id = self.entry_id.clone();

    let end = (self.next_doc + self.limits.max_docs.max(1)).min(self.docs.len());
    for doc in &self.docs[self.next_doc..end] {
      let entry = &self.archive.entries[doc.entry_index];
      let bytes = self.archive.read_entry(entry)?;
      let source_path = entry.meta.path.clone();
      let markdown = String::from_utf8_lossy(&bytes).to_string();
      let entry = entry_from_indexed(entry, bytes);
      let prepared = prepare_markdown(&entry, &markdown);
      let snapshot = match doc_snapshot(
        &doc.doc_id,
        &prepared.title,
        &prepared.content,
        &source_path,
        &self.page_ids_by_path,
        &self.blob_ids_by_path,
      ) {
        Ok(snapshot) => snapshot,
        Err(ImportError::Document(error)) => {
          push_skipped_doc_warning(&mut batch, &source_path, error);
          continue;
        }
        Err(error) => return Err(error),
      };
      batch.docs.push(ImportedDocSnapshot {
        id: doc.doc_id.clone(),
        snapshot,
        meta: prepared.meta,
      });
    }
    add_asset_batch(
      &self.archive,
      &self.limits,
      &mut self.emitted_assets,
      &mut batch,
      &self.assets,
    )?;
    if end == self.docs.len() && has_emitted_assets(&self.emitted_assets, &self.assets) {
      batch.folders.extend(self.folders.clone());
      self.emitted_final = true;
    }
    self.next_doc = end;
    self.completed = end;
    batch.progress = ImportProgress {
      completed: self.completed,
      total: self.total,
    };
    batch.done = self.completed >= self.total && self.emitted_final;
    Ok(Some(batch))
  }
}

fn markdown_folders(docs: &[ImportDocRef], archive: &IndexedArchive) -> Vec<crate::FolderHierarchyDelta> {
  let mut folders = Vec::new();
  for doc in docs {
    let entry = &archive.entries[doc.entry_index];
    folders.extend(folders_for_doc_path(&entry.meta.path, &doc.doc_id, None));
  }
  folders
}

struct PreparedMarkdown {
  title: String,
  content: String,
  meta: Option<ImportedDocMeta>,
}

fn prepare_markdown(entry: &crate::vfs::VfsEntry, markdown: &str) -> PreparedMarkdown {
  let (content, frontmatter) = strip_frontmatter(markdown);
  let title = frontmatter
    .title
    .clone()
    .unwrap_or_else(|| strip_extension(file_name(&entry.path)).to_string());
  PreparedMarkdown {
    meta: frontmatter_meta(&frontmatter, &title),
    title,
    content,
  }
}

#[cfg(test)]
mod tests {
  use super::{
    super::tests::{merge_batch, merge_cursor, unlimited_limits, zip, zip_cursor},
    *,
  };

  fn markdown_zip_cursor(
    bytes: impl AsRef<[u8]>,
    options: ImportOptions,
    limits: ImportBatchLimits,
  ) -> ImportResult<Box<dyn crate::ImportCursor>> {
    zip_cursor("markdownZip", bytes, options, limits)
  }

  fn plan_markdown_zip(bytes: impl AsRef<[u8]>) -> ImportResult<ImportBatch> {
    markdown_zip_cursor(bytes, ImportOptions::default(), unlimited_limits()).and_then(merge_cursor)
  }

  #[test]
  fn markdown_zip_plans_docs_assets_folders_and_progress() {
    let bytes = zip(&[
      ("Export/entry.md", b"entry body"),
      ("Export/folder/target.md", b"target body"),
      ("Export/assets/logo.png", &[137, 80, 78, 71]),
    ]);

    let batch = plan_markdown_zip(bytes).unwrap();

    assert_eq!(batch.docs.len(), 2);
    assert_eq!(batch.blobs.len(), 1);
    assert_eq!(batch.progress.completed, 2);
    assert!(batch.done);
    assert!(
      batch
        .docs
        .iter()
        .any(|doc| doc.snapshot["meta"]["title"] == "entry" && doc.snapshot["type"] == "page")
    );
    assert!(batch.folders.iter().any(|folder| folder.path == "folder"));
  }

  #[test]
  fn markdown_cursor_drains_assets_after_final_doc_batch() {
    let bytes = zip(&[
      ("entry.md", b"![a](a.png)\n![b](b.png)"),
      ("a.png", &[1]),
      ("b.png", &[2]),
    ]);
    let mut cursor = markdown_zip_cursor(
      bytes,
      ImportOptions::default(),
      ImportBatchLimits {
        max_docs: 1,
        max_blobs: 1,
        max_blob_bytes: u64::MAX,
      },
    )
    .unwrap();

    let first = cursor.next_batch().unwrap().unwrap();
    let second = cursor.next_batch().unwrap().unwrap();

    assert_eq!(first.docs.len(), 1);
    assert_eq!(first.blobs.len(), 1);
    assert!(!first.done);
    assert_eq!(second.docs.len(), 0);
    assert_eq!(second.blobs.len(), 1);
    assert!(second.done);
    assert!(cursor.next_batch().unwrap().is_none());
  }

  #[test]
  fn cancel_stops_before_reading_source() {
    let bytes = zip(&[("entry.md", b"entry")]);
    let result = markdown_zip_cursor(
      bytes,
      ImportOptions {
        cancel: true,
        cancel_after_entries: None,
      },
      ImportBatchLimits::default(),
    );

    assert!(matches!(result, Err(ImportError::Cancelled)));
  }

  #[test]
  fn markdown_zip_resolves_links_assets_and_frontmatter() {
    let bytes = zip(&[
      (
        "entry.md",
        b"---\ntitle: Frontmatter Entry\ncreateDate: 2018-04-12T09:51:00.000Z\nupdatedDate: 2018-04-12T10:00:00.000Z\ntags: [a, b]\nfavorite: true\n---\naaa [go](./folder/target.md) ccc\n![logo](./assets/logo.png)\n\n| A | B |\n| --- | --- |\n| 1 | 2 |",
      ),
      ("folder/target.md", b"target"),
      ("assets/logo.png", &[137, 80, 78, 71]),
    ]);

    let batch = plan_markdown_zip(bytes).unwrap();
    let entry = batch
      .docs
      .iter()
      .find(|doc| doc.snapshot["meta"]["title"] == "Frontmatter Entry")
      .unwrap();
    let snapshot = serde_json::to_string(&entry.snapshot).unwrap();

    assert!(snapshot.contains("\"reference\""));
    assert!(snapshot.contains("\"insert\":\"aaa \""));
    assert!(snapshot.contains("\"title\":\"go\""));
    assert!(snapshot.contains("\"insert\":\" \""));
    assert!(snapshot.contains("\"insert\":\" ccc\""));
    assert!(snapshot.contains("\"sourceId\""));
    assert!(snapshot.contains("\"affine:table\""));
    assert_eq!(entry.meta.as_ref().unwrap().create_date, Some(1523526660000));
    assert_eq!(entry.meta.as_ref().unwrap().updated_date, Some(1523527200000));
    assert_eq!(
      entry.meta.as_ref().unwrap().tags.as_ref().unwrap(),
      &vec!["a".to_string(), "b".to_string()]
    );
    assert_eq!(entry.meta.as_ref().unwrap().favorite, Some(true));
  }

  #[test]
  fn markdown_zip_resolves_parent_directory_refs() {
    let bytes = zip(&[
      ("folder/entry.md", b"[root](../target.md)\n![logo](../assets/logo.png)"),
      ("target.md", b"target"),
      ("assets/logo.png", &[137, 80, 78, 71]),
    ]);

    let batch = plan_markdown_zip(bytes).unwrap();
    let entry = batch
      .docs
      .iter()
      .find(|doc| doc.snapshot["meta"]["title"] == "entry")
      .unwrap();
    let snapshot = serde_json::to_string(&entry.snapshot).unwrap();

    assert!(snapshot.contains("\"reference\""));
    assert!(snapshot.contains("\"sourceId\""));
  }

  #[test]
  fn markdown_zip_resolves_link_paths_with_parentheses() {
    let bytes = zip(&[
      (
        "Entry.md",
        b"[Linked item](Folder%20(Archive/Linked%20item%2011111111111111111111111111111111.md)",
      ),
      (
        "Folder (Archive/Linked item 11111111111111111111111111111111.md",
        b"target",
      ),
    ]);

    let batch = plan_markdown_zip(bytes).unwrap();
    let entry = batch
      .docs
      .iter()
      .find(|doc| doc.snapshot["meta"]["title"] == "Entry")
      .unwrap();
    let snapshot = serde_json::to_string(&entry.snapshot).unwrap();

    assert!(snapshot.contains("\"reference\""));
    assert!(snapshot.contains("\"insert\":\" \""));
    assert!(!snapshot.contains("\"insert\":\"Linked item\""));
    assert!(!snapshot.contains("Folder%20"));
  }

  #[test]
  fn markdown_zip_resolves_bare_relation_references() {
    let bytes = zip(&[
      (
        "Objectives/Entry.md",
        b"Related: Alpha task (../Targets/Alpha%20task%2011111111111111111111111111111111.md), Beta task (../Targets/Beta%20task%2022222222222222222222222222222222.md)",
      ),
      (
        "Targets/Alpha task 11111111111111111111111111111111.md",
        b"target",
      ),
      (
        "Targets/Beta task 22222222222222222222222222222222.md",
        b"target",
      ),
    ]);

    let batch = plan_markdown_zip(bytes).unwrap();
    let entry = batch
      .docs
      .iter()
      .find(|doc| doc.snapshot["meta"]["title"] == "Entry")
      .unwrap();
    let snapshot = serde_json::to_string(&entry.snapshot).unwrap();

    assert_eq!(snapshot.matches("\"reference\"").count(), 2);
    assert!(!snapshot.contains("../Targets"));
  }

  #[test]
  fn markdown_zip_resolves_asset_paths_with_parentheses() {
    let bytes = zip(&[
      ("Export/Page.md", b"Before ![cover](Page/image%20(1)%20abc.png) after"),
      ("Export/Page/image (1) abc.png", &[137, 80, 78, 71]),
    ]);

    let batch = plan_markdown_zip(bytes).unwrap();
    let snapshot = serde_json::to_string(&batch.docs[0].snapshot).unwrap();

    assert_eq!(batch.blobs.len(), 1);
    assert!(snapshot.contains("\"sourceId\""));
  }

  #[test]
  fn markdown_zip_skips_too_large_doc_with_warning() {
    let markdown = "a".repeat(1_001_000);
    let bytes = zip(&[("large.md", markdown.as_bytes())]);

    let batch = plan_markdown_zip(bytes).unwrap();

    assert!(batch.docs.is_empty());
    assert_eq!(batch.warnings.len(), 1);
    assert_eq!(batch.warnings[0].code, "skipped_doc");
    assert_eq!(batch.warnings[0].source_path.as_deref(), Some("large.md"));
  }

  #[test]
  fn markdown_zip_skips_too_large_asset_with_warning() {
    let bytes = zip(&[("entry.md", b"![logo](logo.png)"), ("logo.png", &[1, 2, 3, 4])]);
    let mut cursor = markdown_zip_cursor(
      bytes,
      ImportOptions::default(),
      ImportBatchLimits {
        max_docs: usize::MAX,
        max_blobs: usize::MAX,
        max_blob_bytes: 2,
      },
    )
    .unwrap();
    let mut merged = empty_batch(0);
    while let Some(batch) = cursor.next_batch().unwrap() {
      merge_batch(&mut merged, batch);
    }

    assert!(merged.blobs.is_empty());
    assert_eq!(merged.warnings.len(), 1);
    assert_eq!(merged.warnings[0].code, "skipped_asset");
    assert_eq!(merged.warnings[0].source_path.as_deref(), Some("logo.png"));
  }

  #[test]
  fn markdown_zip_zero_blob_limit_still_drains_assets() {
    let bytes = zip(&[("entry.md", b"![logo](logo.png)"), ("logo.png", &[1, 2, 3, 4])]);
    let mut cursor = markdown_zip_cursor(
      bytes,
      ImportOptions::default(),
      ImportBatchLimits {
        max_docs: 1,
        max_blobs: 0,
        max_blob_bytes: u64::MAX,
      },
    )
    .unwrap();
    let mut merged = empty_batch(0);
    let mut batch_count = 0;
    let mut done = false;
    while let Some(batch) = cursor.next_batch().unwrap() {
      batch_count += 1;
      done = batch.done;
      merge_batch(&mut merged, batch);
      assert!(batch_count <= 3);
    }

    assert_eq!(merged.docs.len(), 1);
    assert_eq!(merged.blobs.len(), 1);
    assert!(done);
  }

  #[test]
  fn cancel_stops_during_planning() {
    let bytes = zip(&[("a.md", b"a"), ("b.md", b"b")]);
    let mut cursor = markdown_zip_cursor(
      bytes,
      ImportOptions {
        cancel: false,
        cancel_after_entries: Some(1),
      },
      ImportBatchLimits {
        max_docs: 1,
        max_blobs: usize::MAX,
        max_blob_bytes: u64::MAX,
      },
    )
    .unwrap();

    assert!(cursor.next_batch().unwrap().is_some());
    assert!(matches!(cursor.next_batch(), Err(ImportError::Cancelled)));
  }
}