#!/usr/bin/env python3
"""Render the themed ESPHome pages to PNG on this PC, for checking layouts
without flashing and for example images.

Compiles the page code (components/el512/cyberpunk, konosuba, steins, trek)
natively with g++ against ESPHome's own display and font code (from the
installed esphome package) and the exact font bitmaps of the last
`esphome compile` (its generated main.cpp), then draws each page at given
moments after it was shown. Run `esphome compile esphome/el512-panel-local.yaml`
first, and again after changing fonts in the YAML.

    python tools/render_pages.py                          # all pages, default moments
    python tools/render_pages.py --page ks_explosion --at 0,6900,15000
    python tools/render_pages.py --scenario busy          # also: normal, empty, alert
    python tools/render_pages.py --sheet                  # plus one contact sheet
    python tools/render_pages.py --list                   # page names and moments
    python tools/render_pages.py --examples               # docs/images/pages/<page>.png

Output goes to .render/out (git-ignored) unless --out says otherwise. Needs a
g++ with C++20 (MinGW-w64 on Windows works) and Pillow.
"""

import argparse
import re
import shutil
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
WORK = ROOT / ".render"
PAGE_SOURCES = ["cyberpunk", "konosuba", "steins", "trek"]
PAGE_HEADERS = ["el512_pages.h", "el512_helpers.h", "konosuba_sprites.h", "steins_sprites.h", "trek_sprites.h"]
ESPHOME_PARTS = ["core", "components/display", "components/font", "components/host", "components/time",
                 "components/lvgl", "components/graph"]
ESPHOME_SOURCES = ["components/display/display.cpp", "components/display/rect.cpp",
                   "components/font/font.cpp", "core/time.cpp", "core/color.cpp", "core/helpers.cpp"]
# ESPHome's headers include "defines.h" from their own folder; on the host only
# the display and font code is wanted.
HOST_DEFINES = "#pragma once\n#define USE_HOST\n#define USE_DISPLAY\n#define USE_FONT\n#define ESPHOME_LOG_LEVEL 0\n"
COMPAT = """#pragma once
// MinGW lacks strcasestr, which ESPHome's helpers use.
#ifdef __cplusplus
#include <cctype>
#include <cstring>
inline char *strcasestr(const char *h, const char *n) {
  const size_t len = strlen(n);
  for (; *h; h++)
    if (strncasecmp(h, n, len) == 0)
      return (char *) h;
  return len == 0 ? (char *) h : nullptr;
}
#endif
"""
AMBER = (255, 170, 30)
BACKGROUND = (20, 12, 0)


def run(cmd, **kw):
    r = subprocess.run(cmd, capture_output=True, text=True, **kw)
    if r.returncode != 0:
        sys.stderr.write(r.stdout + r.stderr)
        raise SystemExit(f"failed: {' '.join(str(c) for c in cmd[:3])} ...")
    return r.stdout


def write_if_changed(path, text):
    if not path.exists() or path.read_text(encoding="utf-8") != text:
        path.write_text(text, encoding="utf-8")


def stage_esphome(stage):
    """Copy the needed part of the installed ESPHome sources, with host defines."""
    import esphome

    src = Path(esphome.__file__).parent
    marker = stage / "esphome" / ".version"
    version = getattr(__import__("esphome.const", fromlist=["__version__"]), "__version__", "?")
    if marker.exists() and marker.read_text() == version:
        return version
    shutil.rmtree(stage / "esphome", ignore_errors=True)
    for part in ESPHOME_PARTS:
        dst = stage / "esphome" / part
        dst.mkdir(parents=True, exist_ok=True)
        for f in (src / part).iterdir():
            if f.suffix in (".h", ".cpp", ".c"):
                shutil.copy2(f, dst / f.name)
    (stage / "esphome" / "core" / "defines.h").write_text(HOST_DEFINES)
    shutil.rmtree(stage / "obj", ignore_errors=True)  # rebuild ESPHome objects
    marker.write_text(version)
    return version


