kicad-mcp/kicad_mcp/tools/model3d_tools.py

"""
3D Model Analysis Tools for KiCad MCP Server.

Provides MCP tools for analyzing 3D models, mechanical constraints,
and visualization data from KiCad PCB files.
"""

import json
from typing import Any, Dict

from fastmcp import FastMCP
from kicad_mcp.utils.model3d_analyzer import (
    analyze_pcb_3d_models,
    get_mechanical_constraints,
    Model3DAnalyzer
)
from kicad_mcp.utils.path_validator import validate_kicad_file


def register_model3d_tools(mcp: FastMCP) -> None:
    """Register 3D model analysis tools with the MCP server."""
    
    @mcp.tool()
    def analyze_3d_models(pcb_file_path: str) -> Dict[str, Any]:
        """
        Analyze 3D models and mechanical aspects of a KiCad PCB file.
        
        Extracts 3D component information, board dimensions, clearance violations,
        and generates data suitable for 3D visualization.
        
        Args:
            pcb_file_path: Full path to the .kicad_pcb file to analyze
            
        Returns:
            Dictionary containing 3D analysis results including:
            - board_dimensions: Physical board size and outline  
            - components: List of 3D components with positions and models
            - height_analysis: Component height statistics
            - clearance_violations: Detected mechanical issues
            - stats: Summary statistics
            
        Examples:
            analyze_3d_models("/path/to/my_board.kicad_pcb")
            analyze_3d_models("~/kicad_projects/robot_controller/robot.kicad_pcb")
        """
        try:
            # Validate the PCB file path
            validated_path = validate_kicad_file(pcb_file_path, "pcb")
            
            # Perform 3D analysis
            result = analyze_pcb_3d_models(validated_path)
            
            return {
                "success": True,
                "pcb_file": validated_path,
                "analysis": result
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "pcb_file": pcb_file_path
            }
    
    @mcp.tool()
    def check_mechanical_constraints(pcb_file_path: str) -> Dict[str, Any]:
        """
        Check mechanical constraints and clearances in a KiCad PCB.
        
        Performs comprehensive mechanical analysis including component clearances,
        board edge distances, height constraints, and identifies potential
        manufacturing or assembly issues.
        
        Args:
            pcb_file_path: Path to the .kicad_pcb file to analyze
            
        Returns:
            Dictionary containing mechanical analysis results:
            - constraints: List of constraint violations
            - clearance_violations: Detailed clearance issues
            - board_dimensions: Physical board properties
            - recommendations: Suggested improvements
        """
        try:
            validated_path = validate_kicad_file(pcb_file_path, "pcb")
            
            # Perform mechanical analysis
            analysis = get_mechanical_constraints(validated_path)
            
            # Generate recommendations
            recommendations = []
            
            if analysis.height_analysis["max"] > 5.0:
                recommendations.append("Consider using lower profile components to reduce board height")
            
            if len(analysis.clearance_violations) > 0:
                recommendations.append("Review component placement to resolve clearance violations")
            
            if analysis.board_dimensions.width > 80 or analysis.board_dimensions.height > 80:
                recommendations.append("Large board size may increase manufacturing costs")
            
            return {
                "success": True,
                "pcb_file": validated_path,
                "constraints": analysis.mechanical_constraints,
                "clearance_violations": [
                    {
                        "type": v["type"],
                        "components": [v.get("component1", ""), v.get("component2", ""), v.get("component", "")],
                        "distance": v["distance"],
                        "required": v["required_clearance"],
                        "severity": v["severity"]
                    }
                    for v in analysis.clearance_violations
                ],
                "board_dimensions": {
                    "width_mm": analysis.board_dimensions.width,
                    "height_mm": analysis.board_dimensions.height,
                    "thickness_mm": analysis.board_dimensions.thickness,
                    "area_mm2": analysis.board_dimensions.width * analysis.board_dimensions.height
                },
                "height_analysis": analysis.height_analysis,
                "recommendations": recommendations,
                "component_count": len(analysis.components)
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "pcb_file": pcb_file_path
            }
    
    @mcp.tool()
    def generate_3d_visualization_json(pcb_file_path: str, output_path: str = None) -> Dict[str, Any]:
        """
        Generate JSON data file for 3D visualization of PCB.
        
        Creates a structured JSON file containing all necessary data for
        3D visualization tools, including component positions, board outline,
        and model references.
        
        Args:
            pcb_file_path: Path to the .kicad_pcb file
            output_path: Optional path for output JSON file (defaults to same dir as PCB)
            
        Returns:
            Dictionary with generation results and file path
        """
        try:
            validated_path = validate_kicad_file(pcb_file_path, "pcb")
            
            # Generate visualization data
            viz_data = analyze_pcb_3d_models(validated_path)
            
            # Determine output path
            if not output_path:
                output_path = validated_path.replace('.kicad_pcb', '_3d_viz.json')
            
            # Save visualization data
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(viz_data, f, indent=2)
            
            return {
                "success": True,
                "pcb_file": validated_path,
                "output_file": output_path,
                "component_count": viz_data.get("stats", {}).get("total_components", 0),
                "models_found": viz_data.get("stats", {}).get("components_with_3d_models", 0),
                "board_size": f"{viz_data.get('board_dimensions', {}).get('width', 0):.1f}x{viz_data.get('board_dimensions', {}).get('height', 0):.1f}mm"
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "pcb_file": pcb_file_path
            }
    
    @mcp.tool()
    def component_height_distribution(pcb_file_path: str) -> Dict[str, Any]:
        """
        Analyze the height distribution of components on a PCB.
        
        Provides detailed analysis of component heights, useful for
        determining enclosure requirements and assembly considerations.
        
        Args:
            pcb_file_path: Path to the .kicad_pcb file
            
        Returns:
            Height distribution analysis with statistics and component breakdown
        """
        try:
            validated_path = validate_kicad_file(pcb_file_path, "pcb")
            
            analyzer = Model3DAnalyzer(validated_path)
            components = analyzer.extract_3d_components()
            height_analysis = analyzer.analyze_component_heights(components)
            
            # Categorize components by height
            height_categories = {
                "very_low": [],    # < 1mm
                "low": [],         # 1-2mm
                "medium": [],      # 2-5mm
                "high": [],        # 5-10mm
                "very_high": []    # > 10mm
            }
            
            for comp in components:
                height = analyzer._estimate_component_height(comp)
                
                if height < 1.0:
                    height_categories["very_low"].append((comp.reference, height))
                elif height < 2.0:
                    height_categories["low"].append((comp.reference, height))
                elif height < 5.0:
                    height_categories["medium"].append((comp.reference, height))
                elif height < 10.0:
                    height_categories["high"].append((comp.reference, height))
                else:
                    height_categories["very_high"].append((comp.reference, height))
            
            return {
                "success": True,
                "pcb_file": validated_path,
                "height_statistics": height_analysis,
                "height_categories": {
                    category: [{"component": ref, "height_mm": height} 
                             for ref, height in components]
                    for category, components in height_categories.items()
                },
                "tallest_components": sorted(
                    [(comp.reference, analyzer._estimate_component_height(comp)) 
                     for comp in components],
                    key=lambda x: x[1], reverse=True
                )[:10],  # Top 10 tallest components
                "enclosure_requirements": {
                    "minimum_height_mm": height_analysis["max"] + 2.0,  # Add 2mm clearance
                    "recommended_height_mm": height_analysis["max"] + 5.0  # Add 5mm clearance
                }
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "pcb_file": pcb_file_path
            }
    
    @mcp.tool()
    def check_assembly_feasibility(pcb_file_path: str) -> Dict[str, Any]:
        """
        Analyze PCB assembly feasibility and identify potential issues.
        
        Checks for component accessibility, assembly sequence issues,
        and manufacturing constraints that could affect PCB assembly.
        
        Args:
            pcb_file_path: Path to the .kicad_pcb file
            
        Returns:
            Assembly feasibility analysis with issues and recommendations
        """
        try:
            validated_path = validate_kicad_file(pcb_file_path, "pcb")
            
            analyzer = Model3DAnalyzer(validated_path)
            mechanical_analysis = analyzer.perform_mechanical_analysis()
            components = mechanical_analysis.components
            
            assembly_issues = []
            assembly_warnings = []
            
            # Check for components too close to board edge
            for comp in components:
                edge_distance = analyzer._distance_to_board_edge(
                    comp, mechanical_analysis.board_dimensions
                )
                if edge_distance < 1.0:  # Less than 1mm from edge
                    assembly_warnings.append({
                        "component": comp.reference,
                        "issue": f"Component only {edge_distance:.2f}mm from board edge",
                        "recommendation": "Consider moving component away from edge for easier assembly"
                    })
            
            # Check for very small components that might be hard to place
            small_component_footprints = ["0201", "0402"]
            for comp in components:
                if any(size in (comp.footprint or "") for size in small_component_footprints):
                    assembly_warnings.append({
                        "component": comp.reference,
                        "issue": f"Very small footprint {comp.footprint}",
                        "recommendation": "Verify pick-and-place machine compatibility"
                    })
            
            # Check component density
            board_area = (mechanical_analysis.board_dimensions.width * 
                         mechanical_analysis.board_dimensions.height)
            component_density = len(components) / (board_area / 100)  # Components per cm²
            
            if component_density > 5.0:
                assembly_warnings.append({
                    "component": "Board",
                    "issue": f"High component density: {component_density:.1f} components/cm²",
                    "recommendation": "Consider larger board or fewer components for easier assembly"
                })
            
            return {
                "success": True,
                "pcb_file": validated_path,
                "assembly_feasible": len(assembly_issues) == 0,
                "assembly_issues": assembly_issues,
                "assembly_warnings": assembly_warnings,
                "component_density": component_density,
                "board_utilization": {
                    "component_count": len(components),
                    "board_area_mm2": board_area,
                    "density_per_cm2": component_density
                },
                "recommendations": [
                    "Review component placement for optimal assembly sequence",
                    "Ensure adequate fiducial markers for automated assembly",
                    "Consider component orientation for consistent placement direction"
                ] if assembly_warnings else ["PCB appears suitable for standard assembly processes"]
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "pcb_file": pcb_file_path
            }