kicad-mcp/src/mckicad/utils/lib_resolver.py

"""Unified symbol library resolution for KiCad.

Given a lib_id like ``"Device:R"``, finds the ``.kicad_sym`` file and
extracts the full symbol sexp tree.  Reads both global and project
``sym-lib-table`` files, resolving KiCad path variables.

This module is the eventual replacement for the ad-hoc library search
in sexp_parser.py, providing proper sym-lib-table parsing with variable
substitution and mtime-based cache invalidation.
"""

from __future__ import annotations

import logging
import os
from typing import Any

from mckicad.config import get_kicad_user_dir
from mckicad.utils.sexp_tree import (
    SexpNode,
    find,
    find_named,
    find_one,
    find_recursive,
    get_values,
    parse,
    parse_file,
    text_at,
)

logger = logging.getLogger(__name__)

# Cache: path -> (mtime, parsed table dict)
_lib_table_cache: dict[str, tuple[float, dict[str, str]]] = {}


def resolve_symbol(
    lib_id: str, *, project_dir: str | None = None,
) -> SexpNode | None:
    """Resolve a lib_id to its parsed symbol node.

    Args:
        lib_id: Full library identifier (e.g. ``"Device:R"``).
        project_dir: Project directory for ``${KIPRJMOD}`` substitution.

    Returns:
        The parsed SexpNode for the symbol, or None if not found.
    """
    if ":" not in lib_id:
        return None

    library_name, symbol_name = lib_id.split(":", 1)
    if not library_name or not symbol_name:
        return None

    lib_path = resolve_library_path(library_name, project_dir=project_dir)
    if lib_path is None:
        return None

    try:
        tree = parse_file(lib_path)
    except (OSError, UnicodeDecodeError):
        logger.debug("Failed to parse library file: %s", lib_path)
        return None

    return find_named(tree, "symbol", symbol_name)


def resolve_symbol_with_tree(
    lib_id: str, *, project_dir: str | None = None,
) -> tuple[SexpNode | None, SexpNode | None]:
    """Resolve a lib_id, returning both the symbol node and its library tree.

    The library tree is needed for ``extract_symbol_pins`` and
    ``extract_pin_unit_map`` when the symbol uses ``(extends ...)``
    to inherit pin geometry from a parent symbol.

    Returns:
        ``(symbol_node, lib_tree)`` tuple.  Both are None if resolution fails.
    """
    if ":" not in lib_id:
        return None, None

    library_name, symbol_name = lib_id.split(":", 1)
    if not library_name or not symbol_name:
        return None, None

    lib_path = resolve_library_path(library_name, project_dir=project_dir)
    if lib_path is None:
        return None, None

    try:
        tree = parse_file(lib_path)
    except (OSError, UnicodeDecodeError):
        logger.debug("Failed to parse library file: %s", lib_path)
        return None, None

    sym = find_named(tree, "symbol", symbol_name)
    return sym, tree


def resolve_library_path(
    library_name: str, *, project_dir: str | None = None,
) -> str | None:
    """Resolve a library name to its ``.kicad_sym`` file path.

    Search order:
    1. Project sym-lib-table (if project_dir given)
    2. Global sym-lib-table (KiCad 10, 9, 8, 7)
    3. Direct file search (project dir, libs/, ../libs/)
    4. System symbol directory (built-in libraries like Device)
    """
    filename = f"{library_name}.kicad_sym"

    # 1. Project sym-lib-table
    if project_dir:
        table_path = os.path.join(project_dir, "sym-lib-table")
        if os.path.isfile(table_path):
            table = load_sym_lib_table(table_path)
            if library_name in table:
                uri = _substitute_variables(table[library_name], project_dir)
                if os.path.isfile(uri):
                    return uri

    # 2. Global sym-lib-tables (newest KiCad first)
    for version_dir in ("10.0", "9.0", "8.0", "7.0"):
        config_dir = os.path.expanduser(f"~/.config/kicad/{version_dir}")
        table_path = os.path.join(config_dir, "sym-lib-table")
        if os.path.isfile(table_path):
            table = load_sym_lib_table(table_path)
            if library_name in table:
                uri = _substitute_variables(table[library_name], project_dir)
                if os.path.isfile(uri):
                    return uri

    # 3. Direct file search (fallback for projects without sym-lib-table)
    if project_dir:
        candidates = [
            os.path.join(project_dir, filename),
            os.path.join(project_dir, "libs", filename),
            os.path.join(project_dir, os.pardir, "libs", filename),
        ]
        for candidate in candidates:
            resolved = os.path.normpath(candidate)
            if os.path.isfile(resolved):
                return resolved

    # 4. System symbol directory (built-in libraries, e.g. Device, Connector).
    # This is the authoritative source when no sym-lib-table is configured.
    symbol_dir = get_kicad_symbol_dir()
    if symbol_dir:
        resolved = os.path.join(symbol_dir, filename)
        if os.path.isfile(resolved):
            return resolved

    return None


