kicad-mcp-server/python/commands/dynamic_symbol_loader.py

"""
Dynamic Symbol Loader for KiCad Schematics

Loads symbols from .kicad_sym library files and injects them into schematics
on-the-fly using TEXT MANIPULATION (not sexpdata) to preserve file formatting.

This enables access to all ~10,000+ KiCad symbols dynamically.
"""

import os
import re
import uuid
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger("kicad_interface")


class DynamicSymbolLoader:
    """
    Dynamically loads symbols from KiCad library files and injects them into schematics.

    Uses raw text manipulation instead of sexpdata to avoid corrupting the KiCad file format.

    Key rules for KiCad 9 .kicad_sch format:
    - Top-level symbols in lib_symbols must have library prefix: (symbol "Device:R" ...)
    - Sub-symbols must NOT have library prefix: (symbol "R_0_1" ...), (symbol "R_1_1" ...)
    - Parent symbols must appear BEFORE child symbols that use (extends ...)
    """

    def __init__(self, project_path: Optional[Path] = None):
        self.symbol_cache = {}  # Cache: "lib:symbol" -> raw text block
        self.project_path = project_path  # Project directory for project-specific libraries

    def find_kicad_symbol_libraries(self) -> List[Path]:
        """Find all KiCad symbol library directories"""
        possible_paths = [
            Path("/usr/share/kicad/symbols"),
            Path("/usr/local/share/kicad/symbols"),
            Path("C:/Program Files/KiCad/9.0/share/kicad/symbols"),
            Path("C:/Program Files/KiCad/8.0/share/kicad/symbols"),
            Path("/Applications/KiCad/KiCad.app/Contents/SharedSupport/symbols"),
            Path.home() / ".local" / "share" / "kicad" / "9.0" / "symbols",
            Path.home() / "Documents" / "KiCad" / "9.0" / "3rdparty" / "symbols",
        ]
        for env_var in ["KICAD9_SYMBOL_DIR", "KICAD8_SYMBOL_DIR", "KICAD_SYMBOL_DIR"]:
            if env_var in os.environ:
                possible_paths.insert(0, Path(os.environ[env_var]))

        return [p for p in possible_paths if p.exists() and p.is_dir()]

    def find_library_file(self, library_name: str) -> Optional[Path]:
        """Find the .kicad_sym file for a given library name.

        Search order:
        1. Project-specific sym-lib-table (if project_path is set)
        2. Global KiCad symbol library directories
        """
        # 1. Check project-specific sym-lib-table
        if self.project_path:
            project_table = Path(self.project_path) / "sym-lib-table"
            if project_table.exists():
                resolved = self._resolve_library_from_table(project_table, library_name)
                if resolved:
                    logger.info(f"Found '{library_name}' in project sym-lib-table: {resolved}")
                    return resolved

        # 2. Fall back to global KiCad symbol directories
        for lib_dir in self.find_kicad_symbol_libraries():
            lib_file = lib_dir / f"{library_name}.kicad_sym"
            if lib_file.exists():
                return lib_file

        logger.warning(f"Library file not found: {library_name}.kicad_sym")
        return None

    def _resolve_library_from_table(self, table_path: Path, library_name: str) -> Optional[Path]:
        """Parse a sym-lib-table file and return the resolved path for the given library nickname."""
        try:
            with open(table_path, "r", encoding="utf-8") as f:
                content = f.read()

            lib_pattern = r'\(lib\s+\(name\s+"?([^"\)\s]+)"?\)\s*\(type\s+[^)]+\)\s*\(uri\s+"?([^"\)\s]+)"?'
            for match in re.finditer(lib_pattern, content, re.IGNORECASE):
                nickname = match.group(1)
                if nickname != library_name:
                    continue
                uri = match.group(2)
                resolved = self._resolve_sym_uri(uri)
                if resolved and Path(resolved).exists():
                    return Path(resolved)
        except Exception as e:
            logger.warning(f"Could not parse sym-lib-table {table_path}: {e}")
        return None

    def _resolve_sym_uri(self, uri: str) -> Optional[str]:
        """Resolve environment variables in a sym-lib-table URI."""
        env_map = {
            "KICAD9_SYMBOL_DIR": [
                "C:/Program Files/KiCad/9.0/share/kicad/symbols",
                "/usr/share/kicad/symbols",
                "/Applications/KiCad/KiCad.app/Contents/SharedSupport/symbols",
            ],
            "KICAD8_SYMBOL_DIR": [
                "C:/Program Files/KiCad/8.0/share/kicad/symbols",
            ],
            "KIPRJMOD": [str(self.project_path)] if self.project_path else [],
        }
        result = uri
        for var, candidates in env_map.items():
            if f"${{{var}}}" in result:
                for candidate in candidates:
                    candidate_path = result.replace(f"${{{var}}}", candidate)
                    if Path(candidate_path).exists():
                        return candidate_path
                # Fallback: try OS env
                if var in os.environ:
                    return result.replace(f"${{{var}}}", os.environ[var])
        return result

    def _extract_symbol_block(self, text: str, symbol_name: str) -> Optional[str]:
        """
        Extract a complete symbol block from a library or schematic file by matching
        parentheses depth. Returns the raw text of the symbol definition.
        """
        lines = text.split("\n")
        start = None
        for i, line in enumerate(lines):
            stripped = line.strip()
            # Match exact symbol name (not sub-symbols like Name_0_1)
            if stripped.startswith(f'(symbol "{symbol_name}"') and not re.match(
                r'.*_\d+_\d+"', stripped
            ):
                start = i
                break

        if start is None:
            return None

        depth = 0
        end = None
        for i in range(start, len(lines)):
            for ch in lines[i]:
                if ch == "(":
                    depth += 1
                elif ch == ")":
                    depth -= 1
                    if depth == 0:
                        end = i
                        break
            if end is not None:
                break

        if end is None:
            return None

        return "\n".join(lines[start : end + 1])

    def extract_symbol_from_library(
        self, library_name: str, symbol_name: str
    ) -> Optional[str]:
        """
        Extract a symbol definition from a KiCad .kicad_sym library file.
        Returns the raw text block, ready to be injected into a schematic.

        The returned block has:
        - Top-level name prefixed with library: (symbol "Library:Name" ...)
        - Sub-symbol names WITHOUT prefix: (symbol "Name_0_1" ...)
        """
        cache_key = f"{library_name}:{symbol_name}"
        if cache_key in self.symbol_cache:
            return self.symbol_cache[cache_key]

        lib_path = self.find_library_file(library_name)
        if not lib_path:
            return None

        with open(lib_path, "r", encoding="utf-8") as f:
            lib_content = f.read()

        block = self._extract_symbol_block(lib_content, symbol_name)
        if block is None:
            logger.warning(
                f"Symbol '{symbol_name}' not found in {library_name}.kicad_sym"
            )
            return None

        # Check if this symbol uses (extends "ParentName")
        extends_match = re.search(r'\(extends "([^"]+)"\)', block)
        parent_block = None
        if extends_match:
            parent_name = extends_match.group(1)
            logger.info(
                f"Symbol {symbol_name} extends {parent_name}, extracting parent too"
            )
            parent_block = self._extract_symbol_block(lib_content, parent_name)
            if parent_block:
                # Prefix parent top-level name with library
                parent_block = parent_block.replace(
                    f'(symbol "{parent_name}"',
                    f'(symbol "{library_name}:{parent_name}"',
                    1,  # Only first occurrence (top-level)
                )

        # Prefix top-level symbol name with library
        full_name = f"{library_name}:{symbol_name}"
        block = block.replace(
            f'(symbol "{symbol_name}"',
            f'(symbol "{full_name}"',
            1,  # Only first occurrence (top-level)
        )
        # Sub-symbols like "Name_0_1" keep their short names (already correct from library)

        # Combine parent + child if extends is used
        if parent_block:
            result = parent_block + "\n" + block
        else:
            result = block

        self.symbol_cache[cache_key] = result
        logger.info(f"Extracted symbol {full_name} ({len(result)} chars)")
        return result

    def inject_symbol_into_schematic(
        self, schematic_path: Path, library_name: str, symbol_name: str
    ) -> bool:
        """
        Inject a symbol definition into a schematic's lib_symbols section.
        Uses text manipulation to preserve file formatting.
        """
        full_name = f"{library_name}:{symbol_name}"

        with open(schematic_path, "r", encoding="utf-8") as f:
            content = f.read()

        # Check if symbol already exists
        if f'(symbol "{full_name}"' in content:
            logger.info(f"Symbol {full_name} already exists in schematic")
            return True

        # Extract symbol from library
        symbol_block = self.extract_symbol_from_library(library_name, symbol_name)
        if not symbol_block:
            raise ValueError(
                f"Symbol '{symbol_name}' not found in library '{library_name}'"
            )

        # Indent the block to match lib_symbols indentation (4 spaces for top-level)
        indented_lines = []
        for line in symbol_block.split("\n"):
            # Add 4-space indent for the content inside lib_symbols
            indented_lines.append("    " + line if line.strip() else line)
        indented_block = "\n".join(indented_lines)

        # Find the end of lib_symbols section to insert before closing )
        lines = content.split("\n")
        lib_sym_start = None
        lib_sym_end = None
        depth = 0

        for i, line in enumerate(lines):
            if "(lib_symbols" in line and lib_sym_start is None:
                lib_sym_start = i
                depth = 0
                for ch in line:
                    if ch == "(":
                        depth += 1
                    elif ch == ")":
                        depth -= 1
                continue
            if lib_sym_start is not None and lib_sym_end is None:
                for ch in line:
                    if ch == "(":
                        depth += 1
                    elif ch == ")":
                        depth -= 1
                        if depth == 0:
                            lib_sym_end = i
                            break
                if lib_sym_end is not None:
                    break

        if lib_sym_end is None:
            raise ValueError("No lib_symbols section found in schematic")

        # Insert the symbol block just before the closing ) of lib_symbols
        lines.insert(lib_sym_end, indented_block)

        with open(schematic_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))

        # Handle both Path objects and strings
        sch_name = (
            schematic_path.name
            if hasattr(schematic_path, "name")
            else str(schematic_path)
        )
        logger.info(f"Injected symbol {full_name} into {sch_name}")
        return True

    def create_component_instance(
        self,
        schematic_path: Path,
        library_name: str,
        symbol_name: str,
        reference: str,
        value: str = "",
        footprint: str = "",
        x: float = 0,
        y: float = 0,
    ) -> bool:
        """
        Add a component instance to the schematic.
        This creates the (symbol ...) block with lib_id reference.
        """
        full_lib_id = f"{library_name}:{symbol_name}"
        new_uuid = str(uuid.uuid4())

        instance_block = f"""  (symbol (lib_id "{full_lib_id}") (at {x} {y} 0) (unit 1)
    (in_bom yes) (on_board yes) (dnp no)
    (uuid "{new_uuid}")
    (property "Reference" "{reference}" (at {x} {y - 2.54} 0)
      (effects (font (size 1.27 1.27)))
    )
    (property "Value" "{value or symbol_name}" (at {x} {y + 2.54} 0)
      (effects (font (size 1.27 1.27)))
    )
    (property "Footprint" "{footprint}" (at {x} {y} 0)
      (effects (font (size 1.27 1.27)) (hide yes))
    )
    (property "Datasheet" "~" (at {x} {y} 0)
      (effects (font (size 1.27 1.27)) (hide yes))
    )
  )"""

        with open(schematic_path, "r", encoding="utf-8") as f:
            content = f.read()

        # Insert before (sheet_instances or at end before final )
        lines = content.split("\n")
        insert_pos = None

        for i, line in enumerate(lines):
            if "(sheet_instances" in line:
                insert_pos = i
                break

        if insert_pos is None:
            # Insert before the last closing parenthesis
            for i in range(len(lines) - 1, -1, -1):
                if lines[i].strip() == ")":
                    insert_pos = i
                    break

        if insert_pos is None:
            raise ValueError("Could not find insertion point in schematic")

        lines.insert(insert_pos, instance_block)

        with open(schematic_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))

        logger.info(
            f"Added component instance {reference} ({full_lib_id}) at ({x}, {y})"
        )
        return True

    def load_symbol_dynamically(
        self, schematic_path: Path, library_name: str, symbol_name: str
    ) -> str:
        """
        Complete workflow: inject symbol definition and create a template instance.
        Returns a template reference name.
        """
        logger.info(f"Loading symbol dynamically: {library_name}:{symbol_name}")

        # Step 1: Inject symbol definition into lib_symbols
        self.inject_symbol_into_schematic(schematic_path, library_name, symbol_name)

        # Step 2: Create an offscreen template instance
        lib_clean = library_name.replace("-", "_").replace(".", "_")
        sym_clean = symbol_name.replace("-", "_").replace(".", "_")
        template_ref = f"_TEMPLATE_{lib_clean}_{sym_clean}"

        self.create_component_instance(
            schematic_path,
            library_name,
            symbol_name,
            reference=template_ref,
            value=symbol_name,
            x=-200,
            y=-200,
        )

        logger.info(f"Symbol loaded. Template reference: {template_ref}")
        return template_ref

    def add_component(
        self,
        schematic_path: Path,
        library_name: str,
        symbol_name: str,
        reference: str,
        value: str = "",
        footprint: str = "",
        x: float = 0,
        y: float = 0,
        project_path: Optional[Path] = None,
    ) -> bool:
        """
        High-level: ensure symbol definition exists in schematic, then add an instance.
        This is the main entry point for adding components.

        Args:
            project_path: Optional project directory. When set, project-specific
                          sym-lib-table is also searched for the library file.
        """
        if project_path:
            self.project_path = project_path
        # Ensure symbol definition is in lib_symbols
        self.inject_symbol_into_schematic(schematic_path, library_name, symbol_name)

        # Add the component instance
        return self.create_component_instance(
            schematic_path,
            library_name,
            symbol_name,
            reference=reference,
            value=value,
            footprint=footprint,
            x=x,
            y=y,
        )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    loader = DynamicSymbolLoader()

    print("\n=== Testing Dynamic Symbol Loader (Text-based) ===\n")

    print("1. Finding KiCad symbol library directories...")
    lib_dirs = loader.find_kicad_symbol_libraries()
    print(f"   Found {len(lib_dirs)} directories")

    print("\n2. Extracting symbols...")
    for lib, sym in [
        ("Device", "R"),
        ("Device", "C"),
        ("Device", "LED"),
        ("Device", "Q_NMOS"),
    ]:
        block = loader.extract_symbol_from_library(lib, sym)
        if block:
            print(f"   OK: {lib}:{sym} ({len(block)} chars)")
        else:
            print(f"   FAIL: {lib}:{sym}")

    print("\n3. Testing extends resolution...")
    block = loader.extract_symbol_from_library("Regulator_Switching", "LM2596S-5")
    if block and "LM2596S-12" in block:
        print(f"   OK: LM2596S-5 includes parent LM2596S-12 ({len(block)} chars)")
    else:
        print(f"   FAIL: extends not resolved")

    print("\nAll tests passed!")