affine_importer 0.1.2

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

use pulldown_cmark::{Event, Parser, Tag, TagEnd};
use serde_json::{Map as JsonMap, Value as JsonValue};

use super::{IndexedEntry, file_name, mime_from_path, strip_extension};
use crate::vfs::normalize_import_path;

pub(super) fn rewrite_asset_references(
  markdown: &str,
  current_path: &str,
  blob_ids_by_path: &BTreeMap<String, String>,
) -> String {
  rewrite_markdown_lines(markdown, |line, protected| {
    rewrite_line_asset_references(line, protected, current_path, blob_ids_by_path)
  })
}

pub(super) fn rewrite_markdown_link_destinations(
  markdown: &str,
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
) -> String {
  rewrite_markdown_lines(markdown, |line, protected| {
    rewrite_line_link_destinations(line, protected, current_path, page_ids_by_path)
  })
}

fn rewrite_line_link_destinations(
  line: &str,
  protected: &[Range<usize>],
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
) -> String {
  let mut output = String::with_capacity(line.len());
  let mut rest = line;
  loop {
    let Some(start) = rest.find('[') else {
      output.push_str(rest);
      break;
    };
    output.push_str(&rest[..start]);
    let position = line.len() - rest.len() + start;
    if is_protected(protected, position) {
      output.push('[');
      rest = &rest[start + 1..];
      continue;
    }
    if start > 0 && rest.as_bytes().get(start - 1) == Some(&b'!') {
      output.push('[');
      rest = &rest[start + 1..];
      continue;
    }
    let link = &rest[start..];
    let Some((label, target, consumed)) = parse_link_span(link) else {
      output.push('[');
      rest = &rest[start + 1..];
      continue;
    };
    if page_ids_by_path.contains_key(&resolve_path(current_path, &target)) {
      output.push_str(&format!("[{label}](<{target}>)"));
    } else {
      output.push_str(&link[..consumed]);
    }
    rest = &link[consumed..];
  }
  output
}

pub(super) fn rewrite_bare_page_references(
  markdown: &str,
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
) -> String {
  rewrite_markdown_lines(markdown, |line, protected| {
    rewrite_line_bare_page_references(line, protected, current_path, page_ids_by_path)
  })
}

pub(super) fn rewrite_markdown_lines<F>(markdown: &str, mut rewrite_line: F) -> String
where
  F: FnMut(&str, &[Range<usize>]) -> String,
{
  let protected = markdown_code_ranges(markdown);
  let mut output = String::with_capacity(markdown.len());
  let mut line_start = 0;
  let mut range_start = 0;
  for segment in markdown.split_inclusive('\n') {
    let (line, newline) = split_newline(segment);
    let line_end = line_start + line.len();
    while protected.get(range_start).is_some_and(|range| range.end <= line_start) {
      range_start += 1;
    }
    let line_ranges = protected[range_start..]
      .iter()
      .take_while(|range| range.start < line_end)
      .map(|range| range.start.saturating_sub(line_start)..range.end.min(line_end) - line_start)
      .collect::<Vec<_>>();
    output.push_str(&rewrite_line(line, &line_ranges));
    output.push_str(newline);
    line_start += segment.len();
  }
  output
}

fn markdown_code_ranges(markdown: &str) -> Vec<Range<usize>> {
  let mut ranges = Vec::new();
  let mut code_block_start = None;
  for (event, range) in Parser::new(markdown).into_offset_iter() {
    match event {
      Event::Start(Tag::CodeBlock(_)) => code_block_start = Some(range.start),
      Event::End(TagEnd::CodeBlock) => {
        if let Some(start) = code_block_start.take() {
          ranges.push(start..range.end);
        }
      }
      Event::Code(_) if code_block_start.is_none() => ranges.push(range),
      _ => {}
    }
  }
  if let Some(start) = code_block_start {
    ranges.push(start..markdown.len());
  }
  ranges
}

fn split_newline(segment: &str) -> (&str, &str) {
  segment.strip_suffix('\n').map_or((segment, ""), |line| (line, "\n"))
}

pub(super) fn is_protected(ranges: &[Range<usize>], position: usize) -> bool {
  ranges.iter().any(|range| range.contains(&position))
}

fn rewrite_line_bare_page_references(
  line: &str,
  protected: &[Range<usize>],
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
) -> String {
  let mut output = String::with_capacity(line.len());
  let mut rest = line;
  while let Some(open) = rest.find('(') {
    let position = line.len() - rest.len() + open;
    let Some(close) = rest[open + 1..].find(')').map(|index| open + 1 + index) else {
      break;
    };
    if is_protected(protected, position) {
      output.push_str(&rest[..=open]);
      rest = &rest[open + 1..];
      continue;
    }
    let target = &rest[open + 1..close];
    if !target.to_ascii_lowercase().ends_with(".md")
      || !page_ids_by_path.contains_key(&resolve_path(current_path, target))
    {
      output.push_str(&rest[..=close]);
      rest = &rest[close + 1..];
      continue;
    }
    if open > 0 && rest.as_bytes().get(open - 1) == Some(&b']') {
      output.push_str(&rest[..=close]);
      rest = &rest[close + 1..];
      continue;
    }
    let label_start = bare_reference_label_start(&rest[..open]);
    let label = rest[label_start..open].trim();
    if label.is_empty() {
      output.push_str(&rest[..=close]);
      rest = &rest[close + 1..];
      continue;
    }
    output.push_str(&rest[..label_start]);
    output.push_str(&format!("[{label}](<{target}>)"));
    rest = &rest[close + 1..];
  }
  output.push_str(rest);
  output
}

