SaraView/Client/Diagrams.py

"""Mermaid diagrams (```mermaid blocks), like the web view's. Flet can't draw
them, so the server does (POST /api/diagram, Backend/Diagrams.py) and they're
shown as pictures in this client's theme. In a local folder (no server), or
if the server can't draw them, they stay code blocks.

A DiagramBlock shows its code straight away and swaps in the picture once it
arrives (fetched on a background thread), so opening a note never waits on
the server. Pictures are kept for the session, so the live editor's
re-renders don't ask again."""

import re
import struct
from typing import Callable, Optional

import flet as ft

from Core import CodeStyleSheet, MarkdownCodeTheme, MarkdownSheet, Palette
from Sources import SourceError

DiagramFence = re.compile(r"^ {0,3}(`{3,}|~{3,})[ \t]*mermaid[ \t]*$", re.IGNORECASE)

# Per session: pictures and Mermaid's complaints by (source, font size), and
# whether the server said it can't draw diagrams at all.
Pictures: dict[tuple[str, int], bytes] = {}
Problems: dict[tuple[str, int], str] = {}
ServerCantDraw = False


def DiagramSource(BlockText: str) -> Optional[str]:
    """The diagram inside a ```mermaid code block, or None for other blocks."""
    Lines = BlockText.split("\n")
    Match = DiagramFence.match(Lines[0])
    if not Match:
        return None
    Fence = Match.group(1)
    Body = Lines[1:]
    if Body:
        Last = Body[-1].strip()
        if Last.startswith(Fence) and not Last.strip(Fence[0]):
            Body = Body[:-1]
    return "\n".join(Body)


def ThemePalette() -> dict[str, str]:
    """This client's theme by the server's role names (see Core.Palette)."""
    return {
        "Background": Palette.Base, "Surface": Palette.Mantle, "SurfaceAlt": Palette.Surface0,
        "Border": Palette.Surface1, "Text": Palette.Text, "TextMuted": Palette.Overlay0,
        "Black": Palette.Crust, "White": Palette.White, "Red": Palette.Red, "Orange": Palette.Peach,
        "Yellow": Palette.Yellow, "Green": Palette.Green, "Blue": Palette.Blue, "Purple": Palette.Mauve,
    }


def CanDraw(TheApp) -> bool:
    return TheApp.Source.Kind == "Server" and not ServerCantDraw


class DiagramBlock(ft.Container):
    def __init__(self, TheApp, Source: str, FontSize: int):
        self.TheApp = TheApp
        self.Source = Source
        self.FontSize = FontSize
        self.Key = (Source, FontSize)
        self.Mounted = False
        self.Fetching = False
        super().__init__(content=self._Content())

    def _Code(self) -> ft.Control:
        return ft.Markdown(
            value=f"```\n{self.Source}\n```",
            extension_set=ft.MarkdownExtensionSet.GITHUB_FLAVORED,
            md_style_sheet=MarkdownSheet(self.FontSize),
            code_theme=MarkdownCodeTheme(),
            code_style_sheet=CodeStyleSheet(),
        )

    def _Content(self) -> ft.Control:
        Picture = Pictures.get(self.Key)
        if Picture is not None:
            # The server draws at twice the size; shown at half, it's sharp on HiDPI.
            Width, Height = struct.unpack(">II", Picture[16:24])
            return ft.Row([ft.Image(src=Picture, width=Width / 2, height=Height / 2)],
                          alignment=ft.MainAxisAlignment.CENTER, scroll=ft.ScrollMode.AUTO)
        Problem = Problems.get(self.Key)
        if Problem is not None:
            return ft.Column([self._Code(), ft.Text(f"Diagram error: {Problem}", color=Palette.Red,
                                                    size=self.FontSize - 2, selectable=True)], spacing=4)
        return self._Code()

    def did_mount(self) -> None:
        self.Mounted = True
        if self.Key not in Pictures and self.Key not in Problems and not self.Fetching and CanDraw(self.TheApp):
            self.Fetching = True
            self.page.run_thread(self._Fetch)

    def will_unmount(self) -> None:
        self.Mounted = False

    def _Fetch(self) -> None:
        global ServerCantDraw
        try:
            Pictures[self.Key] = self.TheApp.Source.Diagram(self.Source, ThemePalette(), self.FontSize)
        except SourceError as Error:
            if Error.Status == 422:
                Problems[self.Key] = str(Error)
            elif Error.Status in (404, 503):
                ServerCantDraw = True  # an older server, or one without Chromium: stay code blocks
        finally:
            self.Fetching = False
        if self.Mounted:
            self.content = self._Content()
            self.update()


def MaybeDiagram(TheApp, BlockText: str, FontSize: int) -> Optional[ft.Control]:
    """A DiagramBlock for a ```mermaid code block (LiveMarkdown's blocks), else None."""
    Source = DiagramSource(BlockText)
    if Source is None or not CanDraw(TheApp):
        return None
    return DiagramBlock(TheApp, Source, FontSize)


def WithDiagrams(TheApp, Document: str, FontSize: int, MakeMarkdown: Callable[[str], ft.Control]) -> list[ft.Control]:
    """The rendered view's controls: the Markdown, split around its diagrams."""
    if not CanDraw(TheApp) or not any(DiagramFence.match(Line) for Line in Document.split("\n")):
        return [MakeMarkdown(Document)]
    from LiveMarkdown import SplitBlocks  # LiveMarkdown imports this module

    Controls: list[ft.Control] = []
    Position = 0
    for Block in SplitBlocks(Document):
        Source = DiagramSource(Document[Block.Start:Block.End]) if Block.Kind == "Code" else None
        if Source is None:
            continue
        if Document[Position:Block.Start].strip():
            Controls.append(MakeMarkdown(Document[Position:Block.Start]))
        Controls.append(DiagramBlock(TheApp, Source, FontSize))
        Position = Block.End
    if Document[Position:].strip():
        Controls.append(MakeMarkdown(Document[Position:]))
    return Controls