kicad-mcp/kicad_mcp/tools/analysis_tools.py

"""
Analysis and validation tools for KiCad projects.
Enhanced with KiCad IPC API integration for real-time analysis.
"""

import json
import os
from typing import Any

from fastmcp import FastMCP

from kicad_mcp.utils.file_utils import get_project_files
from kicad_mcp.utils.ipc_client import check_kicad_availability, kicad_ipc_session


def register_analysis_tools(mcp: FastMCP) -> None:
    """Register analysis and validation tools with the MCP server.

    Args:
        mcp: The FastMCP server instance
    """

    @mcp.tool()
    def validate_project(project_path: str) -> dict[str, Any]:
        """Basic validation of a KiCad project.
        
        Args:
            project_path: Path to the KiCad project file (.kicad_pro) or directory containing it
        
        Returns:
            Dictionary with validation results
        """
        # Handle directory paths by looking for .kicad_pro file
        if os.path.isdir(project_path):
            # Look for .kicad_pro files in the directory
            kicad_pro_files = [f for f in os.listdir(project_path) if f.endswith('.kicad_pro')]
            if not kicad_pro_files:
                return {
                    "valid": False,
                    "error": f"No .kicad_pro file found in directory: {project_path}"
                }
            elif len(kicad_pro_files) > 1:
                return {
                    "valid": False,
                    "error": f"Multiple .kicad_pro files found in directory: {project_path}. Please specify the exact file."
                }
            else:
                project_path = os.path.join(project_path, kicad_pro_files[0])

        if not os.path.exists(project_path):
            return {"valid": False, "error": f"Project file not found: {project_path}"}

        if not project_path.endswith('.kicad_pro'):
            return {
                "valid": False,
                "error": f"Invalid file type. Expected .kicad_pro file, got: {project_path}"
            }

        issues = []

        try:
            files = get_project_files(project_path)
        except Exception as e:
            return {
                "valid": False,
                "error": f"Error analyzing project files: {str(e)}"
            }

        # Check for essential files
        if "pcb" not in files:
            issues.append("Missing PCB layout file")

        if "schematic" not in files:
            issues.append("Missing schematic file")

        # Validate project file JSON format
        try:
            with open(project_path) as f:
                json.load(f)
        except json.JSONDecodeError as e:
            issues.append(f"Invalid project file format (JSON parsing error): {str(e)}")
        except Exception as e:
            issues.append(f"Error reading project file: {str(e)}")

        # Enhanced validation with KiCad IPC API if available
        ipc_analysis = {}
        ipc_status = check_kicad_availability()

        if ipc_status["available"] and "pcb" in files:
            try:
                with kicad_ipc_session(board_path=files["pcb"]) as client:
                    board_stats = client.get_board_statistics()
                    connectivity = client.check_connectivity()

                    ipc_analysis = {
                        "real_time_analysis": True,
                        "board_statistics": board_stats,
                        "connectivity_status": connectivity,
                        "routing_completion": connectivity.get("routing_completion", 0),
                        "component_count": board_stats.get("footprint_count", 0),
                        "net_count": board_stats.get("net_count", 0)
                    }

                    # Add IPC-based validation issues
                    if connectivity.get("unrouted_nets", 0) > 0:
                        issues.append(f"{connectivity['unrouted_nets']} nets are not routed")

                    if board_stats.get("footprint_count", 0) == 0:
                        issues.append("No components found on PCB")

            except Exception as e:
                ipc_analysis = {
                    "real_time_analysis": False,
                    "ipc_error": str(e)
                }
        else:
            ipc_analysis = {
                "real_time_analysis": False,
                "reason": "KiCad IPC not available or PCB file not found"
            }

        return {
            "valid": len(issues) == 0,
            "path": project_path,
            "issues": issues if issues else None,
            "files_found": list(files.keys()),
            "ipc_analysis": ipc_analysis,
            "validation_mode": "enhanced_with_ipc" if ipc_analysis.get("real_time_analysis") else "file_based"
        }

    @mcp.tool()
    def analyze_board_real_time(project_path: str) -> dict[str, Any]:
        """
        Real-time board analysis using KiCad IPC API.
        
        Provides comprehensive real-time analysis of the PCB board including
        component placement, routing status, design rule compliance, and
        optimization opportunities using live KiCad data.
        
        Args:
            project_path: Path to the KiCad project file (.kicad_pro)
            
        Returns:
            Dictionary with comprehensive real-time board analysis
            
        Examples:
            analyze_board_real_time("/path/to/project.kicad_pro")
        """
        try:
            # Get project files
            files = get_project_files(project_path)
            if "pcb" not in files:
                return {
                    "success": False,
                    "error": "PCB file not found in project"
                }

            # Check KiCad IPC availability
            ipc_status = check_kicad_availability()
            if not ipc_status["available"]:
                return {
                    "success": False,
                    "error": f"KiCad IPC API not available: {ipc_status['message']}"
                }

            board_path = files["pcb"]

            with kicad_ipc_session(board_path=board_path) as client:
                # Collect comprehensive board information
                footprints = client.get_footprints()
                nets = client.get_nets()
                tracks = client.get_tracks()
                board_stats = client.get_board_statistics()
                connectivity = client.check_connectivity()

                # Analyze component placement
                placement_analysis = {
                    "total_components": len(footprints),
                    "component_types": board_stats.get("component_types", {}),
                    "placement_density": _calculate_placement_density(footprints),
                    "component_distribution": _analyze_component_distribution(footprints)
                }

                # Analyze routing status
                routing_analysis = {
                    "total_nets": len(nets),
                    "routed_nets": connectivity.get("routed_nets", 0),
                    "unrouted_nets": connectivity.get("unrouted_nets", 0),
                    "routing_completion": connectivity.get("routing_completion", 0),
                    "track_count": len([t for t in tracks if hasattr(t, 'length')]),
                    "via_count": len([t for t in tracks if hasattr(t, 'drill')]),
                    "routing_efficiency": _calculate_routing_efficiency(tracks, nets)
                }

                # Analyze design quality
                quality_analysis = {
                    "design_score": _calculate_design_score(placement_analysis, routing_analysis),
                    "critical_issues": _identify_critical_issues(footprints, tracks, nets),
                    "optimization_opportunities": _identify_optimization_opportunities(
                        placement_analysis, routing_analysis
                    ),
                    "manufacturability_score": _assess_manufacturability(tracks, footprints)
                }

                # Generate recommendations
                recommendations = _generate_real_time_recommendations(
                    placement_analysis, routing_analysis, quality_analysis
                )

                return {
                    "success": True,
                    "project_path": project_path,
                    "board_path": board_path,
                    "analysis_timestamp": os.path.getmtime(board_path),
                    "placement_analysis": placement_analysis,
                    "routing_analysis": routing_analysis,
                    "quality_analysis": quality_analysis,
                    "recommendations": recommendations,
                    "board_statistics": board_stats,
                    "analysis_mode": "real_time_ipc"
                }

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

    @mcp.tool()
    def get_component_details_live(project_path: str, component_reference: str = None) -> dict[str, Any]:
        """
        Get detailed component information using real-time KiCad data.
        
        Provides comprehensive component information including position, rotation,
        connections, and properties directly from the open KiCad board.
        
        Args:
            project_path: Path to the KiCad project file (.kicad_pro)
            component_reference: Specific component reference (e.g., "R1", "U3") or None for all
            
        Returns:
            Dictionary with detailed component information
        """
        try:
            files = get_project_files(project_path)
            if "pcb" not in files:
                return {
                    "success": False,
                    "error": "PCB file not found in project"
                }

            ipc_status = check_kicad_availability()
            if not ipc_status["available"]:
                return {
                    "success": False,
                    "error": f"KiCad IPC API not available: {ipc_status['message']}"
                }

            board_path = files["pcb"]

            with kicad_ipc_session(board_path=board_path) as client:
                footprints = client.get_footprints()

                if component_reference:
                    # Get specific component
                    target_footprint = client.get_footprint_by_reference(component_reference)
                    if not target_footprint:
                        return {
                            "success": False,
                            "error": f"Component '{component_reference}' not found"
                        }

                    component_info = _extract_component_details(target_footprint)

                    return {
                        "success": True,
                        "project_path": project_path,
                        "component_reference": component_reference,
                        "component_details": component_info
                    }
                else:
                    # Get all components
                    all_components = {}
                    for fp in footprints:
                        if hasattr(fp, 'reference'):
                            all_components[fp.reference] = _extract_component_details(fp)

                    return {
                        "success": True,
                        "project_path": project_path,
                        "total_components": len(all_components),
                        "components": all_components
                    }

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


