"""Mermaid diagrams as images, for the desktop client (POST /api/diagram).

The web view draws ```mermaid blocks itself (Frontend/src/Mermaid.js), but
Flet can't, so the client asks the server instead. This draws them the same
way, with the same Mermaid, in a headless Chromium (Playwright) that's
started on the first diagram and kept for the next ones, and returns a PNG
at twice the size (sharp on HiDPI screens; the client shows it at half).

Mermaid comes from the frontend build (the Dockerfile copies it to
Backend/Vendor/; for local development, Frontend/node_modules after
`npm install`), the font from Frontend's public/Fonts."""

from __future__ import annotations

import asyncio
import hashlib
import json
import re
import tempfile
import tomllib
from collections import OrderedDict
from pathlib import Path
from typing import Any, Optional

from Settings import DefaultThemesDirectory, PaletteRoles

ApplicationDirectory = Path(__file__).resolve().parent
ProjectDirectory = ApplicationDirectory.parent

MermaidCandidates = [
    ApplicationDirectory / "Vendor" / "mermaid.min.js",
    ProjectDirectory / "Frontend" / "node_modules" / "mermaid" / "dist" / "mermaid.min.js",
]
FontCandidates = [
    ProjectDirectory / "Frontend" / "dist" / "Fonts" / "FSEX302.ttf",
    ProjectDirectory / "Frontend" / "public" / "Fonts" / "FSEX302.ttf",
]

MaxSourceLength = 20_000
RenderTimeout = 20  # seconds
CacheSize = 200
HexColour = re.compile(r"^#[0-9a-fA-F]{6}$")


class DiagramError(Exception):
    """The diagram itself is wrong (Mermaid's message)."""


class DiagramsUnavailable(Exception):
    """This server can't draw diagrams (no Chromium or no Mermaid)."""


def _FirstExisting(Candidates: list[Path]) -> Optional[Path]:
    return next((Candidate for Candidate in Candidates if Candidate.is_file()), None)


def _DefaultPalette() -> dict[str, str]:
    with (DefaultThemesDirectory / "CatppuccinMocha.toml").open("rb") as ThemeFile:
        return tomllib.load(ThemeFile).get("Palette", {})


def CleanPalette(Palette: dict[str, Any]) -> dict[str, str]:
    """The default theme, overridden by whichever roles were sent as #rrggbb."""
    Clean = _DefaultPalette()
    for Role in PaletteRoles:
        Value = Palette.get(Role)
        if isinstance(Value, str) and HexColour.match(Value):
            Clean[Role] = Value
    return Clean


def ThemeVariables(Palette: dict[str, str], FontSize: int) -> dict[str, Any]:
    """Mermaid's colours from the theme's roles: the same mapping as the web
    view's Frontend/src/Mermaid.js, so a diagram looks the same in both."""
    Role = Palette.get
    Background = Role("Background")
    Red, Green, Blue = (int(Background[Index:Index + 2], 16) for Index in (1, 3, 5))
    Luminance = (0.299 * Red + 0.587 * Green + 0.114 * Blue) / 255
    return {
        "darkMode": Luminance < 0.5,
        "background": Background,
        "fontFamily": '"Fixedsys Excelsior", "Cascadia Mono", monospace',
        "fontSize": f"{FontSize}px",
        "textColor": Role("Text"),
        "lineColor": Role("TextMuted"),
        "primaryColor": Role("SurfaceAlt"),
        "primaryTextColor": Role("Text"),
        "primaryBorderColor": Role("Purple"),
        "secondaryColor": Role("Surface"),
        "secondaryTextColor": Role("Text"),
        "secondaryBorderColor": Role("Blue"),
        "tertiaryColor": Role("Background"),
        "tertiaryTextColor": Role("Text"),
        "tertiaryBorderColor": Role("Border"),
        "mainBkg": Role("SurfaceAlt"),
        "nodeBorder": Role("Purple"),
        "clusterBkg": Role("Surface"),
        "clusterBorder": Role("Border"),
        "edgeLabelBackground": Role("Surface"),
        "titleColor": Role("Text"),
        "noteBkgColor": Role("Surface"),
        "noteTextColor": Role("Text"),
        "noteBorderColor": Role("Yellow"),
        "actorBkg": Role("SurfaceAlt"),
        "actorBorder": Role("Purple"),
        "actorTextColor": Role("Text"),
        "actorLineColor": Role("TextMuted"),
        "signalColor": Role("Text"),
        "signalTextColor": Role("Text"),
        "labelBoxBkgColor": Role("SurfaceAlt"),
        "labelBoxBorderColor": Role("Purple"),
        "labelTextColor": Role("Text"),
        "loopTextColor": Role("Text"),
        "activationBkgColor": Role("Surface"),
        "activationBorderColor": Role("Blue"),
        "pie1": Role("Blue"),
        "pie2": Role("Purple"),
        "pie3": Role("Green"),
        "pie4": Role("Yellow"),
        "pie5": Role("Orange"),
        "pie6": Role("Red"),
        "pieTitleTextColor": Role("Text"),
        "pieSectionTextColor": Role("Black"),
        "pieLegendTextColor": Role("Text"),
        "pieStrokeColor": Background,
        "pieOuterStrokeColor": Role("Border"),
    }