def extract_fonts(main_cpp, stage):
    """Turn the fonts of ESPHome's generated main.cpp into fonts.cpp / fonts.h."""
    text = main_cpp.read_text(encoding="utf-8")
    cpp = ['#include "fonts.h"', "", "using esphome::font::Font;", "using esphome::font::Glyph;", ""]
    hdr = ["#pragma once", '#include "esphome/components/font/font.h"', ""]
    for name, glyphs, args in re.findall(r"new\((\w+)\) font::Font\((\w+), ([^;]*)\);", text):
        gdef = re.search(rf"static const font::Glyph {glyphs}\[\] = (\{{.*?\}});\n", text)
        if gdef is None:
            raise SystemExit(f"glyphs of {name} not found in {main_cpp}")
        for data in sorted(set(re.findall(r"\((uint8_t_id\w*) \+", gdef.group(1)))):
            ddef = re.search(rf"static constexpr uint8_t {data}\[\] PROGMEM = (\{{.*?\}});\n", text)
            cpp.append(f"static constexpr uint8_t {data}[] = {ddef.group(1)};")
        cpp.append(f"static const Glyph {glyphs}[] = {gdef.group(1)};")
        cpp.append(f"Font *const {name} = new Font({glyphs}, {args});")
        cpp.append("")
        hdr.append(f"extern esphome::font::Font *const {name};")
    if len(hdr) == 3:
        raise SystemExit(f"no fonts found in {main_cpp}")
    write_if_changed(stage / "fonts.cpp", "\n".join(cpp) + "\n")
    write_if_changed(stage / "fonts.h", "\n".join(hdr) + "\n")


def stage_pages(stage):
    """Page sources, with the host EL512Display in place of the panel's."""
    pages = stage / "el512"
    pages.mkdir(parents=True, exist_ok=True)
    comp = ROOT / "components" / "el512"
    for name in PAGE_HEADERS + [f"{p}.h" for p in PAGE_SOURCES] + [f"{p}.cpp" for p in PAGE_SOURCES]:
        write_if_changed(pages / name, (comp / name).read_text(encoding="utf-8"))
    here = Path(__file__).resolve().parent / "render"
    write_if_changed(pages / "el512_display.h", (here / "el512_display.h").read_text(encoding="utf-8"))
    for name in ("render_main.cpp", "host_stubs.cpp"):
        write_if_changed(stage / name, (here / name).read_text(encoding="utf-8"))
    write_if_changed(stage / "compat.h", COMPAT)


def build(stage, cxx):
    obj = stage / "obj"
    obj.mkdir(exist_ok=True)
    flags = ["-std=gnu++20", "-O1", "-include", str(stage / "compat.h"), "-I", str(stage), "-I", str(stage / "el512")]
    objects = []
    # ESPHome's own sources: compiled once per ESPHome version.
    for src in ESPHOME_SOURCES:
        o = obj / (src.replace("/", "_") + ".o")
        if not o.exists():
            run([cxx, *flags, "-w", "-c", str(stage / "esphome" / src), "-o", str(o)])
        objects.append(o)
    # Ours: rebuilt when they change.
    for src in [stage / "fonts.cpp", stage / "render_main.cpp", stage / "host_stubs.cpp"] + [stage / "el512" / f"{p}.cpp" for p in PAGE_SOURCES]:
        o = obj / (src.stem + ".o")
        deps = [src] + list((stage / "el512").glob("*.h")) + [stage / "fonts.h"]
        if not o.exists() or any(d.stat().st_mtime > o.stat().st_mtime for d in deps):
            warn = ["-w"] if src.stem == "fonts" else ["-Wall"]
            out = run([cxx, *flags, *warn, "-c", str(src), "-o", str(o)])
            sys.stdout.write(out)
        objects.append(o)
    exe = stage / ("render.exe" if sys.platform == "win32" else "render")
    if not exe.exists() or any(o.stat().st_mtime > exe.stat().st_mtime for o in objects):
        run([cxx, *[str(o) for o in objects], "-o", str(exe), "-static"])
    return exe