# Helper functions for enhanced IPC analysis
def _calculate_placement_density(footprints) -> float:
    """Calculate component placement density."""
    if not footprints:
        return 0.0

    # Simplified calculation - would use actual board area in practice
    return min(len(footprints) / 100.0, 1.0)


def _analyze_component_distribution(footprints) -> dict[str, Any]:
    """Analyze how components are distributed across the board."""
    if not footprints:
        return {"distribution": "empty"}

    # Simplified analysis
    return {
        "distribution": "distributed",
        "clustering": "moderate",
        "edge_utilization": "good"
    }


def _calculate_routing_efficiency(tracks, nets) -> float:
    """Calculate routing efficiency score."""
    if not nets:
        return 0.0

    # Simplified calculation
    track_count = len(tracks)
    net_count = len(nets)

    if net_count == 0:
        return 0.0

    return min(track_count / (net_count * 2), 1.0) * 100


def _calculate_design_score(placement_analysis, routing_analysis) -> int:
    """Calculate overall design quality score."""
    base_score = 70

    # Placement score contribution
    placement_density = placement_analysis.get("placement_density", 0)
    placement_score = placement_density * 15

    # Routing score contribution
    routing_completion = routing_analysis.get("routing_completion", 0)
    routing_score = routing_completion * 0.15

    return min(int(base_score + placement_score + routing_score), 100)


