periscope/backend/services/pcb_pipeline.py

"""PCB review pipeline — parallel to analysis (exam, not auto-place).

Stages: ensure_graph → parse_pcb → classify → inventory → checks →
ai_review → write_report. Uses ``pcb_status`` so analysis ``status`` is
untouched. Does not write ``.kicad_pcb`` or packing coordinates.
"""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from pathlib import Path

from backend.periscopex.cad_bridge import (
    annotate_findings_cad,
    build_cad_bridge,
    cad_index_from_graph,
    write_cad_bridge,
)
from backend.periscopex.finding_engine import apply_decisions
from backend.periscopex.graph import build_graph
from backend.periscopex.models import ComponentConstraints, DesignGraph, LayoutGraph, ValidationReport
from backend.periscopex.pcb_checks import assign_pcb_finding_ids, run_pcb_checks
from backend.periscopex.pcb_inventory import build_pcb_inventory
from backend.services import projects as proj_svc
from backend.services.pipeline import PipelineWorkspace, broker
from backend.services.storage import StorageBackend

logger = logging.getLogger(__name__)

_PCB_ACTIVE = frozenset({"queued", "running"})
_ANALYSIS_BUSY = frozenset({
    proj_svc.STATUS_QUEUED,
    proj_svc.STATUS_RUNNING,
})
_PLACEMENT_ACTIVE = frozenset({"queued", "running"})


def _load_constraints_map(
    extracted_dir: Path,
    storage: StorageBackend | None = None,
) -> dict[str, ComponentConstraints]:
    """Project extracted/ plus shared ``library/extracted`` (one store for schema+PCB)."""
    result: dict[str, ComponentConstraints] = {}
    if extracted_dir.is_dir():
        for f in extracted_dir.glob("*.json"):
            try:
                c = ComponentConstraints.model_validate_json(
                    f.read_text(encoding="utf-8"),
                )
            except Exception:
                logger.exception("skipping bad extraction %s", f)
                continue
            result[c.mpn] = c
    if storage is None:
        return result
    try:
        keys = storage.list_recursive("library/extracted/")
    except Exception:
        logger.exception("listing library/extracted failed")
        return result
    for key in keys:
        if not key.endswith(".json"):
            continue
        try:
            data = storage.read_json(key)
            c = ComponentConstraints.model_validate(data)
        except Exception:
            logger.exception("skipping bad library extraction %s", key)
            continue
        result.setdefault(c.mpn, c)
    return result


def _publish(project_id: str, event: str, data: dict) -> None:
    broker.publish(project_id, event, data)


def _step(project_id: str, stage: str, status: str, detail: str = "") -> None:
    payload: dict = {"stage": stage, "status": status}
    if detail:
        payload["detail"] = detail
    _publish(project_id, "pcb_step_update", payload)