def extract_symbol_pins(
    symbol_node: SexpNode,
    *,
    lib_tree: SexpNode | None = None,
) -> list[dict[str, Any]]:
    """Extract pin dicts from a parsed symbol node.

    Output format matches ``sexp_parser.parse_lib_symbol_pins()``:
    ``{number, name, type, shape, x, y, rotation, length}``

    Searches through sub-symbols (children with tag "symbol") to find
    all pin definitions.  If the symbol uses ``(extends "Parent")``,
    follows the chain to find pin definitions in the parent symbol
    (requires ``lib_tree`` — the parsed library file root).
    """
    # Check for extends chain
    effective = _resolve_extends(symbol_node, lib_tree)

    pins: list[dict[str, Any]] = []
    all_pin_nodes = find_recursive(effective, "pin")

    for pin_node in all_pin_nodes:
        pin_dict = _parse_pin_node(pin_node)
        if pin_dict is not None:
            pins.append(pin_dict)

    return pins


def extract_pin_unit_map(
    symbol_node: SexpNode,
    *,
    lib_tree: SexpNode | None = None,
) -> dict[str, int]:
    """Build pin_number -> unit_number map from sub-symbols.

    KiCad sub-symbols are named ``<base>_<unit>_<style>`` (e.g.
    ``TL072_2_1`` for unit 2).  If the symbol uses ``(extends "Parent")``,
    the parent's sub-symbols are used for pin resolution.
    """
    effective = _resolve_extends(symbol_node, lib_tree)
    base_name = effective.values[0] if effective.values else ""
    pin_unit_map: dict[str, int] = {}

    for sub_sym in find(effective, "symbol"):
        if not sub_sym.values:
            continue
        sub_name = sub_sym.values[0]

        # Parse unit number from <base>_<unit>_<style> pattern
        if not sub_name.startswith(base_name + "_"):
            continue
        suffix = sub_name[len(base_name) + 1 :]
        parts = suffix.split("_", 1)
        if not parts or not parts[0].isdigit():
            continue
        unit_num = int(parts[0])

        # Extract pins in this sub-symbol
        for pin_node in find_recursive(sub_sym, "pin"):
            pin_dict = _parse_pin_node(pin_node)
            if pin_dict is not None:
                pin_unit_map[pin_dict["number"]] = unit_num

    return pin_unit_map


def extract_symbol_sexp_text(
    lib_id: str, *, project_dir: str | None = None,
) -> str | None:
    """Return the raw sexp text for a symbol from its library file.

    Used to embed correct lib_symbols in schematics (avoiding
    kicad-sch-api's stale cached versions).
    """
    if ":" not in lib_id:
        return None

    library_name, symbol_name = lib_id.split(":", 1)
    if not library_name or not symbol_name:
        return None

    lib_path = resolve_library_path(library_name, project_dir=project_dir)
    if lib_path is None:
        return None

    try:
        with open(lib_path, encoding="utf-8") as f:
            source = f.read()
    except (OSError, UnicodeDecodeError):
        return None

    tree = parse(source)
    sym = find_named(tree, "symbol", symbol_name)
    if sym is None:
        return None

    return text_at(source, sym.span)


def get_kicad_symbol_dir() -> str | None:
    """Detect the KiCad system symbol directory.

    Checks, in order, the generic ``KICAD_SYMBOL_DIR`` env var, the
    version-specific ``KICAD{10,9,8,7}_SYMBOL_DIR`` vars, then the standard
    install paths for macOS, Linux, and Windows.
    """
    for env_var in (
        "KICAD_SYMBOL_DIR",
        "KICAD10_SYMBOL_DIR",
        "KICAD9_SYMBOL_DIR",
        "KICAD8_SYMBOL_DIR",
        "KICAD7_SYMBOL_DIR",
    ):
        val = os.environ.get(env_var)
        if val and os.path.isdir(val):
            return val

    # Standard per-platform install paths
    candidates = [
        # macOS
        "/Applications/KiCad/KiCad.app/Contents/SharedSupport/symbols",
        # Linux
        "/usr/share/kicad/symbols",
        "/usr/local/share/kicad/symbols",
        # Windows
        r"C:\Program Files\KiCad\10.0\share\kicad\symbols",
        r"C:\Program Files\KiCad\9.0\share\kicad\symbols",
    ]
    for path in candidates:
        if os.path.isdir(path):
            return path

    return None


