# SPDX-License-Identifier: Apache-2.0
"""Cut the printed meshes with a plane and draw the section, with the parts' layout boxes it crosses.

    python v01/cad/section.py posts            # one section through each board post (along x)
    python v01/cad/section.py icon             # through the NFC icon and the post beside it
    python v01/cad/section.py eye              # through the eye, its cradle and pins
    python v01/cad/section.py --y -0.5         # any plane: --x, --y or --z at that coordinate
    python v01/cad/section.py --y -0.5 --around=-40,15 --half 14   # zoom (use = before a negative): centre (along, z) and half-width, mm

Writes out_pebble/sections/<name>.svg (and .png on macOS, through Quick Look). Shells are black, other printed
parts grey, layout boxes dashed (keep-out spaces red). Reads the exported STLs, so it shows what will print.
"""
from __future__ import annotations

import argparse
import json
import shutil
import subprocess
from pathlib import Path

import numpy as np
import trimesh
from trimesh.intersections import mesh_plane

OUT = Path(__file__).parent / "out_pebble"
DIR = OUT / "sections"
PARTS = ["top_shell", "bottom_shell", "battery_deck", "board_straps", "bumper_band", "eye_guide", "eye_lens", "eye_cradle", "power_cap", "mic_cap", "foot"]
AX = {"x": 0, "y": 1, "z": 2}


def cut(axis, at):
    """{part: segments as an (n, 2, 2) array in the plane's 2D coordinates}, and the 2D axes' names."""
    n = np.zeros(3); n[AX[axis]] = 1
    o = np.zeros(3); o[AX[axis]] = at
    keep = [i for i in range(3) if i != AX[axis]]
    if axis == "z":
        keep = [0, 1]
    out = {}
    for p in PARTS:
        f = OUT / f"{p}.stl"
        if f.exists():
            seg = mesh_plane(trimesh.load(f), n, o)
            if len(seg):
                out[p] = seg[:, :, keep]
    return out, ["xyz"[i] for i in keep]


def boxes(axis, at, keep):
    L = json.loads((OUT / "layout.json").read_text())
    a = AX[axis]
    out = []
    for c in L["comps"]:
        if c["at"][a] <= at <= c["at"][a] + c["size"][a]:
            i, j = (AX[k] for k in keep)
            out.append((c["name"], c.get("role"), c["at"][i], c["at"][j], c["size"][i], c["size"][j]))
    return out


def svg(name, axis, at, around=None, half=None, title=""):
    segs, keep = cut(axis, at)
    bx = boxes(axis, at, keep)
    allp = np.concatenate([s.reshape(-1, 2) for s in segs.values()]) if segs else np.zeros((1, 2))
    if around is not None:
        (cu, cv), h = around, half or 12.0
        u0, u1, v0, v1 = cu - h, cu + h, cv - h * 0.66, cv + h * 0.66
    else:
        (u0, v0), (u1, v1) = allp.min(0) - 3, allp.max(0) + 3
    W, H = u1 - u0, v1 - v0
    px = 1200 / W
    sw = 1.2 / px                                         # 1.2 px lines at any zoom
    X = lambda u: u - u0
    Y = lambda v: v1 - v                                  # SVG y runs down
    g = [f'<rect x="0" y="0" width="{W:.3f}" height="{H:.3f}" fill="white"/>']
    for name_, role, bu, bv, su, sv in bx:
        red = role == "keep-out"
        g.append(f'<rect x="{X(bu):.3f}" y="{Y(bv + sv):.3f}" width="{su:.3f}" height="{sv:.3f}" fill="{"#d0342c" if red else "#2f6fd6"}" fill-opacity=".06" '
                 f'stroke="{"#d0342c" if red else "#2f6fd6"}" stroke-width="{sw * .8:.3f}" stroke-dasharray="{sw * 4:.3f} {sw * 3:.3f}"><title>{name_}</title></rect>')
    for p, s in segs.items():
        col = "#111" if p.endswith("shell") else "#777"
        d = " ".join(f"M{X(a[0]):.3f} {Y(a[1]):.3f}L{X(b[0]):.3f} {Y(b[1]):.3f}" for a, b in s   # any segment that crosses the window
                     if min(a[0], b[0]) <= u1 and max(a[0], b[0]) >= u0 and min(a[1], b[1]) <= v1 and max(a[1], b[1]) >= v0)
        g.append(f'<path d="{d}" stroke="{col}" stroke-width="{sw:.3f}" fill="none"><title>{p}</title></path>')
    label = title or f"{axis} = {at:g} mm"
    g.append(f'<text x="{sw * 6:.3f}" y="{sw * 16:.3f}" font-size="{sw * 13:.3f}" font-family="Helvetica, Arial">{label} · {keep[0]} → , {keep[1]} ↑ · '
             f'1 mm = {px:.0f} px</text>')
    DIR.mkdir(exist_ok=True)
    f = DIR / f"{name}.svg"
    f.write_text(f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {W:.3f} {H:.3f}" width="1200" height="{int(1200 * H / W)}">' + "".join(g) + "</svg>")
    if shutil.which("qlmanage"):
        subprocess.run(["qlmanage", "-t", "-s", "1200", "-o", str(DIR), str(f)], capture_output=True)
        png = DIR / f"{name}.svg.png"
        if png.exists():
            png.replace(DIR / f"{name}.png")
    return f


def presets(which):
    L = json.loads((OUT / "layout.json").read_text())
    jobs = []
    if which == "posts":
        for f in L["info"].get("post_flares", []):
            cx, cy = f["c"]
            jobs.append((f"post-{f['post'].split()[1]}", "y", cy, (cx, f["z_ref"]), 12.0, f"{f['post']}: section along x through its axis"))
    elif which == "icon":
        ix, iy = L["nfc_icon"]
        jobs.append(("icon", "y", iy, (ix + 6, L["info"]["post_flares"][0]["z_ref"] + 2), 16.0, "NFC icon and the post beside it, along x"))
    elif which == "eye":
        e = L["eye"]
        jobs.append(("eye", "y", 0.0, (e["x"], e["led_top"]), 18.0, "eye: guide, lens, cradle, along x"))
    return jobs


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("preset", nargs="?", choices=["posts", "icon", "eye"])
    for a in "xyz":
        ap.add_argument(f"--{a}", type=float)
    ap.add_argument("--around", help="zoom centre: along,up in the section's own axes (mm)")
    ap.add_argument("--half", type=float, help="zoom half-width (mm)")
    a = ap.parse_args()
    if a.preset:
        jobs = presets(a.preset)
    else:
        axis = next((k for k in "xyz" if getattr(a, k) is not None), None)
        if axis is None:
            ap.error("give a preset or one of --x/--y/--z")
        around = tuple(float(v) for v in a.around.split(",")) if a.around else None
        jobs = [(f"{axis}{getattr(a, axis):g}", axis, getattr(a, axis), around, a.half, "")]
    for name, axis, at, around, half, title in jobs:
        print(svg(name, axis, at, around, half, title))


if __name__ == "__main__":
    main()