"""Drawing Mermaid diagrams in a local folder, where there's no server to do
it (Diagrams.py): the same way the server's Backend/Diagrams.py does, with
Mermaid in a headless Chromium driven by Playwright.

Settings > General > "Draw diagrams in local folders" (Config.LocalDiagrams,
on by default). Nothing for this ships with SaraView: the first diagram sets
up what's needed, all of it in ToolsDir (~/.local/share/SaraView/DiagramTools):
- Playwright, installed with pip into ToolsDir/Python;
- Mermaid, the version the web view uses, downloaded from npm and checked;
- a headless Chromium (ToolsDir/Browsers), only when no Chromium-based
  browser (Chrome, Chromium, Brave, Edge) is installed or it won't start.
  An installed one is used as it is, never changed.
Turning the setting off deletes ToolsDir: all of that, and nothing else.

Playwright's sync API has to stay on the thread that started it, so one
worker thread does all the drawing; Render() hands it a job and waits."""

import base64
import hashlib
import io
import os
import queue
import shutil
import subprocess
import sys
import tarfile
import threading
from concurrent.futures import Future
from pathlib import Path
from typing import Optional

import httpx

from Core import AssetsDir, ChildEnvironment, DataHome, AppName, IsWindows

ToolsDir = DataHome / AppName / "DiagramTools"
PythonDir = ToolsDir / "Python"
BrowsersDir = ToolsDir / "Browsers"
MermaidPath = ToolsDir / "mermaid.min.js"
PagePath = ToolsDir / "Render.html"

PlaywrightVersion = "1.62.0"  # same as the server's (Backend/requirements.txt)
MermaidVersion = "11.17.2"
MermaidIntegrity = "sha512-V6K3C8EBdEsPFZXSKMJe6ppQOENxuHARr9GvHX4hh47lAbhMRD9qf4oEK7LoaRQxULMa80/qt5gHO73aCleBBg=="
FontPath = AssetsDir / "fonts" / "Fixedsys_Excelsior_Regular.ttf"
RenderTimeout = 20  # seconds, per diagram
SetupTimeout = 900  # seconds, for the downloads

BrowserNames = ["google-chrome", "google-chrome-stable", "chromium", "chromium-browser",
                "brave-browser", "brave", "microsoft-edge", "microsoft-edge-stable"]
WindowsBrowsers = [
    r"Microsoft\Edge\Application\msedge.exe",
    r"Google\Chrome\Application\chrome.exe",
    r"BraveSoftware\Brave-Browser\Application\brave.exe",
]


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


class Unavailable(Exception):
    """Diagrams can't be drawn here (setup failed, or a frozen build)."""


def ThemeVariables(Palette: dict[str, str], FontSize: int) -> dict:
    """Mermaid's colours from the theme's roles: the same mapping as
    Backend/Diagrams.py and the web view's Frontend/src/Mermaid.js."""
    Role = Palette.get
    Background = Role("Background")
    Red, Green, Blue = (int(Background[Index:Index + 2], 16) for Index in (1, 3, 5))
    return {
        "darkMode": (0.299 * Red + 0.587 * Green + 0.114 * Blue) / 255 < 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"),
    }