def load_sym_lib_table(table_path: str) -> dict[str, str]:
    """Parse a sym-lib-table into ``{name: uri}`` dict.

    Uses mtime-based caching to avoid re-parsing unchanged files.
    URIs are returned with variables un-substituted (caller applies
    substitution based on context).
    """
    try:
        mtime = os.path.getmtime(table_path)
    except OSError:
        return {}

    cached = _lib_table_cache.get(table_path)
    if cached is not None and cached[0] == mtime:
        return cached[1]

    try:
        with open(table_path, encoding="utf-8") as f:
            content = f.read()
    except (OSError, UnicodeDecodeError):
        return {}

    tree = parse(content)
    result: dict[str, str] = {}

    for lib_node in find(tree, "lib"):
        name_vals = get_values(lib_node, "name")
        uri_vals = get_values(lib_node, "uri")
        if name_vals and uri_vals:
            result[name_vals[0]] = uri_vals[0]

    _lib_table_cache[table_path] = (mtime, result)
    return result


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _resolve_extends(
    symbol_node: SexpNode, lib_tree: SexpNode | None,
) -> SexpNode:
    """Follow ``(extends "ParentName")`` to find the symbol with actual pins.

    KiCad allows symbols to inherit pin geometry from a parent symbol
    in the same library file (e.g. TL072 extends LM2904).  This
    function walks the chain up to 5 levels to find the root definition.
    """
    current = symbol_node
    for _ in range(5):  # guard against cycles
        extends_node = find_one(current, "extends")
        if extends_node is None or not extends_node.values:
            break
        if lib_tree is None:
            break
        parent_name = extends_node.values[0]
        parent = find_named(lib_tree, "symbol", parent_name)
        if parent is None:
            break
        current = parent
    return current


def _parse_pin_node(pin_node: SexpNode) -> dict[str, Any] | None:
    """Convert a parsed ``(pin type shape ...)`` node into a pin dict."""
    # pin_node.values = [type, shape, ...]
    if len(pin_node.values) < 2:
        return None

    pin_type = pin_node.values[0]
    pin_shape = pin_node.values[1]

    # (at X Y [rotation])
    at_vals = get_values(pin_node, "at")
    if not at_vals or len(at_vals) < 2:
        return None
    x = float(at_vals[0])
    y = float(at_vals[1])
    rotation = float(at_vals[2]) if len(at_vals) > 2 else 0.0

    # (length L)
    length_vals = get_values(pin_node, "length")
    length = float(length_vals[0]) if length_vals else 0.0

    # (name "NAME" ...)
    name_node = find_one(pin_node, "name")
    name = name_node.values[0] if name_node and name_node.values else ""

    # (number "NUM" ...)
    number_node = find_one(pin_node, "number")
    number = number_node.values[0] if number_node and number_node.values else ""

    return {
        "number": number,
        "name": name,
        "type": pin_type,
        "shape": pin_shape,
        "x": x,
        "y": y,
        "rotation": rotation,
        "length": length,
    }


def search_symbols(
    query: str, *, project_dir: str | None = None,
) -> list[dict[str, Any]]:
    """Search all available symbol libraries for symbols matching a query.

    Iterates libraries from sym-lib-table (global + project-local) and
    fuzzy-matches the query against symbol name, keywords, and description.

    Returns list of dicts: {lib_id, name, description, keywords, pin_count}.
    """
    query_lower = query.lower()
    results: list[dict[str, Any]] = []

    seen_lib_ids: set[str] = set()

    for lib_name, lib_path in _iter_all_libraries(project_dir=project_dir):
        try:
            tree = parse_file(lib_path)
        except (OSError, UnicodeDecodeError):
            continue

        for sym_node in find(tree, "symbol"):
            if not sym_node.values:
                continue
            sym_name = sym_node.values[0]
            # Skip sub-symbols (contain _ followed by digit pattern)
            if "_" in sym_name:
                parts = sym_name.rsplit("_", 2)
                if len(parts) >= 3 and parts[-1].isdigit() and parts[-2].isdigit():
                    continue

            lib_id = f"{lib_name}:{sym_name}"
            if lib_id in seen_lib_ids:
                continue

            # Extract searchable fields
            description = ""
            keywords = ""
            for prop in find(sym_node, "property"):
                vals = prop.values
                if len(vals) >= 2:
                    pname = vals[0] if vals[0] != "private" else (vals[1] if len(vals) >= 3 else "")
                    pval = vals[1] if vals[0] != "private" else (vals[2] if len(vals) >= 3 else "")
                    if pname == "ki_keywords":
                        keywords = pval
                    elif pname == "ki_description":
                        description = pval

            # Fuzzy match
            searchable = f"{sym_name} {description} {keywords}".lower()
            if query_lower not in searchable:
                continue

            seen_lib_ids.add(lib_id)
            pin_count = len(extract_symbol_pins(sym_node, lib_tree=tree))
            results.append({
                "lib_id": lib_id,
                "name": sym_name,
                "description": description,
                "keywords": keywords,
                "pin_count": pin_count,
            })

    return results


