kicad-mcp/kicad_mcp/tools/ai_tools.py

"""
AI/LLM Integration Tools for KiCad MCP Server.

Provides intelligent analysis and recommendations for KiCad designs including
smart component suggestions, automated design rule recommendations, and layout optimization.
"""

from typing import Any

from fastmcp import FastMCP

from kicad_mcp.utils.component_utils import ComponentType, get_component_type
from kicad_mcp.utils.file_utils import get_project_files
from kicad_mcp.utils.netlist_parser import parse_netlist_file
from kicad_mcp.utils.pattern_recognition import analyze_circuit_patterns


def register_ai_tools(mcp: FastMCP) -> None:
    """Register AI/LLM integration tools with the MCP server."""

    @mcp.tool()
    def suggest_components_for_circuit(project_path: str, circuit_function: str = None) -> dict[str, Any]:
        """
        Analyze circuit patterns and suggest appropriate components.
        
        Uses circuit analysis to identify incomplete circuits and suggest
        missing components based on common design patterns and best practices.
        
        Args:
            project_path: Path to the KiCad project file (.kicad_pro)
            circuit_function: Optional description of intended circuit function
            
        Returns:
            Dictionary with component suggestions categorized by circuit type
            
        Examples:
            suggest_components_for_circuit("/path/to/project.kicad_pro")
            suggest_components_for_circuit("/path/to/project.kicad_pro", "audio amplifier")
        """
        try:
            # Get project files
            files = get_project_files(project_path)
            if "schematic" not in files:
                return {
                    "success": False,
                    "error": "Schematic file not found in project"
                }

            schematic_file = files["schematic"]

            # Analyze existing circuit patterns
            patterns = analyze_circuit_patterns(schematic_file)

            # Parse netlist for component analysis
            try:
                netlist_data = parse_netlist_file(schematic_file)
                components = netlist_data.get("components", [])
            except:
                components = []

            # Generate suggestions based on patterns
            suggestions = _generate_component_suggestions(patterns, components, circuit_function)

            return {
                "success": True,
                "project_path": project_path,
                "circuit_analysis": {
                    "identified_patterns": list(patterns.keys()),
                    "component_count": len(components),
                    "missing_patterns": _identify_missing_patterns(patterns, components)
                },
                "component_suggestions": suggestions,
                "design_recommendations": _generate_design_recommendations(patterns, components),
                "implementation_notes": [
                    "Review suggested components for compatibility with existing design",
                    "Verify component ratings match circuit requirements",
                    "Consider thermal management for power components",
                    "Check component availability and cost before finalizing"
                ]
            }

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

    @mcp.tool()
    def recommend_design_rules(project_path: str, target_technology: str = "standard") -> dict[str, Any]:
        """
        Generate automated design rule recommendations based on circuit analysis.
        
        Analyzes the circuit topology, component types, and signal characteristics
        to recommend appropriate design rules for the specific application.
        
        Args:
            project_path: Path to the KiCad project file (.kicad_pro)
            target_technology: Target technology ("standard", "hdi", "rf", "automotive")
            
        Returns:
            Dictionary with customized design rule recommendations
            
        Examples:
            recommend_design_rules("/path/to/project.kicad_pro")
            recommend_design_rules("/path/to/project.kicad_pro", "rf")
        """
        try:
            # Get project files
            files = get_project_files(project_path)

            analysis_data = {}

            # Analyze schematic if available
            if "schematic" in files:
                patterns = analyze_circuit_patterns(files["schematic"])
                analysis_data["patterns"] = patterns

                try:
                    netlist_data = parse_netlist_file(files["schematic"])
                    analysis_data["components"] = netlist_data.get("components", [])
                except:
                    analysis_data["components"] = []

            # Analyze PCB if available
            if "pcb" in files:
                pcb_analysis = _analyze_pcb_characteristics(files["pcb"])
                analysis_data["pcb"] = pcb_analysis

            # Generate design rules based on analysis
            design_rules = _generate_design_rules(analysis_data, target_technology)

            return {
                "success": True,
                "project_path": project_path,
                "target_technology": target_technology,
                "circuit_analysis": {
                    "identified_patterns": list(analysis_data.get("patterns", {}).keys()),
                    "component_types": _categorize_components(analysis_data.get("components", [])),
                    "signal_types": _identify_signal_types(analysis_data.get("patterns", {}))
                },
                "recommended_rules": design_rules,
                "rule_justifications": _generate_rule_justifications(design_rules, analysis_data),
                "implementation_priority": _prioritize_rules(design_rules)
            }

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

    @mcp.tool()
    def optimize_pcb_layout(project_path: str, optimization_goals: list[str] = None) -> dict[str, Any]:
        """
        Analyze PCB layout and provide optimization suggestions.
        
        Reviews component placement, routing, and design practices to suggest
        improvements for signal integrity, thermal management, and manufacturability.
        
        Args:
            project_path: Path to the KiCad project file (.kicad_pro)
            optimization_goals: List of optimization priorities (e.g., ["signal_integrity", "thermal", "cost"])
            
        Returns:
            Dictionary with layout optimization recommendations
            
        Examples:
            optimize_pcb_layout("/path/to/project.kicad_pro")
            optimize_pcb_layout("/path/to/project.kicad_pro", ["signal_integrity", "cost"])
        """
        try:
            if not optimization_goals:
                optimization_goals = ["signal_integrity", "thermal", "manufacturability"]

            # Get project files
            files = get_project_files(project_path)

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

            pcb_file = files["pcb"]

            # Analyze current layout
            layout_analysis = _analyze_pcb_layout(pcb_file)

            # Get circuit context from schematic if available
            circuit_context = {}
            if "schematic" in files:
                patterns = analyze_circuit_patterns(files["schematic"])
                circuit_context = {"patterns": patterns}

            # Generate optimization suggestions
            optimizations = _generate_layout_optimizations(
                layout_analysis, circuit_context, optimization_goals
            )

            return {
                "success": True,
                "project_path": project_path,
                "optimization_goals": optimization_goals,
                "layout_analysis": {
                    "component_density": layout_analysis.get("component_density", 0),
                    "routing_utilization": layout_analysis.get("routing_utilization", {}),
                    "thermal_zones": layout_analysis.get("thermal_zones", []),
                    "critical_signals": layout_analysis.get("critical_signals", [])
                },
                "optimization_suggestions": optimizations,
                "implementation_steps": _generate_implementation_steps(optimizations),
                "expected_benefits": _calculate_optimization_benefits(optimizations)
            }

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

    @mcp.tool()
    def analyze_design_completeness(project_path: str) -> dict[str, Any]:
        """
        Analyze design completeness and suggest missing elements.
        
        Performs comprehensive analysis to identify missing components, 
        incomplete circuits, and design gaps that should be addressed.
        
        Args:
            project_path: Path to the KiCad project file (.kicad_pro)
            
        Returns:
            Dictionary with completeness analysis and improvement suggestions
        """
        try:
            files = get_project_files(project_path)

            completeness_analysis = {
                "schematic_completeness": 0,
                "pcb_completeness": 0,
                "design_gaps": [],
                "missing_elements": [],
                "verification_status": {}
            }

            # Analyze schematic completeness
            if "schematic" in files:
                schematic_analysis = _analyze_schematic_completeness(files["schematic"])
                completeness_analysis.update(schematic_analysis)

            # Analyze PCB completeness
            if "pcb" in files:
                pcb_analysis = _analyze_pcb_completeness(files["pcb"])
                completeness_analysis["pcb_completeness"] = pcb_analysis["completeness_score"]
                completeness_analysis["design_gaps"].extend(pcb_analysis["gaps"])

            # Overall completeness score
            overall_score = (
                completeness_analysis["schematic_completeness"] * 0.6 +
                completeness_analysis["pcb_completeness"] * 0.4
            )

            return {
                "success": True,
                "project_path": project_path,
                "completeness_score": round(overall_score, 1),
                "analysis_details": completeness_analysis,
                "priority_actions": _prioritize_completeness_actions(completeness_analysis),
                "design_checklist": _generate_design_checklist(completeness_analysis),
                "recommendations": _generate_completeness_recommendations(completeness_analysis)
            }

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


