#!/usr/bin/env python3
"""One-command panel-descriptor capture.

Resolves a GxEPD2 driver class to its vendored source, builds and runs a host
harness that *observes* the driver, and prints a candidate EpdPanelDescriptor
(the body of a `components/gxepd2/panel_<model>.h`). No ESP toolchain needed -
just a host C++ compiler. See tools/gxepd2_harness/README.md and CONTRIBUTING.md.

    tools/capture.py GxEPD2_154_Z90c             # candidate -> stdout
    tools/capture.py GxEPD2_154_Z90c -o panel.h  # candidate -> file
    tools/capture.py --list                      # capturable driver classes
    tools/capture.py GxEPD2_154_Z90c --debug     # also dump the capture structure

The output is a *candidate*: review it, fill any TODO fields still left
(plane_width/line_order for dual-region), save it as
components/gxepd2/panel_<model>.h, add it to display.py's MODELS, compile,
and bench-verify before marking it verified.
4-colour (epd4c/) and 7-colour (epd7c/) panels have
bits_per_pixel/palette/blank_code auto-filled by observing the vendor
GxEPD2_4C/GxEPD2_7C wrapper (see gxepd2_harness/emit_palette_4c.h /
emit_palette_7c.h).
"""

import argparse
import difflib
import json
import pathlib
import re
import subprocess
import sys
import tempfile

ROOT = pathlib.Path(__file__).resolve().parent.parent  # code/
CATALOG = ROOT / "tools" / "gxepd2_catalog" / "catalog.json"
VSRC = ROOT / "third_party" / "GxEPD2" / "src"
HARNESS = ROOT / "tools" / "gxepd2_harness"
DISPLAY_PY = ROOT / "components" / "gxepd2" / "display.py"
PREFIX = "GxEPD2_"


def macro_of(model: str) -> str:
    return "PANEL_" + model[len(PREFIX):].upper()


def filename_of(model: str) -> str:
    return "panel_" + model[len(PREFIX):].lower() + ".h"


def _fmt_of(src_rel: str) -> str:
    """Panel colour format from the vendored source subdir (epd3c/epd4c/epd7c/
    gdem3c/...). Everything else is mono. The directory is GxEPD2's own
    classification, so this never guesses."""
    sub = src_rel.replace("\\", "/").split("/")[0].lower()
    if "7c" in sub:
        return "7-color"
    if "4c" in sub:
        return "4-color"
    if "3c" in sub:
        return "3-color"
    return "mono"


def _header_meta(hdr_path: pathlib.Path) -> tuple:
    """(part, controller) scraped from the driver header's `// Panel:` /
    `// Controller:` comment lines, or ("TODO", "TODO") when absent/unusable.
    `part` is only trusted when it reads as a Good Display GDExNNNN part number
    (the silkscreen label); a prose description (e.g. 565c's "5.65inch ACeP")
    or a missing line stays TODO for the dev to fill from the physical panel."""
    part, controller = "TODO", "TODO"
    try:
        text = hdr_path.read_text()
    except OSError:
        return part, controller
    m = re.search(r"^//\s*Panel:\s*(\S+)", text, re.M)
    if m and re.match(r"GDE[A-Z][0-9A-Z]+$", m.group(1)):
        part = m.group(1)
    m = re.search(r"^//\s*Controller:\s*([^,:\s]+)", text, re.M)
    if m and m.group(1).lower() != "unknown":
        controller = m.group(1)
    return part, controller


def registry_line(model: str, src_rel: str) -> str:
    """The exact `PANELS` entry to paste into display.py - every mechanical
    field (key, descriptor macro, format) filled so it can't be mistyped, with
    part/controller scraped where the header allows and `status` starting at
    "structural" (a human flips it to "verified" after the bench)."""
    part, controller = _header_meta(VSRC / (src_rel[:-4] + ".h"))
    return (f'    "{model}": Panel("{macro_of(model)}", '
            f'"{part}", "{controller}", "{_fmt_of(src_rel)}", "structural"),')