def get_symbol_info(
    lib_id: str, *, project_dir: str | None = None,
) -> dict[str, Any] | None:
    """Get detailed info about a specific symbol by lib_id.

    Returns dict with name, description, keywords, pin_count,
    footprint_filters, and units count. Returns None if not found.
    """
    sym_node, lib_tree = resolve_symbol_with_tree(lib_id, project_dir=project_dir)
    if sym_node is None:
        return None

    name = sym_node.values[0] if sym_node.values else ""
    description = ""
    keywords = ""
    fp_filters = ""

    for prop in find(sym_node, "property"):
        vals = prop.values
        if len(vals) >= 2:
            pname = vals[0] if vals[0] != "private" else (vals[1] if len(vals) >= 3 else "")
            pval = vals[1] if vals[0] != "private" else (vals[2] if len(vals) >= 3 else "")
            if pname == "ki_keywords":
                keywords = pval
            elif pname == "ki_description":
                description = pval
            elif pname == "ki_fp_filters":
                fp_filters = pval

    pins = extract_symbol_pins(sym_node, lib_tree=lib_tree)

    # Count units from sub-symbol enumeration
    units: set[int] = set()
    base_name = name
    for sub in find(sym_node, "symbol"):
        if sub.values and sub.values[0].startswith(base_name + "_"):
            suffix = sub.values[0][len(base_name) + 1:]
            parts = suffix.split("_", 1)
            if parts and parts[0].isdigit():
                units.add(int(parts[0]))

    return {
        "lib_id": lib_id,
        "name": name,
        "description": description,
        "keywords": keywords,
        "pin_count": len(pins),
        "pins": pins,
        "footprint_filters": fp_filters,
        "units": max(len(units), 1),
    }


def _iter_all_libraries(
    *, project_dir: str | None = None,
) -> list[tuple[str, str]]:
    """Yield (library_name, file_path) for all available symbol libraries."""
    result: list[tuple[str, str]] = []
    seen: set[str] = set()

    # Project sym-lib-table
    if project_dir:
        table_path = os.path.join(project_dir, "sym-lib-table")
        if os.path.isfile(table_path):
            table = load_sym_lib_table(table_path)
            for name, uri in table.items():
                if name not in seen:
                    resolved = _substitute_variables(uri, project_dir)
                    if os.path.isfile(resolved):
                        result.append((name, resolved))
                        seen.add(name)

    # Global sym-lib-tables
    for version_dir in ("9.0", "8.0"):
        config_dir = os.path.expanduser(f"~/.config/kicad/{version_dir}")
        table_path = os.path.join(config_dir, "sym-lib-table")
        if os.path.isfile(table_path):
            table = load_sym_lib_table(table_path)
            for name, uri in table.items():
                if name not in seen:
                    resolved = _substitute_variables(uri, project_dir)
                    if os.path.isfile(resolved):
                        result.append((name, resolved))
                        seen.add(name)

    return result


def _substitute_variables(uri: str, project_dir: str | None = None) -> str:
    """Replace KiCad path variables in a URI."""
    if "${KIPRJMOD}" in uri and project_dir:
        uri = uri.replace("${KIPRJMOD}", project_dir)

    symbol_dir = get_kicad_symbol_dir()
    if symbol_dir:
        uri = uri.replace("${KICAD9_SYMBOL_DIR}", symbol_dir)
        uri = uri.replace("${KICAD8_SYMBOL_DIR}", symbol_dir)

    kicad_user_dir = get_kicad_user_dir()
    if kicad_user_dir:
        uri = uri.replace("${KICAD_USER_DIR}", kicad_user_dir)

    return os.path.normpath(uri)