pinakes/crates/pinakes-ui/src/components/markdown_viewer.rs

use dioxus::{document::eval, prelude::*};

/// Event handler for wikilink clicks. Called with the target note name.
pub type WikilinkClickHandler = EventHandler<String>;

#[component]
pub fn MarkdownViewer(
  content_url: String,
  media_type: String,
  #[props(default)] on_wikilink_click: Option<WikilinkClickHandler>,
) -> Element {
  let mut rendered_html = use_signal(String::new);
  let mut frontmatter_html = use_signal(|| Option::<String>::None);
  let mut raw_content = use_signal(String::new);
  let mut loading = use_signal(|| true);
  let mut error = use_signal(|| Option::<String>::None);
  let mut show_preview = use_signal(|| true);

  // Fetch content on mount
  let url = content_url.clone();
  let mtype = media_type.clone();
  use_effect(move || {
    let url = url.clone();
    let mtype = mtype.clone();
    spawn(async move {
      loading.set(true);
      error.set(None);
      match reqwest::get(&url).await {
        Ok(resp) => {
          match resp.text().await {
            Ok(text) => {
              raw_content.set(text.clone());
              if mtype == "md" || mtype == "markdown" {
                let (fm_html, body_html) =
                  render_markdown_with_frontmatter(&text);
                frontmatter_html.set(fm_html);
                rendered_html.set(body_html);
              } else {
                frontmatter_html.set(None);
                rendered_html.set(render_plaintext(&text));
              };
            },
            Err(e) => error.set(Some(format!("Failed to read content: {e}"))),
          }
        },
        Err(e) => error.set(Some(format!("Failed to fetch: {e}"))),
      }
      loading.set(false);
    });
  });

  // Set up global wikilink click handler that the inline onclick attributes
  // will call This bridges JavaScript -> Rust communication
  use_effect(move || {
    if let Some(handler) = on_wikilink_click {
      spawn(async move {
        // Set up a global function that wikilink onclick handlers can call
        // The function stores the clicked target in localStorage
        let setup_js = r#"
                    window.__dioxus_wikilink_click = function(target) {
                        console.log('Wikilink clicked:', target);
                        localStorage.setItem('__wikilink_clicked', target);
                    };
                "#;

        let _ = eval(setup_js).await;

        // Poll localStorage to detect wikilink clicks
        loop {
          tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

          let check_js = r#"
                        (function() {
                            const target = localStorage.getItem('__wikilink_clicked');
                            if (target) {
                                localStorage.removeItem('__wikilink_clicked');
                                return target;
                            }
                            return '';
                        })();
                    "#;

          if let Ok(result) = eval(check_js).await
            && let Some(target) = result.as_str()
            && !target.is_empty()
          {
            handler.call(target.to_string());
          }
        }
      });
    }
  });

  let is_loading = *loading.read();
  let is_preview = *show_preview.read();

  rsx! {
      div { class: "markdown-viewer",
          // View toggle toolbar
          if !is_loading && error.read().is_none() {
              div { class: "markdown-toolbar",
                  button {
                      class: if is_preview { "toolbar-btn active" } else { "toolbar-btn" },
                      onclick: move |_| show_preview.set(true),
                      title: "Preview Mode",
                      "Preview"
                  }
                  button {
                      class: if !is_preview { "toolbar-btn active" } else { "toolbar-btn" },
                      onclick: move |_| show_preview.set(false),
                      title: "Source Mode",
                      "Source"
                  }
              }
          }

          if is_loading {
              div { class: "loading-overlay",
                  div { class: "spinner" }
                  "Loading content..."
              }
          }

          if let Some(ref err) = *error.read() {
              div { class: "error-banner",
                  span { class: "error-icon", "\u{26a0}" }
                  "{err}"
              }
          }

          if !is_loading && error.read().is_none() {
              if is_preview {
                  // Preview mode - show rendered markdown
                  if let Some(ref fm) = *frontmatter_html.read() {
                      div {
                          class: "frontmatter-card",
                          dangerous_inner_html: "{fm}",
                      }
                  }
                  div {
                      class: "markdown-content",
                      dangerous_inner_html: "{rendered_html}",
                  }
              } else {
                  // Source mode - show raw markdown
                  pre { class: "markdown-source",
                      code { "{raw_content}" }
                  }
              }
          }
      }
  }
}