async def run_pcb_pipeline(
    storage: StorageBackend, user_id: str, project_id: str,
) -> None:
    meta = proj_svc.get_project(storage, user_id, project_id)
    if not meta:
        raise ValueError(f"Project {project_id} not found")
    if (meta.pcb_status or "draft") not in _PCB_ACTIVE:
        logger.warning(
            "pcb worker booted with pcb_status=%s for %s; exiting",
            meta.pcb_status, project_id,
        )
        return

    proj_svc.update_project(
        storage, user_id, project_id,
        pcb_status="running",
        pcb_cancel_requested=False,
        pcb_state=None,
    )

    try:
        async with PipelineWorkspace(storage, user_id, project_id) as ws:
            if _cancelled(storage, user_id, project_id):
                _finish_cancelled(storage, user_id, project_id)
                return

            graph = await _ensure_graph(ws, meta, project_id)
            if _cancelled(storage, user_id, project_id):
                _finish_cancelled(storage, user_id, project_id)
                return

            _step(project_id, "parse_pcb", "running")
            layout = _load_layout(ws)
            if layout is None:
                raise FileNotFoundError(
                    "Missing or unreadable uploads/pcb.kicad_pcb"
                )
            _step(
                project_id, "parse_pcb", "complete",
                f"{len(layout.footprints)} footprints, {len(layout.nets)} nets",
            )

            if _cancelled(storage, user_id, project_id):
                _finish_cancelled(storage, user_id, project_id)
                return

            _step(project_id, "classify", "running", "domains and groups")
            cmap = _load_constraints_map(ws.local_path("extracted"), storage)
            plan = build_placement_plan(graph, cmap)
            fg_path = ws.local_path("functional_groups.json")
            fg_path.write_text(plan.model_dump_json(indent=2) + "\n")
            ws._upload_file("functional_groups.json")
            _step(
                project_id, "classify", "complete",
                f"{len(plan.domains)} domains, {len(plan.groups)} groups",
            )

            if _cancelled(storage, user_id, project_id):
                _finish_cancelled(storage, user_id, project_id)
                return

            _step(project_id, "inventory", "running")
            z0_by_net: dict[str, float] = {}
            try:
                from backend.periscopex.impedance_traces import analyze_where_needed

                zrep = analyze_where_needed(layout, graph)
                for row in zrep.get("nets") or []:
                    if not isinstance(row, dict) or row.get("error"):
                        continue
                    name = str(row.get("net_name") or row.get("name") or "")
                    z = row.get("z0_ohm") or row.get("mean_z0") or row.get("z0")
                    if name and isinstance(z, (int, float)):
                        z0_by_net[name] = float(z)
            except Exception:
                logger.exception("Z0 inventory skipped")
            inventory = build_pcb_inventory(
                layout, graph,
                domain_ids=[d.domain_id for d in plan.domains],
                group_count=len(plan.groups),
                z0_by_net=z0_by_net,
            )
            inv_path = ws.local_path("pcb_inventory.json")
            inv_path.write_text(inventory.model_dump_json(indent=2) + "\n")
            ws._upload_file("pcb_inventory.json")
            _step(
                project_id, "inventory", "complete",
                f"{len(inventory.nets)} routed nets",
            )

            if _cancelled(storage, user_id, project_id):
                _finish_cancelled(storage, user_id, project_id)
                return

            _step(project_id, "checks", "running")
            findings = run_pcb_checks(graph, cmap, layout)
            _step(
                project_id, "checks", "complete",
                f"{len(findings)} deterministic findings",
            )

            if _cancelled(storage, user_id, project_id):
                _finish_cancelled(storage, user_id, project_id)
                return

            _step(project_id, "ai_review", "running", "layout vs shared library extraction")
            coverage: dict[str, list[str]] = {}
            skipped: list[dict] = []
            try:
                from backend.services.api_logs import ApiLogger
                from backend.services.pcb_validation import review_pcb_ics

                pdf_dir = ws.local_path("uploads/datasheets")
                pdf_dir.mkdir(parents=True, exist_ok=True)
                logger_api = ApiLogger()

                async def _prog(ref, turn, tool, detail):
                    _step(
                        project_id, "ai_review", "running",
                        f"{ref} {tool} {detail}".strip(),
                    )

                ai_findings, coverage, skipped = await review_pcb_ics(
                    graph, cmap, layout, plan, inventory,
                    pdf_dir, storage=storage, api_logger=logger_api,
                    on_progress=_prog,
                )
                findings.extend(ai_findings)
                logger_api.flush(storage, user_id, project_id)
                detail = (
                    f"{len(ai_findings)} AI findings, {len(skipped)} skipped"
                )
            except Exception:
                logger.exception("PCB AI review failed — keeping deterministic findings")
                detail = "AI skipped (error)"
            _step(project_id, "ai_review", "complete", detail)

            if _cancelled(storage, user_id, project_id):
                _finish_cancelled(storage, user_id, project_id)
                return

            _step(project_id, "write_report", "running")
            assign_pcb_finding_ids(findings)
            annotate_findings_cad(findings, cad_index_from_graph(graph))
            dec_path = ws.local_path("decisions.json")
            if dec_path.is_file():
                try:
                    apply_decisions(findings, json.loads(dec_path.read_text()))
                except Exception:
                    logger.exception("decisions.json apply failed")
            summary: dict[str, int] = {"ERROR": 0, "WARNING": 0, "INFO": 0}
            for f in findings:
                summary[f.status] = summary.get(f.status, 0) + 1
            report = ValidationReport(
                project=project_id,
                timestamp=datetime.now(timezone.utc).isoformat(),
                findings=findings,
                summary=summary,
                coverage=coverage,
                not_reviewed=skipped,
            )
            report_path = ws.local_path("pcb_report.json")
            report_path.write_text(report.model_dump_json(indent=2) + "\n")
            ws._upload_file("pcb_report.json")

            bridge = build_cad_bridge(report, project_id)
            write_cad_bridge(ws.local_path("periscope-findings.json"), bridge)
            ws._upload_file("periscope-findings.json")
            _step(project_id, "write_report", "complete", f"{len(findings)} findings")

        proj_svc.update_project(
            storage, user_id, project_id,
            pcb_status="complete",
            pcb_state={
                "findings": len(findings),
                "domains": len(plan.domains),
                "groups": len(plan.groups),
                "skipped": len(skipped),
            },
            pcb_cancel_requested=False,
        )
        _publish(project_id, "pcb_complete", {
            "findings": len(findings),
            "domains": len(plan.domains),
            "groups": len(plan.groups),
        })
    except Exception as e:
        logger.exception("pcb pipeline failed for %s", project_id)
        proj_svc.update_project(
            storage, user_id, project_id,
            pcb_status="error",
            pcb_state={"error": str(e)},
        )
        _publish(project_id, "pcb_error", {"error": str(e)})