fn bare_reference_label_start(prefix: &str) -> usize {
  [", ", ": "]
    .iter()
    .filter_map(|delimiter| prefix.rfind(delimiter).map(|index| index + delimiter.len()))
    .max()
    .unwrap_or(0)
}

fn rewrite_line_asset_references(
  line: &str,
  protected: &[Range<usize>],
  current_path: &str,
  blob_ids_by_path: &BTreeMap<String, String>,
) -> String {
  let mut output = String::with_capacity(line.len());
  let mut rest = line;
  loop {
    let Some(start) = rest.find("![") else {
      output.push_str(rest);
      break;
    };
    output.push_str(&rest[..start]);
    let position = line.len() - rest.len() + start;
    if is_protected(protected, position) {
      output.push_str("![");
      rest = &rest[start + 2..];
      continue;
    }
    let image = &rest[start..];
    let Some((alt, target, consumed)) = parse_image_span(image) else {
      output.push_str(&rest[start..start + 2]);
      rest = &rest[start + 2..];
      continue;
    };
    if let Some(source_id) = blob_ids_by_path.get(&resolve_path(current_path, &target)) {
      output.push_str(&format!("![{alt}](blob://{source_id})"));
    } else {
      output.push_str(&image[..consumed]);
    }
    rest = &image[consumed..];
  }
  output
}

fn parse_image_span(value: &str) -> Option<(String, String, usize)> {
  if value.starts_with("![") {
    return parse_link_like_span(value, 2);
  }
  None
}

pub(super) fn parse_link_span(value: &str) -> Option<(String, String, usize)> {
  if value.starts_with('[') {
    return parse_link_like_span(value, 1);
  }
  None
}

fn parse_link_like_span(value: &str, label_start: usize) -> Option<(String, String, usize)> {
  let bytes = value.as_bytes();
  let mut index = label_start;
  let mut escaped = false;
  while index < bytes.len() {
    let byte = bytes[index];
    if escaped {
      escaped = false;
    } else if byte == b'\\' {
      escaped = true;
    } else if byte == b']' {
      break;
    }
    index += 1;
  }
  if index >= bytes.len() || bytes.get(index + 1) != Some(&b'(') {
    return None;
  }
  let label = value[label_start..index].to_string();
  index += 2;
  let target_start = index;
  let mut depth = 0usize;
  let mut escaped = false;
  let mut last_close = None;
  while index < bytes.len() {
    let byte = bytes[index];
    if escaped {
      escaped = false;
    } else if byte == b'\\' {
      escaped = true;
    } else if byte == b'(' {
      depth += 1;
    } else if byte == b')' {
      last_close = Some(index);
      if depth == 0 {
        let target = value[target_start..index].to_string();
        return Some((label, target, index + 1));
      }
      depth -= 1;
    }
    index += 1;
  }
  let close = last_close?;
  let target = value[target_start..close].to_string();
  Some((label, target, close + 1))
}

pub(super) fn rewrite_page_references(
  value: &mut JsonValue,
  current_path: &str,
  page_ids_by_path: &BTreeMap<String, String>,
) {
  rewrite_link_references(value, |target, title| {
    let resolved = resolve_path(current_path, target);
    page_ids_by_path
      .get(&resolved)
      .map(|page_id| linked_page_reference(page_id, Some(title), None))
  });
}

pub(super) fn rewrite_link_references(value: &mut JsonValue, mut resolve: impl FnMut(&str, &str) -> Option<JsonValue>) {
  rewrite_link_references_inner(value, &mut resolve);
}

fn rewrite_link_references_inner(value: &mut JsonValue, resolve: &mut impl FnMut(&str, &str) -> Option<JsonValue>) {
  match value {
    JsonValue::Object(map) => {
      rewrite_link_reference(map, resolve);
      for value in map.values_mut() {
        rewrite_link_references_inner(value, resolve);
      }
    }
    JsonValue::Array(values) => {
      for value in values {
        rewrite_link_references_inner(value, resolve);
      }
    }
    _ => {}
  }
}

pub(super) fn linked_page_reference(page_id: &str, title: Option<&str>, block_id: Option<&str>) -> JsonValue {
  let mut reference = serde_json::json!({
    "type": "LinkedPage",
    "pageId": page_id,
  });
  if let Some(object) = reference.as_object_mut() {
    if let Some(title) = title {
      object.insert("title".to_string(), JsonValue::String(title.to_string()));
    }
    if let Some(block_id) = block_id {
      object.insert(
        "params".to_string(),
        serde_json::json!({ "mode": "page", "blockIds": [block_id] }),
      );
    }
  }
  reference
}