def shipped_models() -> set:
    """Model keys already wired into display.py's MODELS, to flag in --list."""
    try:
        text = DISPLAY_PY.read_text()
    except OSError:
        return set()
    return set(re.findall(r'"(GxEPD2_[^"]+)":\s*Panel\(', text))


def load_catalog() -> dict:
    if not CATALOG.exists():
        sys.exit(f"catalog not found: {CATALOG}\n(run tools/gxepd2_catalog/generate_catalog.py)")
    return json.loads(CATALOG.read_text())


def capturable(entry: dict) -> bool:
    # The emitter instantiates Driver(cs, dc, rst, busy) - only the standard
    # 4-arg constructor works without special-casing.
    return entry.get("ctor_args") == 4 and (VSRC / entry["source"]).exists()


def do_list(catalog: dict) -> None:
    shipped = shipped_models()
    rows = []
    for model, entry in sorted(catalog.items()):
        if not capturable(entry):
            continue
        mark = "shipped" if model in shipped else "-"
        rows.append((model, mark, entry.get("comment", "")))
    w = max(len(m) for m, _, _ in rows)
    print(f"{len(rows)} capturable driver classes (ctor_args=4, source present):\n")
    for model, mark, comment in rows:
        print(f"  {model:<{w}}  {mark:<7}  {comment}")


def resolve_source(model: str, catalog: dict):
    """(source relpath, ctor_args|None, real_class_name) or (None, None, None).

    Catalog first (it has ctor_args), else glob the vendored tree - so a driver
    the fork added but GxEPD2's example headers don't list (e.g. the F52) is
    still capturable. Matching is case-insensitive and returns the driver's
    REAL class name (the `model:` key may differ in case, e.g. GxEPD2_213_Z98C
    vs the class GxEPD2_213_Z98c)."""
    low = model.lower()
    for k, e in catalog.items():
        if k.lower() == low:
            return e["source"], e.get("ctor_args"), e.get("driver", k)
    hits = [p for p in VSRC.glob("**/GxEPD2_*.cpp") if p.stem.lower() == low]
    if len(hits) == 1:
        return str(hits[0].relative_to(VSRC)), None, hits[0].stem
    if len(hits) > 1:
        sys.exit(f"{model}: multiple sources found: {[str(h) for h in hits]}")
    return None, None, None