/// Parse frontmatter and render markdown body. Returns (frontmatter_html,
/// body_html).
fn render_markdown_with_frontmatter(text: &str) -> (Option<String>, String) {
  use gray_matter::{Matter, engine::YAML};

  let matter = Matter::<YAML>::new();
  let Ok(result) = matter.parse(text) else {
    // If frontmatter parsing fails, just render the whole text as markdown
    return (None, render_markdown(text));
  };

  let fm_html = result.data.and_then(|data| render_frontmatter_card(&data));

  let body_html = render_markdown(&result.content);
  (fm_html, body_html)
}

/// Render frontmatter fields as an HTML card.
fn render_frontmatter_card(data: &gray_matter::Pod) -> Option<String> {
  let gray_matter::Pod::Hash(map) = data else {
    return None;
  };

  if map.is_empty() {
    return None;
  }

  let mut html = String::from("<dl class=\"frontmatter-fields\">");

  for (key, value) in map {
    let display_value = pod_to_display(value);
    let escaped_key = escape_html(key);
    html.push_str(&format!("<dt>{escaped_key}</dt><dd>{display_value}</dd>"));
  }

  html.push_str("</dl>");
  Some(html)
}

fn pod_to_display(pod: &gray_matter::Pod) -> String {
  match pod {
    gray_matter::Pod::String(s) => escape_html(s),
    gray_matter::Pod::Integer(n) => n.to_string(),
    gray_matter::Pod::Float(f) => f.to_string(),
    gray_matter::Pod::Boolean(b) => b.to_string(),
    gray_matter::Pod::Array(arr) => {
      let items: Vec<String> = arr.iter().map(pod_to_display).collect();
      items.join(", ")
    },
    gray_matter::Pod::Hash(map) => {
      let items: Vec<String> = map
        .iter()
        .map(|(k, v)| format!("{}: {}", escape_html(k), pod_to_display(v)))
        .collect();
      items.join("; ")
    },
    gray_matter::Pod::Null => String::new(),
  }
}

fn render_markdown(text: &str) -> String {
  use pulldown_cmark::{Options, Parser, html};

  // First, convert wikilinks to standard markdown links
  let text_with_links = convert_wikilinks(text);

  let mut options = Options::empty();
  options.insert(Options::ENABLE_TABLES);
  options.insert(Options::ENABLE_STRIKETHROUGH);
  options.insert(Options::ENABLE_TASKLISTS);
  options.insert(Options::ENABLE_FOOTNOTES);
  options.insert(Options::ENABLE_HEADING_ATTRIBUTES);

  let parser = Parser::new_ext(&text_with_links, options);
  let mut html_output = String::new();
  html::push_html(&mut html_output, parser);

  // Sanitize HTML using ammonia with a safe allowlist
  sanitize_html(&html_output)
}