# Helper functions for component suggestions
def _generate_component_suggestions(patterns: dict, components: list, circuit_function: str = None) -> dict[str, list]:
    """Generate component suggestions based on circuit analysis."""
    suggestions = {
        "power_management": [],
        "signal_conditioning": [],
        "protection": [],
        "filtering": [],
        "interface": [],
        "passive_components": []
    }

    # Analyze existing components
    component_types = [get_component_type(comp.get("value", "")) for comp in components]

    # Power management suggestions
    if "power_supply" in patterns:
        if ComponentType.VOLTAGE_REGULATOR not in component_types:
            suggestions["power_management"].append({
                "component": "Voltage Regulator",
                "suggestion": "Add voltage regulator for stable power supply",
                "examples": ["LM7805", "AMS1117-3.3", "LM2596"]
            })

        if ComponentType.CAPACITOR not in component_types:
            suggestions["power_management"].append({
                "component": "Decoupling Capacitors",
                "suggestion": "Add decoupling capacitors near power pins",
                "examples": ["100nF ceramic", "10uF tantalum", "1000uF electrolytic"]
            })

    # Signal conditioning suggestions
    if "amplifier" in patterns:
        if not any("op" in comp.get("value", "").lower() for comp in components):
            suggestions["signal_conditioning"].append({
                "component": "Operational Amplifier",
                "suggestion": "Consider op-amp for signal amplification",
                "examples": ["LM358", "TL072", "OPA2134"]
            })

    # Protection suggestions
    if "microcontroller" in patterns or "processor" in patterns:
        if ComponentType.FUSE not in component_types:
            suggestions["protection"].append({
                "component": "Fuse or PTC Resettable Fuse",
                "suggestion": "Add overcurrent protection",
                "examples": ["1A fuse", "PPTC 0.5A", "Polyfuse 1A"]
            })

        if not any("esd" in comp.get("value", "").lower() for comp in components):
            suggestions["protection"].append({
                "component": "ESD Protection",
                "suggestion": "Add ESD protection for I/O pins",
                "examples": ["TVS diode", "ESD suppressors", "Varistors"]
            })

    # Filtering suggestions
    if any(pattern in patterns for pattern in ["switching_converter", "motor_driver"]):
        suggestions["filtering"].append({
            "component": "EMI Filter",
            "suggestion": "Add EMI filtering for switching circuits",
            "examples": ["Common mode choke", "Ferrite beads", "Pi filter"]
        })

    # Interface suggestions based on circuit function
    if circuit_function:
        function_lower = circuit_function.lower()
        if "audio" in function_lower:
            suggestions["interface"].extend([
                {
                    "component": "Audio Jack",
                    "suggestion": "Add audio input/output connector",
                    "examples": ["3.5mm jack", "RCA connector", "XLR"]
                },
                {
                    "component": "Audio Coupling Capacitor",
                    "suggestion": "AC coupling for audio signals",
                    "examples": ["10uF", "47uF", "100uF"]
                }
            ])

        if "usb" in function_lower or "communication" in function_lower:
            suggestions["interface"].append({
                "component": "USB Connector",
                "suggestion": "Add USB interface for communication",
                "examples": ["USB-A", "USB-C", "Micro-USB"]
            })

    return suggestions