def capture(model: str, catalog: dict, debug: bool, cxx: str, keep: bool) -> str:
    src_rel, ctor_args, real_class = resolve_source(model, catalog)
    if src_rel is None:
        names = set(catalog) | {p.stem for p in VSRC.glob("**/GxEPD2_*.cpp")}
        hint = difflib.get_close_matches(model, names, n=3)
        msg = f"unknown model: {model} (no {model}.cpp under {VSRC})"
        if hint:
            msg += "\ndid you mean: " + ", ".join(hint)
        sys.exit(msg)
    if ctor_args not in (None, 4):
        sys.exit(f"{model}: non-standard constructor (ctor_args={ctor_args}); "
                 "the emitter only handles the 4-arg GxEPD2_EPD(cs,dc,rst,busy) form.")
    src = VSRC / src_rel
    hdr_rel = src_rel[:-4] + ".h"  # epd3c/GxEPD2_154_Z90c.h
    if not (VSRC / hdr_rel).exists():
        sys.exit(f"driver header not found: {VSRC / hdr_rel}")

    macro = macro_of(model)
    # 4-colour (epd4c/) panels: emit via the 4C entry, which observes the
    # GxEPD2_4C wrapper to fill bits_per_pixel/palette/blank_code instead of
    # leaving them as TODO. The wrapper references WIDTH/HEIGHT unqualified, so
    # those must be defined before GxEPD2_4C.h (pulled in by emit_palette_4c.h).
    is_4c = src_rel.replace("\\", "/").startswith("epd4c/")
    # 7-colour (epd7c/) panels: same approach as 4-colour, via the GxEPD2_7C
    # wrapper (see gxepd2_harness/emit_palette_7c.h).
    is_7c = src_rel.replace("\\", "/").startswith("epd7c/")
    if is_4c:
        main_src = (
            f'#include "{hdr_rel}"\n'
            f'constexpr int16_t WIDTH = {real_class}::WIDTH;\n'
            f'constexpr int16_t HEIGHT = {real_class}::HEIGHT;\n'
            f'#include "emit_palette_4c.h"\n'
            f'int main() {{ emit::emit_descriptor_4c<{real_class}>("{real_class}", "{macro}"); return 0; }}\n'
        )
    elif is_7c:
        main_src = (
            f'#include "{hdr_rel}"\n'
            f'constexpr int16_t WIDTH = {real_class}::WIDTH;\n'
            f'constexpr int16_t HEIGHT = {real_class}::HEIGHT;\n'
            f'#include "emit_palette_7c.h"\n'
            f'int main() {{ emit::emit_descriptor_7c<{real_class}>("{real_class}", "{macro}"); return 0; }}\n'
        )
    else:
        main_src = (
            f'#include "emit_descriptor.h"\n'
            f'#include "{hdr_rel}"\n'
            f'int main() {{ emit::emit_descriptor<{real_class}>("{real_class}", "{macro}"); return 0; }}\n'
        )
    with tempfile.TemporaryDirectory() as td:
        tdp = pathlib.Path(td)
        main_cpp = tdp / "emit_main.cpp"
        main_cpp.write_text(main_src)
        binary = tdp / "emit_main"
        compile_cmd = [
            cxx, "-std=c++17",
            "-I", str(HARNESS / "spy"), "-I", str(VSRC), "-I", str(HARNESS),
            str(main_cpp), str(VSRC / "GxEPD2_EPD.cpp"), str(src),
            "-o", str(binary),
        ]
        print(f"[capture] {model}  <-  third_party/GxEPD2/src/{src_rel}", file=sys.stderr)
        cc = subprocess.run(compile_cmd, capture_output=True, text=True)
        if cc.returncode != 0:
            sys.exit(f"compile failed:\n{cc.stderr}")
        env = {"EMIT_DEBUG": "1"} if debug else None
        run = subprocess.run([str(binary)], capture_output=True, text=True, env=env)
        if keep:
            kept = HARNESS / f"run_{model[len(PREFIX):].lower()}_emit.cpp"
            kept.write_text(main_cpp.read_text())
            print(f"[capture] kept emit main: {kept}", file=sys.stderr)
    if debug and run.stderr:
        print(run.stderr, file=sys.stderr)
    if run.returncode != 0:
        sys.exit(f"capture run failed (exit {run.returncode}):\n{run.stdout}\n{run.stderr}")

    marker = "// GENERATED by emit_descriptor"
    idx = run.stdout.find(marker)
    if idx < 0:
        sys.exit(f"no descriptor in output:\n{run.stdout}")
    descriptor = run.stdout[idx:]
    print(f"[capture] ok -> save as components/gxepd2/{filename_of(model)}", file=sys.stderr)
    print("[capture] then paste this line into display.py's PANELS. Key, descriptor "
          "macro and format are exact; fill any TODO and sanity-check part/controller "
          "(scraped from the driver header) against the physical panel:", file=sys.stderr)
    print(registry_line(model, src_rel), file=sys.stderr)
    return descriptor


# mode -> (harness header to include, template to invoke). Both headers are
# templates over Driver and take (name) - see run_harness.h / run_diag_harness.h.
_PROGRAMS = {
    "sweep": ("run_harness.h", "run_harness"),
    "diag": ("run_diag_harness.h", "run_diag_harness"),
}