# The same drawing as Backend/Diagrams.py's RenderScript.
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;
  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 {};
}"""


# --- Setup ----------------------------------------------------------------

def _Run(Command: list[str], Environment: Optional[dict] = None) -> None:
    Result = subprocess.run(Command, env=Environment or ChildEnvironment(), capture_output=True, text=True,
                            timeout=SetupTimeout)
    if Result.returncode != 0:
        Detail = (Result.stderr or Result.stdout).strip().splitlines()
        raise Unavailable(Detail[-1] if Detail else f"{Command[0]} failed")


def _PlaywrightEnvironment() -> dict:
    Environment = ChildEnvironment()
    Environment["PYTHONPATH"] = str(PythonDir)
    Environment["PLAYWRIGHT_BROWSERS_PATH"] = str(BrowsersDir)
    return Environment


def _InstallPlaywright() -> None:
    if (PythonDir / "playwright").is_dir():
        return
    if getattr(sys, "frozen", False):
        raise Unavailable("This build of SaraView can't install what diagrams need; "
                          "the SaraSuite install can.")
    _Run([sys.executable, "-m", "pip", "install", "--quiet", "--disable-pip-version-check",
          "--target", str(PythonDir), f"playwright=={PlaywrightVersion}"])


def _InstallMermaid() -> None:
    if MermaidPath.is_file():
        return
    Url = f"https://registry.npmjs.org/mermaid/-/mermaid-{MermaidVersion}.tgz"
    try:
        Archive = httpx.get(Url, timeout=120, follow_redirects=True).raise_for_status().content
    except httpx.HTTPError as Error:
        raise Unavailable(f"Couldn't download Mermaid: {Error}")
    Expected = base64.b64decode(MermaidIntegrity.removeprefix("sha512-"))
    if hashlib.sha512(Archive).digest() != Expected:
        raise Unavailable("The Mermaid download didn't match its checksum.")
    with tarfile.open(fileobj=io.BytesIO(Archive), mode="r:gz") as Tar:
        Script = Tar.extractfile("package/dist/mermaid.min.js")
        if Script is None:
            raise Unavailable("The Mermaid download didn't have mermaid.min.js.")
        Partial = MermaidPath.with_suffix(".part")
        Partial.write_bytes(Script.read())
        Partial.replace(MermaidPath)


def _InstalledBrowser() -> Optional[str]:
    if IsWindows:
        for Base in (os.environ.get("ProgramFiles(x86)"), os.environ.get("ProgramFiles"),
                     os.environ.get("LOCALAPPDATA")):
            for Relative in WindowsBrowsers:
                if Base and (Path(Base) / Relative).is_file():
                    return str(Path(Base) / Relative)
        return None
    return next((Found for Name in BrowserNames if (Found := shutil.which(Name))), None)


def _InstallBrowser() -> None:
    _Run([sys.executable, "-m", "playwright", "install", "--only-shell", "chromium"], _PlaywrightEnvironment())


def _WritePage() -> None:
    FontFace = f'@font-face {{ font-family: "Fixedsys Excelsior"; src: url("{FontPath.as_uri()}"); }}' if FontPath.is_file() else ""
    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='{MermaidPath.as_uri()}'></script></head><body><div id='Out'></div></body></html>",
        encoding="utf-8",
    )


# --- The worker -----------------------------------------------------------

class _Worker:
    def __init__(self) -> None:
        self.Jobs: "queue.Queue[Optional[tuple]]" = queue.Queue()
        self.Thread: Optional[threading.Thread] = None
        self.Lock = threading.Lock()
        self.Failure: Optional[str] = None  # why it can't draw, once setup failed
        self.SettingUp = False

    def Start(self) -> None:
        with self.Lock:
            if self.Thread is None or not self.Thread.is_alive():
                self.Thread = threading.Thread(target=self._Run, name="LocalDiagrams", daemon=True)
                self.Thread.start()

    def Stop(self) -> None:
        with self.Lock:
            Thread = self.Thread
            self.Thread = None
        if Thread is not None and Thread.is_alive():
            self.Jobs.put(None)
            Thread.join(timeout=30)

    def _Setup(self):
        """Everything ToolsDir needs, then a page ready to draw on."""
        self.SettingUp = True
        try:
            ToolsDir.mkdir(parents=True, exist_ok=True)
            _InstallPlaywright()
            _InstallMermaid()
            _WritePage()
            if str(PythonDir) not in sys.path:
                sys.path.insert(0, str(PythonDir))
            os.environ["PLAYWRIGHT_BROWSERS_PATH"] = str(BrowsersDir)
            from playwright.sync_api import sync_playwright

            Playwright = sync_playwright().start()
            Browser = None
            Installed = _InstalledBrowser()
            if Installed:
                try:
                    Browser = Playwright.chromium.launch(executable_path=Installed)
                except Exception:
                    Browser = None  # e.g. a Snap Chromium that can't be driven: use our own
            if Browser is None:
                if not BrowsersDir.is_dir():
                    _InstallBrowser()
                Browser = Playwright.chromium.launch()
            Page = Browser.new_context(device_scale_factor=2, viewport={"width": 1600, "height": 1200}).new_page()
            Page.route("**/*", lambda Route: Route.continue_() if Route.request.url.startswith("file:") else Route.abort())
            Page.goto(PagePath.as_uri())
            return Playwright, Browser, Page
        finally:
            self.SettingUp = False

    def _Run(self) -> None:
        Playwright = Browser = Page = None
        Count = 0
        try:
            while True:
                Job = self.Jobs.get()
                if Job is None:
                    return
                Source, Palette, FontSize, Result = Job
                try:
                    if Page is None:
                        Playwright, Browser, Page = self._Setup()
                    Count += 1
                    Outcome = Page.evaluate(RenderScript, {"Source": Source, "Variables": ThemeVariables(Palette, FontSize),
                                                           "Id": f"Diagram{Count}"})
                    if Outcome.get("Error"):
                        raise DiagramError(Outcome["Error"])
                    Result.set_result(Page.locator("#Out").screenshot(omit_background=True, timeout=RenderTimeout * 1000))
                except DiagramError as Error:
                    Result.set_exception(Error)
                except Unavailable as Error:
                    self.Failure = str(Error)
                    Result.set_exception(Error)
                except Exception as Error:
                    if Page is None:  # setup itself broke
                        self.Failure = str(Error)
                        Result.set_exception(Unavailable(str(Error)))
                    else:
                        Result.set_exception(DiagramError(f"Couldn't draw the diagram ({Error})."))
        finally:
            for Closing in (Browser, Playwright):
                try:
                    if Closing is not None:
                        Closing.close() if Closing is Browser else Closing.stop()
                except Exception:
                    pass


_SharedWorker = _Worker()


def IsSetUp() -> bool:
    """Whether the first diagram will draw without downloading anything."""
    return (PythonDir / "playwright").is_dir() and MermaidPath.is_file()


def IsSettingUp() -> bool:
    return _SharedWorker.SettingUp


def Render(Source: str, Palette: dict[str, str], FontSize: int) -> bytes:
    """A diagram as a PNG at twice the size. Blocks: call it off the UI thread.
    Raises DiagramError for a wrong diagram, Unavailable if it can't draw."""
    if _SharedWorker.Failure is not None:
        raise Unavailable(_SharedWorker.Failure)
    _SharedWorker.Start()
    Result: Future = Future()
    _SharedWorker.Jobs.put((Source, Palette, FontSize, Result))
    return Result.result(timeout=SetupTimeout + RenderTimeout)


def Remove() -> None:
    """Stops drawing and deletes ToolsDir: everything set up for local diagrams."""
    _SharedWorker.Stop()
    _SharedWorker.Failure = None
    shutil.rmtree(ToolsDir, ignore_errors=True)


def ToolsSize() -> int:
    """Bytes in ToolsDir, for the setting's note."""
    return sum(File.stat().st_size for File in ToolsDir.rglob("*") if File.is_file()) if ToolsDir.is_dir() else 0