FacereDataset/tools/epro2/kicad/__main__.py

"""CLI: convert EPRO2 SCH_PAGE docs to KiCad ``.kicad_sch`` files.

Usage:
    uv run python -m tools.epro2.kicad <project_dir> --doc <sch_uuid> --out <dir>
    uv run python -m tools.epro2.kicad <project_dir> --all-sch --out <dir>
    uv run python -m tools.epro2.kicad <project_dir> --all-sch --flat --out <dir>

By default ``--all-sch`` produces a hierarchical project: pages are grouped
by their parent SCH (``META.schematic``) and each group gets a directory
with a root ``.kicad_sch`` plus its child pages. ``kicad-cli sch erc`` on
the root then resolves cross-page nets via ``(global_label)``. Pass
``--flat`` to fall back to one file per page in a single directory (the
old behaviour, still useful when you only care about a single page).
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

from ..project_relations import ProjectRelations
from ..replay import Document, Project, replay_project
from .pcb_writer import write_pcb
from .root_sch_writer import ChildSheet, new_sheet_uuid, write_root_sheet
from .sch_writer import write_sch_page


_SAFE_CHARS = re.compile(r"[^A-Za-z0-9._\-一-鿿]+")


def _safe_filename(s: str) -> str:
    s = _SAFE_CHARS.sub("_", s).strip("_")
    return s or "untitled"


def _page_title(doc: Document) -> str:
    return (doc.objects.get("META") or {}).get("title") or doc.doc_uuid[:12]


def _print_stats(name: str) -> None:
    stats = getattr(write_sch_page, "last_stats", None)
    if stats:
        print(
            f"  {name}: wires={stats.wires} symbols={stats.symbol_placements} "
            f"text={stats.text} labels={stats.labels} skipped={stats.skipped} "
            f"lib_emb={stats.lib_symbols_embedded} lib_miss={stats.lib_symbols_missing}"
        )


def _convert_one_flat(
    proj: Project,
    doc_uuid: str,
    out_dir: Path,
    pr: ProjectRelations | None = None,
) -> Path:
    """Standalone single-page export — no hierarchy. Useful for inspecting one page."""
    if doc_uuid not in proj.documents:
        candidates = [u for u in proj.documents if u.startswith(doc_uuid)]
        if len(candidates) != 1:
            raise SystemExit(f"no unique match for {doc_uuid!r} (candidates: {candidates[:5]})")
        doc_uuid = candidates[0]
    doc = proj.documents[doc_uuid]
    if doc.doc_type != "SCH_PAGE":
        raise SystemExit(f"doc {doc_uuid} is {doc.doc_type!r}, not SCH_PAGE")

    text = write_sch_page(doc, project_relations=pr)
    out_path = out_dir / f"{_safe_filename(_page_title(doc))}_{doc_uuid[:8]}.kicad_sch"
    out_path.write_text(text, encoding="utf-8")
    _print_stats(out_path.name)
    return out_path


def _group_pages_by_sch(proj: Project) -> dict[str, list[str]]:
    """Group SCH_PAGE doc uuids by their parent SCH uuid (META.schematic).

    Pages without a META.schematic land in the bucket keyed by their own
    uuid (treated as a single-page schematic). Deleted SCHs (``DELETE_DOC``)
    are skipped — their pages are dropped silently.
    """
    deleted_schs = {
        u for u, d in proj.documents.items()
        if d.doc_type == "SCH" and (d.objects.get("DELETE_DOC") or {}).get("isDelete")
    }
    out: dict[str, list[str]] = {}
    for u, d in proj.documents.items():
        if d.doc_type != "SCH_PAGE":
            continue
        meta = d.objects.get("META") or {}
        sch_uuid = str(meta.get("schematic") or u)
        if sch_uuid in deleted_schs:
            continue
        out.setdefault(sch_uuid, []).append(u)
    # Sort pages within each SCH by zIndex (the editor's tab order)
    for sch_uuid, pages in out.items():
        pages.sort(key=lambda u: (proj.documents[u].objects.get("META") or {}).get("zIndex") or 0)
    return out


def _convert_hierarchical(
    proj: Project,
    out_dir: Path,
    pr: ProjectRelations | None,
) -> int:
    """Emit one root + N children per SCH group. Returns root count."""
    groups = _group_pages_by_sch(proj)
    if not groups:
        print("no SCH_PAGE docs in this project", file=sys.stderr)
        return 0

    root_count = 0
    for sch_uuid, page_uuids in groups.items():
        sch_doc = proj.documents.get(sch_uuid)
        sch_title = (
            (sch_doc.objects.get("META") or {}).get("title")
            if sch_doc and sch_doc.doc_type == "SCH"
            else None
        ) or _page_title(proj.documents[page_uuids[0]])

        sch_dir = out_dir / _safe_filename(sch_title)
        sch_dir.mkdir(parents=True, exist_ok=True)
        print(f"Schematic '{sch_title}' ({len(page_uuids)} page(s)) → {sch_dir}")

        # Pre-assign each child a uuid and a filename; the child page
        # writer needs the uuid to declare its hierarchy path, and the
        # root needs both to emit its (sheet ...) entries.
        children: list[ChildSheet] = []
        for page_uuid in page_uuids:
            page_doc = proj.documents[page_uuid]
            title = _page_title(page_doc)
            children.append(ChildSheet(
                filename=f"{_safe_filename(title)}_{page_uuid[:8]}.kicad_sch",
                title=title,
                sheet_uuid=new_sheet_uuid(),
            ))

        for idx, (page_uuid, child) in enumerate(zip(page_uuids, children)):
            page_doc = proj.documents[page_uuid]
            text = write_sch_page(
                page_doc,
                project_relations=pr,
                sheet_path=f"/{child.sheet_uuid}",
                page_num=idx + 2,  # root takes page 1
            )
            (sch_dir / child.filename).write_text(text, encoding="utf-8")
            _print_stats(child.filename)

        root_text = write_root_sheet(sch_title, children, project_name=_safe_filename(sch_title))
        root_path = sch_dir / f"{_safe_filename(sch_title)}.kicad_sch"
        root_path.write_text(root_text, encoding="utf-8")
        print(f"  ROOT: {root_path}")
        root_count += 1
    return root_count


def _convert_all_pcb(proj: Project, out_dir: Path, pr: ProjectRelations) -> int:
    """Emit one .kicad_pcb per PCB doc. Each is named after its META.title
    and dropped into a sibling directory of the matching SCH for parity
    with the schematic export layout."""
    pcb_uuids = [u for u, d in proj.documents.items() if d.doc_type == "PCB"]
    if not pcb_uuids:
        print("no PCB docs in this project", file=sys.stderr)
        return 0
    print(f"Converting {len(pcb_uuids)} PCB doc(s) → {out_dir}")
    n = 0
    for u in pcb_uuids:
        doc = proj.documents[u]
        title = (doc.objects.get("META") or {}).get("title") or u[:12]
        try:
            text = write_pcb(doc, project_relations=pr)
        except Exception as e:  # noqa: BLE001
            print(f"  FAIL {u[:12]}: {e}", file=sys.stderr)
            continue
        out_path = out_dir / f"{_safe_filename(title)}.kicad_pcb"
        out_path.write_text(text, encoding="utf-8")
        stats = getattr(write_pcb, "last_stats", None)
        if stats:
            print(
                f"  {out_path.name}: nets={stats.nets} fps={stats.footprints} "
                f"fps_unresolved={stats.footprints_unresolved} "
                f"segments={stats.segments} vias={stats.vias} "
                f"zones={stats.zones} edge={stats.edge_cuts}"
            )
        n += 1
    return n


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="EPRO2 → KiCad schematic + PCB exporter")
    ap.add_argument("project_dir", type=Path)
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--doc", help="SCH_PAGE doc uuid (or unique prefix) to convert")
    g.add_argument("--all-sch", action="store_true", help="convert every SCH_PAGE")
    g.add_argument("--all-pcb", action="store_true", help="convert every PCB doc to .kicad_pcb")
    ap.add_argument("--out", type=Path, default=Path("data/processed/kicad_sch"))
    ap.add_argument(
        "--no-lib-symbols",
        action="store_true",
        help="skip lib_symbols generation (Phase 1 mode — placements only, red ?)",
    )
    ap.add_argument(
        "--flat",
        action="store_true",
        help="emit one flat .kicad_sch per page (no hierarchy); only with --all-sch or --doc",
    )
    args = ap.parse_args(argv)

    proj = replay_project(args.project_dir)
    args.out.mkdir(parents=True, exist_ok=True)
    pr = None if args.no_lib_symbols else ProjectRelations.build(proj)

    if args.all_pcb:
        if pr is None:
            pr = ProjectRelations.build(proj)
        _convert_all_pcb(proj, args.out, pr=pr)
        return 0

    if args.doc:
        _convert_one_flat(proj, args.doc, args.out, pr=pr)
        return 0

    if args.flat:
        targets = [u for u, d in proj.documents.items() if d.doc_type == "SCH_PAGE"]
        if not targets:
            print("no SCH_PAGE docs in this project", file=sys.stderr)
            return 1
        print(f"Flat mode: converting {len(targets)} SCH_PAGE docs → {args.out}")
        for u in targets:
            try:
                _convert_one_flat(proj, u, args.out, pr=pr)
            except Exception as e:  # noqa: BLE001
                print(f"  FAIL {u[:12]}: {e}", file=sys.stderr)
        return 0

    _convert_hierarchical(proj, args.out, pr=pr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())