def run_program(model: str, catalog: dict, mode: str, cxx: str) -> str:
    """Build+run run_harness.h ('sweep': full capture + the window-encoder
    sweep) or run_diag_harness.h ('diag': Pattern A row/byte-order fill) for
    `model`, with no per-panel run_<panel>.cpp to hand-write - same resolution
    `capture()` uses. Previously every panel needed its own 4-line
    run_<panel>[_diag].cpp; both headers are already templates over Driver, so
    a single generic main works for any of them."""
    src_rel, ctor_args, real_class = resolve_source(model, catalog)
    if src_rel is None:
        sys.exit(f"unknown model: {model} (no {model}.cpp under {VSRC})")
    if ctor_args not in (None, 4):
        sys.exit(f"{model}: non-standard constructor (ctor_args={ctor_args}); "
                 "the harness only handles the 4-arg GxEPD2_EPD(cs,dc,rst,busy) form.")
    src = VSRC / src_rel
    hdr_rel = src_rel[:-4] + ".h"
    if not (VSRC / hdr_rel).exists():
        sys.exit(f"driver header not found: {VSRC / hdr_rel}")
    if mode == "power":
        # emit_descriptor.h's print_power_sequences, namespaced and keyed by
        # macro (not name) like emit_descriptor() itself - not run_harness.h/
        # run_diag_harness.h's plain-global, name-keyed shape, so it doesn't
        # fit _PROGRAMS below.
        main_src = (
            f'#include "emit_descriptor.h"\n'
            f'#include "{hdr_rel}"\n'
            f'int main() {{ emit::print_power_sequences<{real_class}>("{macro_of(model)}"); return 0; }}\n'
        )
    else:
        header, fn = _PROGRAMS[mode]
        main_src = (
            f'#include "{header}"\n'
            f'#include "{hdr_rel}"\n'
            f'int main() {{ {fn}<{real_class}>("{real_class}"); return 0; }}\n'
        )
    with tempfile.TemporaryDirectory() as td:
        tdp = pathlib.Path(td)
        main_cpp = tdp / "main.cpp"
        main_cpp.write_text(main_src)
        binary = tdp / "main"
        compile_cmd = [
            cxx, "-std=c++17",
            "-I", str(HARNESS / "spy"), "-I", str(VSRC), "-I", str(HARNESS),
            str(main_cpp), str(VSRC / "GxEPD2_EPD.cpp"), str(src),
            "-o", str(binary),
        ]
        print(f"[{mode}] {model}  <-  third_party/GxEPD2/src/{src_rel}", file=sys.stderr)
        cc = subprocess.run(compile_cmd, capture_output=True, text=True)
        if cc.returncode != 0:
            sys.exit(f"compile failed:\n{cc.stderr}")
        run = subprocess.run([str(binary)], capture_output=True, text=True)
    if run.returncode != 0:
        sys.exit(f"{mode} run failed (exit {run.returncode}):\n{run.stdout}\n{run.stderr}")
    return run.stdout


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("model", nargs="?", help="GxEPD2 driver class, e.g. GxEPD2_154_Z90c")
    ap.add_argument("-o", "--out", help="write the candidate descriptor to this file (default: stdout)")
    ap.add_argument("--list", action="store_true", help="list capturable driver classes and exit")
    ap.add_argument("--debug", action="store_true", help="dump the image-stage command/burst structure to stderr")
    ap.add_argument("--keep", action="store_true", help="keep the generated emit main under tools/gxepd2_harness/")
    ap.add_argument("--sweep", action="store_true",
                     help="run the window-encoder sweep (run_harness.h) instead of emitting a descriptor")
    ap.add_argument("--diag", action="store_true",
                     help="run the Pattern-A row/byte-order diagnostic fill (run_diag_harness.h) instead of emitting")
    ap.add_argument("--power", action="store_true",
                     help="print just power_off_sequence/hibernate_sequence (for splicing into an existing "
                          "panel_*.h), instead of emitting a full descriptor")
    ap.add_argument("--cxx", default="g++", help="C++ compiler (default: g++)")
    args = ap.parse_args()

    catalog = load_catalog()
    if args.list:
        do_list(catalog)
        return
    if not args.model:
        ap.error("a model is required (or use --list)")
    if args.sweep or args.diag or args.power:
        mode = "sweep" if args.sweep else "diag" if args.diag else "power"
        out = run_program(args.model, catalog, mode, args.cxx)
        if args.out:
            pathlib.Path(args.out).write_text(out)
            print(f"[{mode}] wrote {args.out}", file=sys.stderr)
        else:
            sys.stdout.write(out)
        return
    descriptor = capture(args.model, catalog, args.debug, args.cxx, args.keep)
    if args.out:
        pathlib.Path(args.out).write_text(descriptor)
        print(f"[capture] wrote {args.out}", file=sys.stderr)
    else:
        sys.stdout.write(descriptor)


if __name__ == "__main__":
    main()