kicad-mcp/kicad_mcp/tools/pattern_tools.py

"""
Circuit pattern recognition tools for KiCad schematics.
"""

import os
from typing import Dict, List, Any, Optional
from mcp.server.fastmcp import FastMCP, Context

from kicad_mcp.utils.file_utils import get_project_files
from kicad_mcp.utils.netlist_parser import extract_netlist, analyze_netlist
from kicad_mcp.utils.pattern_recognition import (
    identify_power_supplies,
    identify_amplifiers,
    identify_filters,
    identify_oscillators,
    identify_digital_interfaces,
    identify_microcontrollers,
    identify_sensor_interfaces,
)


def register_pattern_tools(mcp: FastMCP) -> None:
    """Register circuit pattern recognition tools with the MCP server.

    Args:
        mcp: The FastMCP server instance
    """

    @mcp.tool()
    async def identify_circuit_patterns(schematic_path: str, ctx: Context) -> Dict[str, Any]:
        """Identify common circuit patterns in a KiCad schematic.

        This tool analyzes a schematic to recognize common circuit blocks such as:
        - Power supply circuits (linear regulators, switching converters)
        - Amplifier circuits (op-amps, transistor amplifiers)
        - Filter circuits (RC, LC, active filters)
        - Digital interfaces (I2C, SPI, UART)
        - Microcontroller circuits
        - And more

        Args:
            schematic_path: Path to the KiCad schematic file (.kicad_sch)
            ctx: MCP context for progress reporting

        Returns:
            Dictionary with identified circuit patterns
        """
        if not os.path.exists(schematic_path):
            ctx.info(f"Schematic file not found: {schematic_path}")
            return {"success": False, "error": f"Schematic file not found: {schematic_path}"}

        # Report progress
        await ctx.report_progress(10, 100)
        ctx.info(f"Loading schematic file: {os.path.basename(schematic_path)}")

        try:
            # Extract netlist information
            await ctx.report_progress(20, 100)
            ctx.info("Parsing schematic structure...")

            netlist_data = extract_netlist(schematic_path)

            if "error" in netlist_data:
                ctx.info(f"Error extracting netlist: {netlist_data['error']}")
                return {"success": False, "error": netlist_data["error"]}

            # Analyze components and nets
            await ctx.report_progress(30, 100)
            ctx.info("Analyzing components and connections...")

            components = netlist_data.get("components", {})
            nets = netlist_data.get("nets", {})

            # Start pattern recognition
            await ctx.report_progress(50, 100)
            ctx.info("Identifying circuit patterns...")

            identified_patterns = {
                "power_supply_circuits": [],
                "amplifier_circuits": [],
                "filter_circuits": [],
                "oscillator_circuits": [],
                "digital_interface_circuits": [],
                "microcontroller_circuits": [],
                "sensor_interface_circuits": [],
                "other_patterns": [],
            }

            # Identify power supply circuits
            await ctx.report_progress(60, 100)
            identified_patterns["power_supply_circuits"] = identify_power_supplies(components, nets)

            # Identify amplifier circuits
            await ctx.report_progress(70, 100)
            identified_patterns["amplifier_circuits"] = identify_amplifiers(components, nets)

            # Identify filter circuits
            await ctx.report_progress(75, 100)
            identified_patterns["filter_circuits"] = identify_filters(components, nets)

            # Identify oscillator circuits
            await ctx.report_progress(80, 100)
            identified_patterns["oscillator_circuits"] = identify_oscillators(components, nets)

            # Identify digital interface circuits
            await ctx.report_progress(85, 100)
            identified_patterns["digital_interface_circuits"] = identify_digital_interfaces(
                components, nets
            )

            # Identify microcontroller circuits
            await ctx.report_progress(90, 100)
            identified_patterns["microcontroller_circuits"] = identify_microcontrollers(components)

            # Identify sensor interface circuits
            await ctx.report_progress(95, 100)
            identified_patterns["sensor_interface_circuits"] = identify_sensor_interfaces(
                components, nets
            )

            # Build result
            result = {
                "success": True,
                "schematic_path": schematic_path,
                "component_count": netlist_data["component_count"],
                "identified_patterns": identified_patterns,
            }

            # Count total patterns
            total_patterns = sum(len(patterns) for patterns in identified_patterns.values())
            result["total_patterns_found"] = total_patterns

            # Complete progress
            await ctx.report_progress(100, 100)
            ctx.info(f"Pattern recognition complete. Found {total_patterns} circuit patterns.")

            return result

        except Exception as e:
            ctx.info(f"Error identifying circuit patterns: {str(e)}")
            return {"success": False, "error": str(e)}

    @mcp.tool()
    def analyze_project_circuit_patterns(project_path: str) -> Dict[str, Any]:
        """Identify circuit patterns in a KiCad project's schematic.

        Args:
            project_path: Path to the KiCad project file (.kicad_pro)

        Returns:
            Dictionary with identified circuit patterns
        """
        if not os.path.exists(project_path):
            return {"success": False, "error": f"Project not found: {project_path}"}

        # Get the schematic file
        try:
            files = get_project_files(project_path)

            if "schematic" not in files:
                return {"success": False, "error": "Schematic file not found in project"}

            schematic_path = files["schematic"]

            # Identify patterns in the schematic - call synchronous version
            if not os.path.exists(schematic_path):
                return {"success": False, "error": f"Schematic file not found: {schematic_path}"}

            # Extract netlist data
            netlist_data = extract_netlist(schematic_path)
            if not netlist_data:
                return {"success": False, "error": "Failed to extract netlist from schematic"}

            components, nets = analyze_netlist(netlist_data)
            
            # Identify patterns
            identified_patterns = {}
            identified_patterns["power_supply_circuits"] = identify_power_supplies(components, nets)
            identified_patterns["amplifier_circuits"] = identify_amplifiers(components, nets)
            identified_patterns["filter_circuits"] = identify_filters(components, nets)
            identified_patterns["oscillator_circuits"] = identify_oscillators(components, nets)
            identified_patterns["digital_interface_circuits"] = identify_digital_interfaces(components, nets)
            identified_patterns["microcontroller_circuits"] = identify_microcontrollers(components)
            identified_patterns["sensor_interface_circuits"] = identify_sensor_interfaces(components, nets)

            result = {
                "success": True,
                "schematic_path": schematic_path,
                "patterns": identified_patterns,
                "total_patterns_found": sum(len(patterns) for patterns in identified_patterns.values())
            }

            # Add project path to result
            if "success" in result and result["success"]:
                result["project_path"] = project_path

            return result

        except Exception as e:
            return {"success": False, "error": str(e)}