MarkdownUtils.fs

// --------------------------------------------------------------------------------------
// F# Markdown
// --------------------------------------------------------------------------------------

namespace rec FSharp.Formatting.Markdown

open System
open System.Collections.Generic
open System.Xml.Linq
open FSharp.Formatting.Templating

/// Internal utilities shared across all Markdown formatting back-ends (HTML, Latex, Markdown, .fsx, .ipynb).
/// Provides helpers for inspecting and transforming the <see cref="T:FSharp.Formatting.Markdown.MarkdownParagraph"/>
/// AST, applying template substitutions, and resolving cross-references.
module internal MarkdownUtils =
    /// Returns true when the paragraph is a code block or inline HTML block.
    let isCode =
        (function
        | CodeBlock _
        | InlineHtmlBlock _ -> true
        | _ -> false)

    /// Returns true when the paragraph is an output block (e.g. a notebook cell output).
    let isCodeOutput =
        (function
        | OutputBlock _ -> true
        | _ -> false)

    /// Extracts the optional execution count from a code or inline-HTML block; returns None for other paragraphs.
    let getExecutionCount =
        (function
        | CodeBlock(executionCount = executionCount)
        | InlineHtmlBlock(executionCount = executionCount) -> executionCount
        | _ -> None)

    /// Extracts the source-code string from a code or inline-HTML block; raises if called on other paragraphs.
    let getCode =
        (function
        | CodeBlock(code = code) -> code
        | InlineHtmlBlock(code = code) -> code
        | _ -> failwith "unreachable")

    /// Extracts the output text and MIME kind from an output block; raises if called on other paragraphs.
    let getCodeOutput =
        (function
        | OutputBlock(code, kind, _) -> code, kind
        | _ -> failwith "unreachable")

    /// Splits a paragraph list at the first code block, returning the leading code (or prose) section
    /// paired with the remaining paragraphs. Used by back-ends that process documents cell-by-cell.
    let splitParagraphs paragraphs =
        let firstCode = paragraphs |> List.tryFindIndex isCode

        match firstCode with
        | Some 0 ->
            let code = paragraphs.[0]
            let codeLines = getCode code
            let otherParagraphs = paragraphs.[1..]

            // Collect the code output(s) that follows this cell if any
            let codeOutput = otherParagraphs |> List.takeWhile isCodeOutput |> List.map getCodeOutput

            let otherParagraphs = otherParagraphs |> List.skipWhile isCodeOutput

            Choice1Of2(codeLines, codeOutput, getExecutionCount code), otherParagraphs

        | Some _
        | None ->
            let markdownParagraphs = paragraphs |> List.takeWhile (isCode >> not)

            let otherParagraphs = paragraphs |> List.skipWhile (isCode >> not)

            Choice2Of2 markdownParagraphs, otherParagraphs

    /// Lookup a specified key in a dictionary, possibly
    /// ignoring newlines or spaces in the key.
    let (|LookupKey|_|) (dict: IDictionary<_, _>) (key: string) =
        [ key; key.Replace("\r\n", ""); key.Replace("\r\n", " "); key.Replace("\n", ""); key.Replace("\n", " ") ]
        |> List.tryPick (fun key ->
            match dict.TryGetValue(key) with
            | true, v -> Some v
            | _ -> None)

    /// Context passed around while formatting
    type FormattingContext =
        {
            Links: IDictionary<string, string * string option>
            Newline: string
            /// Additional replacements to be made in content
            Substitutions: Substitutions
            /// Helper to resolve `cref:T:TypeName` references in markdown
            CodeReferenceResolver: string -> (string * string) option
            /// Helper to resolve `[foo](file.md)` references in markdown (where file.md is producing file.fsx)
            MarkdownDirectLinkResolver: string -> string option
            DefineSymbol: string
        }

    /// Format a MarkdownSpan
    let rec formatSpan (ctx: FormattingContext) span =
        match span with
        | LatexInlineMath(body, _) -> sprintf "$%s$" body
        | LatexDisplayMath(body, _) -> sprintf "$$%s$$" body
        | EmbedSpans(cmd, _) -> formatSpans ctx (cmd.Render())
        | Literal(str, _) -> str
        | HardLineBreak(_) -> "  " + ctx.Newline

        | AnchorLink _ -> ""
        | DirectLink(body, link, title, _) ->
            let t =
                title
                |> Option.map (fun t -> sprintf " \"%s\"" (t.Replace("\"", "\\\"")))
                |> Option.defaultValue ""

            "[" + formatSpans ctx body + "](" + link + t + ")"

        | IndirectLink(body, _, LookupKey ctx.Links (link, _), _)
        | IndirectLink(body, link, _, _) -> "[" + formatSpans ctx body + "](" + link + ")"

        | IndirectImage(body, _, LookupKey ctx.Links (link, _), _) -> sprintf "![%s](%s)" body link
        | IndirectImage(body, _, key, _) -> sprintf "![%s][%s]" body key

        | DirectImage(body, link, title, _) ->
            let t =
                title
                |> Option.map (fun t -> sprintf " \"%s\"" (t.Replace("\"", "\\\"")))
                |> Option.defaultValue ""

            sprintf "![%s](%s)" body (link + t)
        | Strong(body, _) -> "**" + formatSpans ctx body + "**"
        | InlineCode(body, _) ->
            // Pick the shortest backtick fence that does not appear in the body.
            // E.g. body "``h``" needs a triple-backtick fence; body "a`b" needs double.
            let maxConsecutiveBackticks =
                body
                |> Seq.fold
                    (fun (maxR, run) c ->
                        if c = '`' then
                            let run' = run + 1
                            (max maxR run'), run'
                        else
                            maxR, 0)
                    (0, 0)
                |> fst

            let fence = String.replicate (maxConsecutiveBackticks + 1) "`"
            // Surround with spaces when the body starts or ends with a backtick so the
            // fence and content do not merge (e.g. `` ``h`` `` would look like 4-backtick).
            if body.Length > 0 && (body.[0] = '`' || body.[body.Length - 1] = '`') then
                fence + " " + body + " " + fence
            else
                fence + body + fence
        | Emphasis(body, _) -> "*" + formatSpans ctx body + "*"

    /// Format a list of MarkdownSpan
    and formatSpans ctx spans =
        spans |> List.map (formatSpan ctx) |> String.concat ""

    /// Format a MarkdownParagraph
    let rec formatParagraph (ctx: FormattingContext) paragraph =
        // Shared helper for both ordered and unordered list blocks.
        // getPrefix receives the 0-based item index and returns the leading string (e.g. "* " or "1. ").
        let formatListBlock paragraphsl (getPrefix: int -> string) =
            let isTight =
                paragraphsl
                |> List.forall (function
                    | [ Span _ ] -> true
                    | _ -> false)

            [ for (n, paragraphs) in List.indexed paragraphsl do
                  for (i, paragraph) in List.indexed paragraphs do
                      let lines: string list = formatParagraph ctx paragraph
                      let lines = if lines.IsEmpty then [ "" ] else lines

                      for (j, line) in List.indexed lines do
                          if i = 0 && j = 0 then
                              yield getPrefix n + line
                          else
                              yield "  " + line

                  if not isTight then
                      yield ""

              if isTight then
                  yield "" ]

        [ match paragraph with
          | LatexBlock(env, lines, _) ->
              // Single-line equation blocks are rendered with the compact $$...$$ notation
              // (which is also valid markdown and what most authors write). Multi-line or
              // non-standard environments keep the \begin{env}...\end{env} form.
              if env = "equation" && lines.Length = 1 then
                  yield sprintf "$$%s$$" lines.[0]
              else
                  yield sprintf "\\begin{%s}" env

                  for line in lines do
                      yield line

                  yield sprintf "\\end{%s}" env

              yield ""

          | Heading(n, spans, _) ->
              yield String.replicate n "#" + " " + formatSpans ctx spans

              yield ""
          | Paragraph(spans, _) ->
              yield String.concat "" [ for span in spans -> formatSpan ctx span ]
              yield ""

          | HorizontalRule(c, _) ->
              yield String.replicate 3 (string c)
              yield ""
          | CodeBlock(code = code; fence = fence; language = language) ->
              // Indented code blocks (fence = None) are serialised as fenced blocks so
              // that the round-trip is valid — raw indented code without a '> ' prefix
              // or 4-space indent would be parsed as a paragraph, not a code block.
              let f = defaultArg fence "```"
              yield f + language

              yield code

              yield f

              yield ""
          | ListBlock(Unordered, paragraphsl, _) -> yield! formatListBlock paragraphsl (fun _ -> "* ")
          | ListBlock(Ordered, paragraphsl, _) -> yield! formatListBlock paragraphsl (fun n -> $"{n + 1}. ")
          | TableBlock(headers, alignments, rows, _) ->

              match headers with
              | Some headers ->
                  yield
                      headers
                      |> List.collect (fun hs -> [ for h in hs -> String.concat "" (formatParagraph ctx h) ])
                      |> String.concat " | "

              | None -> ()

              yield
                  [ for a in alignments ->
                        match a with
                        | AlignLeft -> ":---"
                        | AlignCenter -> ":---:"
                        | AlignRight -> "---:"
                        | AlignDefault -> "---" ]
                  |> String.concat " | "

              let replaceEmptyWith x s =
                  if System.String.IsNullOrWhiteSpace s then x else Some s

              yield
                  [ for r in rows do
                        [ for ps in r do
                              let x =
                                  [ for p in ps do
                                        yield
                                            formatParagraph ctx p
                                            |> Seq.choose (replaceEmptyWith (Some ""))
                                            |> String.concat "" ]

                              yield x |> Seq.choose (replaceEmptyWith (Some "")) |> String.concat "<br />" ]
                        |> Seq.choose (replaceEmptyWith (Some "&#32;"))
                        |> String.concat " | " ]
                  |> String.concat "\n"

              yield "\n"

          | OutputBlock(output, "text/html", _executionCount) ->
              yield (output.Trim())
              yield ""
          | OutputBlock(output, _, _executionCount) ->
              yield "```"
              yield output
              yield "```"
              yield ""
          | OtherBlock(lines, _) -> yield! List.map fst lines
          | InlineHtmlBlock(code, _, _) ->
              let lines = code.Replace("\r\n", "\n").Split('\n') |> Array.toList
              yield! lines
          | YamlFrontmatter(lines, _) ->
              yield "---"

              for line in lines do
                  yield line

              yield "---"
              yield ""
          | Span(body = body) -> yield formatSpans ctx body
          | QuotedBlock(paragraphs = paragraphs) ->
              for paragraph in paragraphs do
                  let lines = formatParagraph ctx paragraph

                  for line in lines do
                      yield "> " + line

                  yield ""
          | EmbedParagraphs(cmd, _) -> yield! cmd.Render() |> Seq.collect (formatParagraph ctx) ]

    /// Strips <c>#if SYMBOL</c> / <c>#endif // SYMBOL</c> conditional compilation lines from an .fsx code block
    /// so that format-specific sections are removed from non-target output formats.
    let adjustFsxCodeForConditionalDefines (defineSymbol, newLine) (code: string) =
        // Inside literate code blocks we conditionally remove some special lines to get nicer output for
        // load sections for different formats. We remove this:
        //   #if IPYNB
        //   #endif // IPYNB
        let sym1 = sprintf "#if %s" defineSymbol
        let sym2 = sprintf "#endif // %s" defineSymbol

        let lines = code.Replace("\r\n", "\n").Split('\n') |> Array.toList

        let lines = lines |> List.filter (fun line -> line.Trim() <> sym1 && line.Trim() <> sym2)

        let code2 = String.concat newLine lines
        code2

    /// Applies template substitutions to a plain text string using the context's substitution table.
    let applySubstitutionsInText ctx (text: string) =
        SimpleTemplating.ApplySubstitutionsInText ctx.Substitutions text

    /// Resolves a <c>cref:</c> inline-code span to a hyperlink if the reference is known; otherwise leaves it as inline code.
    let applyCodeReferenceResolver ctx (code, range) =
        match ctx.CodeReferenceResolver code with
        | None -> InlineCode(code, range)
        | Some(niceName, link) -> DirectLink([ Literal(niceName, range) ], link, None, range)

    /// Resolves a direct link target through the context's Markdown link resolver, returning the mapped URL.
    let applyDirectLinkResolver ctx link =
        match ctx.MarkdownDirectLinkResolver link with
        | None -> link
        | Some newLink -> newLink

    /// Extracts the text-transformation function from a triple of span-mapping functions.
    let mapText (f, _, _) text = f text
    /// Extracts the inline-code transformation function from a triple of span-mapping functions.
    let mapInlineCode (_, f, _) (code, range) = f (code, range)
    /// Applies the text function to the body and then the link function to the result.
    let mapDirectLink (fText, _, fLink) text = fLink (fText text)

    /// Recursively maps a triple of transformation functions over all spans in a span list.
    let rec mapSpans fs (md: MarkdownSpans) =
        md
        |> List.map (function
            | Literal(text, range) -> Literal(mapText fs text, range)
            | Strong(spans, range) -> Strong(mapSpans fs spans, range)
            | Emphasis(spans, range) -> Emphasis(mapSpans fs spans, range)
            | AnchorLink(link, range) -> AnchorLink(mapText fs link, range)
            | DirectLink(spans, link, title, range) ->
                DirectLink(mapSpans fs spans, mapDirectLink fs link, Option.map (mapText fs) title, range)
            | IndirectLink(spans, original, key, range) -> IndirectLink(mapSpans fs spans, original, key, range)
            | DirectImage(body, link, title, range) ->
                DirectImage(mapText fs body, mapText fs link, Option.map (mapText fs) title, range)
            | IndirectImage(body, original, key, range) -> IndirectImage(mapText fs body, original, key, range)
            | HardLineBreak(range) -> HardLineBreak(range)
            | InlineCode(code, range) -> mapInlineCode fs (code, range)

            // NOTE: substitutions not applied to Latex math, embedded spans or inline code
            | LatexInlineMath(code, range) -> LatexInlineMath(code, range)
            | LatexDisplayMath(code, range) -> LatexDisplayMath(code, range)
            | EmbedSpans(customSpans, range) -> EmbedSpans(customSpans, range))

    /// Recursively maps a triple of transformation functions over all paragraphs in a paragraph list,
    /// including nested paragraphs inside list items, block-quotes, and tables.
    let rec mapParagraphs f (md: MarkdownParagraphs) =
        md
        |> List.map (function
            | Heading(size, body, range) -> Heading(size, mapSpans f body, range)
            | Paragraph(body, range) -> Paragraph(mapSpans f body, range)
            | CodeBlock(code, count, fence, language, ignoredLine, range) ->
                CodeBlock(mapText f code, count, fence, language, ignoredLine, range)
            | OutputBlock(output, kind, count) -> OutputBlock(output, kind, count)
            | ListBlock(kind, items, range) -> ListBlock(kind, List.map (mapParagraphs f) items, range)
            | QuotedBlock(paragraphs, range) -> QuotedBlock(mapParagraphs f paragraphs, range)
            | Span(spans, range) -> MarkdownParagraph.Span(mapSpans f spans, range)
            | LatexBlock(env, body, range) -> LatexBlock(env, List.map (mapText f) body, range)
            | HorizontalRule(character, range) -> HorizontalRule(character, range)
            | YamlFrontmatter(lines, range) -> YamlFrontmatter(List.map (mapText f) lines, range)
            | TableBlock(headers, alignments, rows, range) ->
                TableBlock(
                    Option.map (List.map (mapParagraphs f)) headers,
                    alignments,
                    List.map (List.map (mapParagraphs f)) rows,
                    range
                )
            | OtherBlock(lines: (string * MarkdownRange) list, range) ->
                OtherBlock(lines |> List.map (fun (line, range) -> (mapText f line, range)), range)
            | InlineHtmlBlock(code, count, range) ->
                try
                    let fText, _, fLink = f

                    if
                        code.StartsWith("<pre", StringComparison.Ordinal)
                        || code.StartsWith("<table class=\"pre\"", StringComparison.Ordinal)
                    then
                        // Skip check for non-user html
                        // Should be even run that code through `fText`?
                        InlineHtmlBlock(fText code, count, range)
                    else
                        let tempRoot = "fsdocs-secret-temp-root"
                        // We can't be sure code is a single html element, we could get multiple elements.
                        let element = XElement.Parse($"<%s{tempRoot}>%s{code}</%s{tempRoot}>")
                        // ends-with is XPath 2.0 only, https://stackoverflow.com/questions/1525299/xpath-and-xslt-2-0-for-net
                        let attributes =
                            match System.Xml.XPath.Extensions.XPathEvaluate(element, "//*/@*[contains(., '.md')]") with
                            | :? System.Collections.IEnumerable as enumerable ->
                                enumerable |> Seq.cast<XAttribute> |> Seq.toArray
                            | _ -> Array.empty

                        if Array.isEmpty attributes then
                            InlineHtmlBlock(fText code, count, range)
                        else
                            for attribute in attributes do
                                if attribute.Value.EndsWith(".md", StringComparison.Ordinal) then
                                    attribute.SetValue(fLink attribute.Value)

                            let html = element.Elements() |> Seq.map string<XElement> |> String.concat "" |> fText
                            InlineHtmlBlock(html, count, range)
                with ex ->
                    InlineHtmlBlock(mapText f code, count, range)

            // NOTE: substitutions are not currently applied to embedded LiterateParagraph which are in any case eliminated
            // before substitutions are applied.
            | EmbedParagraphs(customParagraphs, range) ->
                //let customParagraphsR = { new MarkdownEmbedParagraphs with member _.Render() = customParagraphs.Render() |> mapParagraphs f }
                EmbedParagraphs(customParagraphs, range))

    /// Applies all context substitutions (text replacements, cref resolution, and direct-link mapping)
    /// to every span and paragraph in a Markdown document.
    let applySubstitutionsInMarkdown ctx md =
        mapParagraphs (applySubstitutionsInText ctx, applyCodeReferenceResolver ctx, applyDirectLinkResolver ctx) md