def _identify_missing_patterns(patterns: dict, components: list) -> list[str]:
    """Identify common circuit patterns that might be missing."""
    missing_patterns = []

    has_digital_components = any(
        comp.get("value", "").lower() in ["microcontroller", "processor", "mcu"]
        for comp in components
    )

    if has_digital_components:
        if "crystal_oscillator" not in patterns:
            missing_patterns.append("crystal_oscillator")
        if "reset_circuit" not in patterns:
            missing_patterns.append("reset_circuit")
        if "power_supply" not in patterns:
            missing_patterns.append("power_supply")

    return missing_patterns


def _generate_design_recommendations(patterns: dict, components: list) -> list[str]:
    """Generate general design recommendations."""
    recommendations = []

    if "power_supply" not in patterns and len(components) > 5:
        recommendations.append("Consider adding dedicated power supply regulation")

    if len(components) > 20 and "decoupling" not in patterns:
        recommendations.append("Add decoupling capacitors for noise reduction")

    if any("high_freq" in str(pattern) for pattern in patterns):
        recommendations.append("Consider transmission line effects for high-frequency signals")

    return recommendations


# Helper functions for design rules
def _analyze_pcb_characteristics(pcb_file: str) -> dict[str, Any]:
    """Analyze PCB file for design rule recommendations."""
    # This is a simplified analysis - in practice would parse the PCB file
    return {
        "layer_count": 2,  # Default assumption
        "min_trace_width": 0.1,
        "min_via_size": 0.2,
        "component_density": "medium"
    }