/// Convert wikilinks [[target]] and [[target|display]] to styled HTML links.
/// Uses a special URL scheme that can be intercepted by click handlers.
fn convert_wikilinks(text: &str) -> String {
  use regex::Regex;

  // Match embeds ![[target]] first, convert to a placeholder image/embed span
  let embed_re = Regex::new(r"!\[\[([^\]|]+)(?:\|([^\]]+))?\]\]").unwrap();
  let text = embed_re.replace_all(text, |caps: &regex::Captures| {
    let target = caps.get(1).unwrap().as_str().trim();
    let alt = caps.get(2).map(|m| m.as_str().trim()).unwrap_or(target);
    format!(
      "<span class=\"wikilink-embed\" data-target=\"{}\" title=\"Embed: \
       {}\">[Embed: {}]</span>",
      escape_html_attr(target),
      escape_html_attr(target),
      escape_html(alt)
    )
  });

  // Match wikilinks [[target]] or [[target|display]]
  let wikilink_re = Regex::new(r"\[\[([^\]|]+)(?:\|([^\]]+))?\]\]").unwrap();
  let text = wikilink_re.replace_all(&text, |caps: &regex::Captures| {
    let target = caps.get(1).unwrap().as_str().trim();
    let display = caps.get(2).map(|m| m.as_str().trim()).unwrap_or(target);
    // Create a styled link that uses a special pseudo-protocol scheme
    // This makes it easier to intercept clicks via JavaScript
    format!(
      "<a href=\"javascript:void(0)\" class=\"wikilink\" \
       data-wikilink-target=\"{target}\" \
       onclick=\"if(window.__dioxus_wikilink_click){{window.\
       __dioxus_wikilink_click('{target_escaped}')}}\">{display}</a>",
      target = escape_html_attr(target),
      target_escaped =
        escape_html_attr(&target.replace('\\', "\\\\").replace('\'', "\\'")),
      display = escape_html(display)
    )
  });

  text.to_string()
}

fn render_plaintext(text: &str) -> String {
  let escaped = escape_html(text);
  format!("<pre><code>{escaped}</code></pre>")
}

/// Escape text for display in HTML content.
fn escape_html(text: &str) -> String {
  text
    .replace('&', "&amp;")
    .replace('<', "&lt;")
    .replace('>', "&gt;")
    .replace('"', "&quot;")
}

/// Escape text for use in HTML attributes (includes single quotes).
fn escape_html_attr(text: &str) -> String {
  text
    .replace('&', "&amp;")
    .replace('<', "&lt;")
    .replace('>', "&gt;")
    .replace('"', "&quot;")
    .replace('\'', "&#x27;")
}

/// Sanitize HTML using ammonia with a safe allowlist.
/// This prevents XSS attacks by removing dangerous elements and attributes.
fn sanitize_html(html: &str) -> String {
  use std::collections::HashSet;

  use ammonia::Builder;

  // Build a custom sanitizer that allows safe markdown elements
  // but strips all event handlers and dangerous elements
  let mut builder = Builder::default();

  // Allow common markdown elements
  let allowed_tags: HashSet<&str> = [
    "a",
    "abbr",
    "acronym",
    "b",
    "blockquote",
    "br",
    "code",
    "dd",
    "del",
    "details",
    "div",
    "dl",
    "dt",
    "em",
    "h1",
    "h2",
    "h3",
    "h4",
    "h5",
    "h6",
    "hr",
    "i",
    "img",
    "ins",
    "kbd",
    "li",
    "mark",
    "ol",
    "p",
    "pre",
    "q",
    "s",
    "samp",
    "small",
    "span",
    "strong",
    "sub",
    "summary",
    "sup",
    "table",
    "tbody",
    "td",
    "tfoot",
    "th",
    "thead",
    "tr",
    "u",
    "ul",
    "var",
    // Task list support
    "input",
  ]
  .into_iter()
  .collect();

  // Allow safe attributes
  let allowed_attrs: HashSet<&str> = [
    "href",
    "src",
    "alt",
    "title",
    "class",
    "id",
    "name",
    "width",
    "height",
    "align",
    "valign",
    "colspan",
    "rowspan",
    "scope",
    // Data attributes for wikilinks (safe - no code execution)
    "data-target",
    "data-wikilink-target",
    // Task list checkbox support
    "type",
    "checked",
    "disabled",
  ]
  .into_iter()
  .collect();

  builder
        .tags(allowed_tags)
        .generic_attributes(allowed_attrs)
        // Allow relative URLs and fragment-only URLs for internal links
        .url_schemes(["http", "https", "mailto"].into_iter().collect())
        .link_rel(Some("noopener noreferrer"))
        // Strip all event handler attributes (onclick, onerror, etc.)
        .strip_comments(true)
        .clean(html)
        .to_string()
}