# Draws #Out from the source; the page has Mermaid and the font loaded.
RenderScript = """async ({ Source, Variables, Id }) => {
  await document.fonts.ready;
  mermaid.initialize({ startOnLoad: false, securityLevel: "strict", theme: "base", themeVariables: Variables });
  try {
    await mermaid.parse(Source);
  } catch (Error) {
    return { Error: String(Error.message || Error) };
  }
  const { svg } = await mermaid.render(Id, Source);
  const Out = document.getElementById("Out");
  Out.innerHTML = svg;
  // Mermaid shrinks diagrams to fit their container; here they get their own size.
  const Picture = Out.querySelector("svg");
  const Box = Picture.viewBox.baseVal;
  Picture.style.maxWidth = "none";
  Picture.setAttribute("width", Box.width);
  Picture.setAttribute("height", Box.height);
  return {};
}"""


class _Renderer:
    def __init__(self) -> None:
        self.Lock = asyncio.Lock()
        self.Playwright = None
        self.Browser = None
        self.Page = None
        self.PageDirectory: Optional[tempfile.TemporaryDirectory] = None
        self.Count = 0
        self.Cache: OrderedDict[str, bytes] = OrderedDict()

    async def _EnsurePage(self) -> None:
        if self.Page is not None and not self.Page.is_closed():
            return
        Mermaid = _FirstExisting(MermaidCandidates)
        if Mermaid is None:
            raise DiagramsUnavailable("Diagrams aren't set up on this server (Mermaid is missing).")
        try:
            from playwright.async_api import async_playwright
        except ImportError:
            raise DiagramsUnavailable("Diagrams aren't set up on this server (Playwright is missing).")

        if self.Browser is None or not self.Browser.is_connected():
            try:
                self.Playwright = self.Playwright or await async_playwright().start()
                self.Browser = await self.Playwright.chromium.launch()
            except Exception as Error:
                raise DiagramsUnavailable(f"Diagrams aren't set up on this server ({Error}).")

        Font = _FirstExisting(FontCandidates)
        FontFace = f'@font-face {{ font-family: "Fixedsys Excelsior"; src: url("{Font.as_uri()}"); }}' if Font else ""
        self.PageDirectory = self.PageDirectory or tempfile.TemporaryDirectory(prefix="SaraViewDiagrams")
        PagePath = Path(self.PageDirectory.name) / "Render.html"
        PagePath.write_text(
            "<!doctype html><html><head><meta charset='utf-8'>"
            f"<style>{FontFace} body {{ margin: 0; }} #Out {{ display: inline-block; padding: 8px; }}</style>"
            f"<script src='{Mermaid.as_uri()}'></script></head><body><div id='Out'></div></body></html>",
            encoding="utf-8",
        )
        Context = await self.Browser.new_context(device_scale_factor=2, viewport={"width": 1600, "height": 1200})
        self.Page = await Context.new_page()
        # Only the page's own files: a diagram can't make the server fetch anything.
        async def OnlyLocalFiles(Route) -> None:
            if Route.request.url.startswith("file:"):
                await Route.continue_()
            else:
                await Route.abort()

        await self.Page.route("**/*", OnlyLocalFiles)
        await self.Page.goto(PagePath.as_uri())

    async def Render(self, Source: str, Palette: dict[str, Any], FontSize: int) -> bytes:
        if not Source.strip():
            raise DiagramError("The diagram is empty.")
        if len(Source) > MaxSourceLength:
            raise DiagramError(f"The diagram is too long (over {MaxSourceLength} characters).")
        Variables = ThemeVariables(CleanPalette(Palette), max(8, min(int(FontSize), 40)))
        Key = hashlib.sha256(json.dumps([Source, Variables], sort_keys=True).encode()).hexdigest()

        async with self.Lock:
            if Key in self.Cache:
                self.Cache.move_to_end(Key)
                return self.Cache[Key]
            await self._EnsurePage()
            self.Count += 1
            try:
                Result = await asyncio.wait_for(
                    self.Page.evaluate(RenderScript, {"Source": Source, "Variables": Variables, "Id": f"Diagram{self.Count}"}),
                    RenderTimeout,
                )
                if Result.get("Error"):
                    raise DiagramError(Result["Error"])
                Png = await self.Page.locator("#Out").screenshot(omit_background=True, timeout=RenderTimeout * 1000)
            except DiagramError:
                raise
            except Exception as Error:
                # A stuck or crashed page: start a fresh one next time.
                if self.Page is not None:
                    await self.Page.context.close()
                self.Page = None
                raise DiagramError(f"Couldn't draw the diagram ({Error}).")

            self.Cache[Key] = Png
            if len(self.Cache) > CacheSize:
                self.Cache.popitem(last=False)
            return Png


_SharedRenderer = _Renderer()


async def RenderDiagram(Source: str, Palette: dict[str, Any], FontSize: int) -> bytes:
    return await _SharedRenderer.Render(Source, Palette, FontSize)