def _generate_design_rules(analysis_data: dict, target_technology: str) -> dict[str, dict]:
    """Generate design rules based on analysis and technology target."""
    base_rules = {
        "trace_width": {"min": 0.1, "preferred": 0.15, "unit": "mm"},
        "via_size": {"min": 0.2, "preferred": 0.3, "unit": "mm"},
        "clearance": {"min": 0.1, "preferred": 0.15, "unit": "mm"},
        "annular_ring": {"min": 0.05, "preferred": 0.1, "unit": "mm"}
    }

    # Adjust rules based on technology
    if target_technology == "hdi":
        base_rules["trace_width"]["min"] = 0.075
        base_rules["via_size"]["min"] = 0.1
        base_rules["clearance"]["min"] = 0.075
    elif target_technology == "rf":
        base_rules["trace_width"]["preferred"] = 0.2
        base_rules["clearance"]["preferred"] = 0.2
    elif target_technology == "automotive":
        base_rules["trace_width"]["min"] = 0.15
        base_rules["clearance"]["min"] = 0.15

    # Adjust based on patterns
    patterns = analysis_data.get("patterns", {})
    if "power_supply" in patterns:
        base_rules["power_trace_width"] = {"min": 0.3, "preferred": 0.5, "unit": "mm"}

    if "high_speed" in patterns:
        base_rules["differential_impedance"] = {"target": 100, "tolerance": 10, "unit": "ohm"}
        base_rules["single_ended_impedance"] = {"target": 50, "tolerance": 5, "unit": "ohm"}

    return base_rules


def _categorize_components(components: list) -> dict[str, int]:
    """Categorize components by type."""
    categories = {}
    for comp in components:
        comp_type = get_component_type(comp.get("value", ""))
        category_name = comp_type.name.lower() if comp_type != ComponentType.UNKNOWN else "other"
        categories[category_name] = categories.get(category_name, 0) + 1

    return categories


def _identify_signal_types(patterns: dict) -> list[str]:
    """Identify signal types based on circuit patterns."""
    signal_types = []

    if "power_supply" in patterns:
        signal_types.append("power")
    if "amplifier" in patterns:
        signal_types.append("analog")
    if "microcontroller" in patterns:
        signal_types.extend(["digital", "clock"])
    if "crystal_oscillator" in patterns:
        signal_types.append("high_frequency")

    return list(set(signal_types))


def _generate_rule_justifications(design_rules: dict, analysis_data: dict) -> dict[str, str]:
    """Generate justifications for recommended design rules."""
    justifications = {}

    patterns = analysis_data.get("patterns", {})

    if "trace_width" in design_rules:
        justifications["trace_width"] = "Based on current carrying capacity and manufacturing constraints"

    if "power_supply" in patterns and "power_trace_width" in design_rules:
        justifications["power_trace_width"] = "Wider traces for power distribution to reduce voltage drop"

    if "high_speed" in patterns and "differential_impedance" in design_rules:
        justifications["differential_impedance"] = "Controlled impedance required for high-speed signals"

    return justifications


def _prioritize_rules(design_rules: dict) -> list[str]:
    """Prioritize design rules by implementation importance."""
    priority_order = []

    if "clearance" in design_rules:
        priority_order.append("clearance")
    if "trace_width" in design_rules:
        priority_order.append("trace_width")
    if "via_size" in design_rules:
        priority_order.append("via_size")
    if "power_trace_width" in design_rules:
        priority_order.append("power_trace_width")
    if "differential_impedance" in design_rules:
        priority_order.append("differential_impedance")

    return priority_order


# Helper functions for layout optimization
def _analyze_pcb_layout(pcb_file: str) -> dict[str, Any]:
    """Analyze PCB layout for optimization opportunities."""
    # Simplified analysis - would parse actual PCB file
    return {
        "component_density": 0.6,
        "routing_utilization": {"top": 0.4, "bottom": 0.3},
        "thermal_zones": ["high_power_area"],
        "critical_signals": ["clock", "reset", "power"]
    }


def _generate_layout_optimizations(layout_analysis: dict, circuit_context: dict, goals: list[str]) -> dict[str, list]:
    """Generate layout optimization suggestions."""
    optimizations = {
        "placement": [],
        "routing": [],
        "thermal": [],
        "signal_integrity": [],
        "manufacturability": []
    }

    if "signal_integrity" in goals:
        optimizations["signal_integrity"].extend([
            "Keep high-speed traces short and direct",
            "Minimize via count on critical signals",
            "Use ground planes for return current paths"
        ])

    if "thermal" in goals:
        optimizations["thermal"].extend([
            "Spread heat-generating components across the board",
            "Add thermal vias under power components",
            "Consider copper pour for heat dissipation"
        ])

    if "cost" in goals or "manufacturability" in goals:
        optimizations["manufacturability"].extend([
            "Use standard via sizes and trace widths",
            "Minimize layer count where possible",
            "Avoid blind/buried vias unless necessary"
        ])

    return optimizations


