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, LocalDiagrams.py
draws them on this computer instead (Settings > General, on by default). If
neither can, 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

import LocalDiagrams
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 IsLocal(TheApp) -> bool:
    return TheApp.Source.Kind != "Server"


def CanDraw(TheApp) -> bool:
    if IsLocal(TheApp):
        return TheApp.Config.LocalDiagrams
    return 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 self._WithNote(f"Diagram error: {Problem}", Palette.Red)
        if self.Fetching and IsLocal(self.TheApp) and not LocalDiagrams.IsSetUp():
            return self._WithNote("Setting up diagrams (first time only, can take a few minutes)…", Palette.Overlay0)
        return self._Code()

    def _WithNote(self, Note: str, Colour: str) -> ft.Control:
        return ft.Column([self._Code(), ft.Text(Note, color=Colour, size=self.FontSize - 2, selectable=True)], spacing=4)

    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
            if IsLocal(self.TheApp) and not LocalDiagrams.IsSetUp():
                self.content = self._Content()
                self.update()
            self.page.run_thread(self._Fetch)

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

    def _Fetch(self) -> None:
        global ServerCantDraw
        try:
            if IsLocal(self.TheApp):
                Pictures[self.Key] = LocalDiagrams.Render(self.Source, ThemePalette(), self.FontSize)
            else:
                Pictures[self.Key] = self.TheApp.Source.Diagram(self.Source, ThemePalette(), self.FontSize)
        except LocalDiagrams.DiagramError as Error:
            Problems[self.Key] = str(Error)
        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
        except Exception as Error:
            # Can't draw on this computer (setup failed, or it timed out): say why.
            Problems[self.Key] = f"Can't draw diagrams here: {Error or type(Error).__name__}"
        finally:
            self._Finish()

    def _Finish(self) -> None:
        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