async def _ensure_graph(ws: PipelineWorkspace, meta, project_id: str) -> DesignGraph:
    graph_path = ws.local_path("design_graph.json")
    if graph_path.is_file():
        _step(project_id, "ensure_graph", "running", "reusing design_graph.json")
        graph = DesignGraph.model_validate_json(graph_path.read_text(encoding="utf-8"))
        _step(
            project_id, "ensure_graph", "complete",
            f"{len(graph.components)} components (cached)",
        )
        return graph

    _step(project_id, "ensure_graph", "running", "building design graph")
    bom_path = ws.local_path("uploads/bom.csv")
    netlist_path = ws.netlist_local_path()
    if not bom_path.is_file() or not Path(netlist_path).is_file():
        raise FileNotFoundError("Missing BOM or netlist for PCB graph_build")

    col_map = meta.bom_columns or {}
    graph = build_graph(
        str(netlist_path),
        str(bom_path),
        str(ws.local_path("extracted")),
        str(ws.local_path("patterns")),
        str(ws.local_path("models")),
        reference_col=col_map.get("reference", "Reference"),
        mpn_col=col_map.get("mpn", "Manufacturer Part Number"),
        include_subdesigns=(
            set(meta.netlist_subdesigns)
            if meta.netlist_subdesigns is not None
            else None
        ),
        pcb_path=ws.local_path("uploads/pcb.kicad_pcb"),
    )
    graph_path.write_text(graph.model_dump_json(indent=2) + "\n")
    ws._upload_file("design_graph.json")
    _step(
        project_id, "ensure_graph", "complete",
        f"{len(graph.components)} components, {len(graph.nets)} nets",
    )
    return graph


def _load_layout(ws: PipelineWorkspace) -> LayoutGraph | None:
    cached = ws.local_path("layout_graph.json")
    if cached.is_file():
        try:
            return LayoutGraph.model_validate_json(
                cached.read_text(encoding="utf-8"),
            )
        except Exception:
            logger.exception("bad layout_graph.json — trying pcb parse")

    pcb = ws.local_path("uploads/pcb.kicad_pcb")
    if not pcb.is_file():
        return None
    try:
        from backend.periscopex.parsers_kicad_pcb import parse_kicad_pcb

        layout = parse_kicad_pcb(pcb)
        cached.write_text(layout.model_dump_json(indent=2) + "\n")
        ws._upload_file("layout_graph.json")
        return layout
    except Exception:
        logger.exception("kicad_pcb parse failed during PCB review")
        return None


def _cancelled(storage: StorageBackend, user_id: str, project_id: str) -> bool:
    meta = proj_svc.get_project(storage, user_id, project_id)
    return bool(meta and meta.pcb_cancel_requested)


def _finish_cancelled(storage: StorageBackend, user_id: str, project_id: str) -> None:
    proj_svc.update_project(
        storage, user_id, project_id,
        pcb_status="cancelled",
        pcb_cancel_requested=False,
    )
    _publish(project_id, "pcb_cancelled", {})


def pcb_busy(meta: proj_svc.ProjectMeta) -> bool:
    return (meta.pcb_status or "draft") in _PCB_ACTIVE


def analysis_busy(meta: proj_svc.ProjectMeta) -> bool:
    return meta.status in _ANALYSIS_BUSY


def placement_busy(meta: proj_svc.ProjectMeta) -> bool:
    return (meta.placement_status or "draft") in _PLACEMENT_ACTIVE