def _generate_implementation_steps(optimizations: dict) -> list[str]:
    """Generate step-by-step implementation guide."""
    steps = []

    if optimizations.get("placement"):
        steps.append("1. Review component placement for optimal positioning")

    if optimizations.get("routing"):
        steps.append("2. Re-route critical signals following guidelines")

    if optimizations.get("thermal"):
        steps.append("3. Implement thermal management improvements")

    if optimizations.get("signal_integrity"):
        steps.append("4. Optimize signal integrity aspects")

    steps.append("5. Run DRC and electrical rules check")
    steps.append("6. Verify design meets all requirements")

    return steps


def _calculate_optimization_benefits(optimizations: dict) -> dict[str, str]:
    """Calculate expected benefits from optimizations."""
    benefits = {}

    if optimizations.get("signal_integrity"):
        benefits["signal_integrity"] = "Improved noise margin and reduced EMI"

    if optimizations.get("thermal"):
        benefits["thermal"] = "Better thermal performance and component reliability"

    if optimizations.get("manufacturability"):
        benefits["manufacturability"] = "Reduced manufacturing cost and higher yield"

    return benefits


# Helper functions for design completeness
def _analyze_schematic_completeness(schematic_file: str) -> dict[str, Any]:
    """Analyze schematic completeness."""
    try:
        patterns = analyze_circuit_patterns(schematic_file)
        netlist_data = parse_netlist_file(schematic_file)
        components = netlist_data.get("components", [])

        completeness_score = 70  # Base score
        missing_elements = []

        # Check for essential patterns
        if "power_supply" in patterns:
            completeness_score += 10
        else:
            missing_elements.append("power_supply_regulation")

        if len(components) > 5:
            if "decoupling" not in patterns:
                missing_elements.append("decoupling_capacitors")
            else:
                completeness_score += 10

        return {
            "schematic_completeness": min(completeness_score, 100),
            "missing_elements": missing_elements,
            "design_gaps": [],
            "verification_status": {"nets": "checked", "components": "verified"}
        }

    except Exception:
        return {
            "schematic_completeness": 50,
            "missing_elements": ["analysis_failed"],
            "design_gaps": [],
            "verification_status": {"status": "error"}
        }


def _analyze_pcb_completeness(pcb_file: str) -> dict[str, Any]:
    """Analyze PCB completeness."""
    # Simplified analysis
    return {
        "completeness_score": 80,
        "gaps": ["silkscreen_labels", "test_points"]
    }


def _prioritize_completeness_actions(analysis: dict) -> list[str]:
    """Prioritize actions for improving design completeness."""
    actions = []

    if "power_supply_regulation" in analysis.get("missing_elements", []):
        actions.append("Add power supply regulation circuit")

    if "decoupling_capacitors" in analysis.get("missing_elements", []):
        actions.append("Add decoupling capacitors near ICs")

    if analysis.get("schematic_completeness", 0) < 80:
        actions.append("Complete schematic design")

    if analysis.get("pcb_completeness", 0) < 80:
        actions.append("Finish PCB layout")

    return actions


def _generate_design_checklist(analysis: dict) -> list[dict[str, Any]]:
    """Generate design verification checklist."""
    checklist = [
        {"item": "Schematic review complete", "status": "complete" if analysis.get("schematic_completeness", 0) > 90 else "pending"},
        {"item": "Component values verified", "status": "complete" if "components" in analysis.get("verification_status", {}) else "pending"},
        {"item": "Power supply design", "status": "complete" if "power_supply_regulation" not in analysis.get("missing_elements", []) else "pending"},
        {"item": "Signal integrity considerations", "status": "pending"},
        {"item": "Thermal management", "status": "pending"},
        {"item": "Manufacturing readiness", "status": "pending"}
    ]

    return checklist


def _generate_completeness_recommendations(analysis: dict) -> list[str]:
    """Generate recommendations for improving completeness."""
    recommendations = []

    completeness = analysis.get("schematic_completeness", 0)

    if completeness < 70:
        recommendations.append("Focus on completing core circuit functionality")
    elif completeness < 85:
        recommendations.append("Add protective and filtering components")
    else:
        recommendations.append("Review design for optimization opportunities")

    if analysis.get("missing_elements"):
        recommendations.append(f"Address missing elements: {', '.join(analysis['missing_elements'])}")

    return recommendations