def to_png(pbm, png, scale):
    from PIL import Image

    # PBM stores lit pixels as 1, which Pillow reads as black: invert for the mask.
    lit = Image.open(pbm).convert("L").point(lambda v: 255 - v)
    img = Image.new("RGB", lit.size, BACKGROUND)
    img.paste(AMBER, mask=lit)
    if scale != 1:
        img = img.resize((img.width * scale, img.height * scale), Image.NEAREST)
    img.save(png)
    pbm.unlink()


def contact_sheet(pngs, path, cols=2):
    from PIL import Image, ImageDraw

    imgs = [Image.open(p) for p in pngs]
    w, h = imgs[0].size
    label = 18
    sheet = Image.new("RGB", (cols * (w + 8) + 8, ((len(imgs) + cols - 1) // cols) * (h + label + 8) + 8), (0, 0, 0))
    draw = ImageDraw.Draw(sheet)
    for i, (p, im) in enumerate(zip(pngs, imgs)):
        x, y = 8 + (i % cols) * (w + 8), 8 + (i // cols) * (h + label + 8)
        draw.text((x, y), p.stem, fill=(180, 120, 20))
        sheet.paste(im, (x, y + label))
    sheet.save(path)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--page", action="append", help="page name (repeatable; default: all)")
    ap.add_argument("--at", help="moments in ms after the page was shown, e.g. 0,2000 (default: per page)")
    ap.add_argument("--scenario", default="normal", choices=["normal", "empty", "alert", "busy"])
    ap.add_argument("--out", type=Path, default=WORK / "out")
    ap.add_argument("--scale", type=int, default=2)
    ap.add_argument("--sheet", action="store_true", help="also write sheet.png with all renders")
    ap.add_argument("--list", action="store_true", help="list pages and their default moments")
    ap.add_argument("--examples", action="store_true",
                    help="one image per page (its first moment) as docs/images/pages/<page>.png")
    ap.add_argument("--build-dir", type=Path, default=ROOT / "esphome/.esphome/build/el512-panel",
                    help="ESPHome build folder with src/main.cpp")
    ap.add_argument("--cxx", default=shutil.which("g++") or "g++")
    args = ap.parse_args()

    main_cpp = args.build_dir / "src" / "main.cpp"
    if not main_cpp.exists():
        raise SystemExit(f"{main_cpp} not found: run `esphome compile esphome/el512-panel-local.yaml` first")
    WORK.mkdir(exist_ok=True)
    stage_esphome(WORK)
    extract_fonts(main_cpp, WORK)
    stage_pages(WORK)
    exe = build(WORK, args.cxx)

    listing = {}
    for line in run([str(exe), "list"]).splitlines():
        name, _, moments = line.partition(" ")
        listing[name] = [int(m) for m in moments.split(",") if m]
    if args.list:
        for name, moments in listing.items():
            print(f"{name:14} {', '.join(map(str, moments))}")
        return
    if args.examples:
        args.out, args.at = ROOT / "docs/images/pages", None
        listing = {n: m[:1] for n, m in listing.items()}
    names = args.page or list(listing)
    for n in names:
        if n not in listing:
            raise SystemExit(f"unknown page {n}; pages: {', '.join(listing)}")
    jobs = [f"{n}:{m}" for n in names
            for m in ([int(x) for x in args.at.split(",")] if args.at else listing[n])]
    args.out.mkdir(parents=True, exist_ok=True)
    pngs = []
    for pbm in run([str(exe), str(args.out), args.scenario, *jobs]).splitlines():
        pbm = Path(pbm)
        stem = pbm.stem.rsplit("_", 1)[0] if args.examples else pbm.stem
        png = pbm.with_name(f"{stem}{'' if args.scenario == 'normal' else '_' + args.scenario}.png")
        to_png(pbm, png, args.scale)
        pngs.append(png)
    if args.sheet and pngs:
        contact_sheet(pngs, args.out / f"sheet_{args.scenario}.png")
        pngs.append(args.out / f"sheet_{args.scenario}.png")
    print(f"{len(pngs)} image(s) in {args.out}")


if __name__ == "__main__":
    main()