pub(super) struct ImportedAssetMetadata {
  pub(super) file_name: String,
  pub(super) mime: String,
  pub(super) size: u64,
}

pub(super) fn imported_asset_metadata(entry: &IndexedEntry) -> ImportedAssetMetadata {
  ImportedAssetMetadata {
    file_name: file_name(&entry.meta.path).to_string(),
    mime: mime_from_path(&entry.meta.path).to_string(),
    size: entry.meta.uncompressed_size,
  }
}

pub(super) fn rewrite_non_image_embeds_to_attachments(
  value: &mut JsonValue,
  blob_by_id: &BTreeMap<String, ImportedAssetMetadata>,
) {
  rewrite_attachment_blocks(value, blob_by_id);
}

fn rewrite_attachment_blocks(value: &mut JsonValue, blob_by_id: &BTreeMap<String, ImportedAssetMetadata>) {
  match value {
    JsonValue::Object(map) => {
      let attachment = map
        .get("flavour")
        .and_then(JsonValue::as_str)
        .filter(|flavour| *flavour == "affine:image")
        .and_then(|_| map.get("props"))
        .and_then(JsonValue::as_object)
        .and_then(|props| props.get("sourceId"))
        .and_then(JsonValue::as_str)
        .and_then(|source_id| blob_by_id.get(source_id))
        .filter(|blob| !blob.mime.starts_with("image/"));

      if let Some(blob) = attachment {
        map.insert(
          "flavour".to_string(),
          JsonValue::String("affine:attachment".to_string()),
        );
        if let Some(props) = map.get_mut("props").and_then(JsonValue::as_object_mut) {
          props.insert("name".to_string(), JsonValue::String(blob.file_name.clone()));
          props.insert("size".to_string(), JsonValue::Number(blob.size.into()));
          props.insert("type".to_string(), JsonValue::String(blob.mime.clone()));
          props.insert("embed".to_string(), JsonValue::Bool(false));
          props.insert("style".to_string(), JsonValue::String("horizontalThin".to_string()));
          props.insert("footnoteIdentifier".to_string(), JsonValue::Null);
        }
      }
      for value in map.values_mut() {
        rewrite_attachment_blocks(value, blob_by_id);
      }
    }
    JsonValue::Array(values) => {
      for value in values {
        rewrite_attachment_blocks(value, blob_by_id);
      }
    }
    _ => {}
  }
}

fn rewrite_link_reference(
  map: &mut JsonMap<String, JsonValue>,
  resolve: &mut impl FnMut(&str, &str) -> Option<JsonValue>,
) {
  let title = map.get("insert").and_then(JsonValue::as_str).unwrap_or("").to_string();
  let Some(target) = map
    .get("attributes")
    .and_then(JsonValue::as_object)
    .and_then(|attributes| attributes.get("link"))
    .and_then(JsonValue::as_str)
  else {
    return;
  };
  let Some(reference) = resolve(target, &title) else {
    return;
  };
  let Some(attributes) = map.get_mut("attributes").and_then(JsonValue::as_object_mut) else {
    return;
  };
  attributes.remove("link");
  attributes.insert("reference".to_string(), reference);
  map.insert("insert".to_string(), JsonValue::String(" ".to_string()));
}

pub(super) fn resolve_path(current_path: &str, target: &str) -> String {
  let target = target.trim().trim_start_matches('<').trim_end_matches('>');
  if target.contains("://") {
    return target.to_string();
  }
  let target = percent_decode(target.split(['?', '#']).next().unwrap_or(target));
  let base = current_path.rsplit_once('/').map(|(base, _)| base).unwrap_or("");
  normalize_import_path(&format!("{base}/{target}"))
}

pub(super) fn register_page_path(page_ids_by_path: &mut BTreeMap<String, String>, path: &str, doc_id: &str) {
  let path = normalize_import_path(path);
  page_ids_by_path.insert(path.clone(), doc_id.to_string());
  page_ids_by_path.insert(strip_extension(&path).to_string(), doc_id.to_string());
}

pub(super) fn register_csv_path(csv_entries_by_path: &mut BTreeMap<String, usize>, path: &str, entry_index: usize) {
  let path = normalize_import_path(path);
  csv_entries_by_path.insert(path, entry_index);
}

fn percent_decode(value: &str) -> String {
  let bytes = value.as_bytes();
  let mut output = Vec::with_capacity(bytes.len());
  let mut index = 0;
  while index < bytes.len() {
    if bytes[index] == b'%'
      && index + 2 < bytes.len()
      && let (Some(high), Some(low)) = (hex_value(bytes[index + 1]), hex_value(bytes[index + 2]))
    {
      output.push(high * 16 + low);
      index += 3;
      continue;
    }
    output.push(bytes[index]);
    index += 1;
  }
  String::from_utf8(output).unwrap_or_else(|_| value.to_string())
}

fn hex_value(byte: u8) -> Option<u8> {
  match byte {
    b'0'..=b'9' => Some(byte - b'0'),
    b'a'..=b'f' => Some(byte - b'a' + 10),
    b'A'..=b'F' => Some(byte - b'A' + 10),
    _ => None,
  }
}