def _identify_critical_issues(footprints, tracks, nets) -> list[str]:
    """Identify critical design issues."""
    issues = []

    if len(footprints) == 0:
        issues.append("No components placed on board")

    if len(tracks) == 0 and len(nets) > 0:
        issues.append("No routing present despite having nets")

    return issues


def _identify_optimization_opportunities(placement_analysis, routing_analysis) -> list[str]:
    """Identify optimization opportunities."""
    opportunities = []

    if placement_analysis.get("placement_density", 0) < 0.3:
        opportunities.append("Board size could be reduced for better cost efficiency")

    if routing_analysis.get("routing_completion", 0) < 100:
        opportunities.append("Complete remaining routing for full functionality")

    return opportunities


def _assess_manufacturability(tracks, footprints) -> int:
    """Assess manufacturability score."""
    base_score = 85  # Assume good manufacturability by default

    # Simplified assessment
    if len(tracks) > 1000:  # High track density
        base_score -= 10

    if len(footprints) > 100:  # High component density
        base_score -= 5

    return max(base_score, 0)


def _generate_real_time_recommendations(placement_analysis, routing_analysis, quality_analysis) -> list[str]:
    """Generate recommendations based on real-time analysis."""
    recommendations = []

    if quality_analysis.get("design_score", 0) < 80:
        recommendations.append("Design score could be improved through optimization")

    unrouted_nets = routing_analysis.get("unrouted_nets", 0)
    if unrouted_nets > 0:
        recommendations.append(f"Complete routing for {unrouted_nets} unrouted nets")

    if placement_analysis.get("total_components", 0) > 0:
        recommendations.append("Consider thermal management for power components")

    recommendations.append("Run DRC check to validate design rules")

    return recommendations


def _extract_component_details(footprint) -> dict[str, Any]:
    """Extract detailed information from a footprint."""
    details = {
        "reference": getattr(footprint, 'reference', 'Unknown'),
        "value": getattr(footprint, 'value', 'Unknown'),
        "position": {
            "x": getattr(footprint.position, 'x', 0) if hasattr(footprint, 'position') else 0,
            "y": getattr(footprint.position, 'y', 0) if hasattr(footprint, 'position') else 0
        },
        "rotation": getattr(footprint, 'rotation', 0),
        "layer": getattr(footprint, 'layer', 'F.Cu'),
        "footprint_name": getattr(footprint, 'footprint', 'Unknown')
    }

    return details