kicad-mcp/kicad_mcp/utils/freerouting_engine.py

"""
FreeRouting Integration Engine

Provides automated PCB routing capabilities using the FreeRouting autorouter.
This module handles DSN file generation from KiCad boards, FreeRouting execution,
and importing the routed results back into KiCad via the IPC API.

FreeRouting: https://www.freerouting.app/
GitHub: https://github.com/freerouting/freerouting
"""

import json
import logging
import os
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import requests
from kipy.board_types import BoardLayer

from kicad_mcp.utils.ipc_client import KiCadIPCClient, kicad_ipc_session

logger = logging.getLogger(__name__)


class FreeRoutingError(Exception):
    """Custom exception for FreeRouting operations."""
    pass


class FreeRoutingEngine:
    """
    Engine for automated PCB routing using FreeRouting.
    
    Handles the complete workflow:
    1. Export DSN file from KiCad board
    2. Process with FreeRouting autorouter
    3. Import routed SES file back to KiCad
    4. Optimize and validate routing results
    """
    
    def __init__(
        self,
        freerouting_jar_path: Optional[str] = None,
        java_executable: str = "java",
        working_directory: Optional[str] = None
    ):
        """
        Initialize FreeRouting engine.
        
        Args:
            freerouting_jar_path: Path to FreeRouting JAR file
            java_executable: Java executable command
            working_directory: Working directory for temporary files
        """
        self.freerouting_jar_path = freerouting_jar_path
        self.java_executable = java_executable
        self.working_directory = working_directory or tempfile.gettempdir()
        
        # Default routing parameters
        self.routing_config = {
            "via_costs": 50,
            "plane_via_costs": 5,
            "start_ripup_costs": 100,
            "automatic_layer_dimming": True,
            "ignore_conduction": False,
            "automatic_neckdown": True,
            "postroute_optimization": True,
            "max_iterations": 1000,
            "improvement_threshold": 0.01
        }
        
        # Layer configuration
        self.layer_config = {
            "signal_layers": [BoardLayer.BL_F_Cu, BoardLayer.BL_B_Cu],
            "power_layers": [],
            "preferred_direction": {
                BoardLayer.BL_F_Cu: "horizontal",
                BoardLayer.BL_B_Cu: "vertical"
            }
        }
    
    def find_freerouting_jar(self) -> Optional[str]:
        """
        Attempt to find FreeRouting JAR file in common locations.
        
        Returns:
            Path to FreeRouting JAR if found, None otherwise
        """
        common_paths = [
            "freerouting.jar",
            "freerouting-1.9.0.jar",
            "/usr/local/bin/freerouting.jar",
            "/opt/freerouting/freerouting.jar",
            os.path.expanduser("~/freerouting.jar"),
            os.path.expanduser("~/bin/freerouting.jar"),
            os.path.expanduser("~/Downloads/freerouting.jar")
        ]
        
        for path in common_paths:
            if os.path.isfile(path):
                logger.info(f"Found FreeRouting JAR at: {path}")
                return path
        
        return None
    
    def check_freerouting_availability(self) -> Dict[str, Any]:
        """
        Check if FreeRouting is available and working.
        
        Returns:
            Dictionary with availability status
        """
        if not self.freerouting_jar_path:
            self.freerouting_jar_path = self.find_freerouting_jar()
        
        if not self.freerouting_jar_path:
            return {
                "available": False,
                "message": "FreeRouting JAR file not found",
                "jar_path": None
            }
        
        if not os.path.isfile(self.freerouting_jar_path):
            return {
                "available": False,
                "message": f"FreeRouting JAR file not found at: {self.freerouting_jar_path}",
                "jar_path": self.freerouting_jar_path
            }
        
        # Test Java and FreeRouting
        try:
            result = subprocess.run(
                [self.java_executable, "-jar", self.freerouting_jar_path, "-help"],
                capture_output=True,
                text=True,
                timeout=30
            )
            
            if result.returncode == 0 or "freerouting" in result.stdout.lower():
                return {
                    "available": True,
                    "message": "FreeRouting is available and working",
                    "jar_path": self.freerouting_jar_path,
                    "java_executable": self.java_executable
                }
            else:
                return {
                    "available": False,
                    "message": f"FreeRouting test failed: {result.stderr}",
                    "jar_path": self.freerouting_jar_path
                }
        
        except subprocess.TimeoutExpired:
            return {
                "available": False,
                "message": "FreeRouting test timed out",
                "jar_path": self.freerouting_jar_path
            }
        except Exception as e:
            return {
                "available": False,
                "message": f"Error testing FreeRouting: {e}",
                "jar_path": self.freerouting_jar_path
            }
    
    def export_dsn_from_kicad(
        self,
        board_path: str,
        dsn_output_path: str,
        routing_options: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        Export DSN file from KiCad board using KiCad CLI.
        
        Args:
            board_path: Path to .kicad_pcb file
            dsn_output_path: Output path for DSN file
            routing_options: Optional routing configuration
            
        Returns:
            True if export successful
        """
        try:
            # Use KiCad CLI to export DSN
            cmd = [
                "kicad-cli", "pcb", "export", "specctra-dsn",
                "--output", dsn_output_path,
                board_path
            ]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0 and os.path.isfile(dsn_output_path):
                logger.info(f"DSN exported successfully to: {dsn_output_path}")
                
                # Post-process DSN file with routing options if provided
                if routing_options:
                    self._customize_dsn_file(dsn_output_path, routing_options)
                
                return True
            else:
                logger.error(f"DSN export failed: {result.stderr}")
                return False
        
        except subprocess.TimeoutExpired:
            logger.error("DSN export timed out")
            return False
        except Exception as e:
            logger.error(f"Error exporting DSN: {e}")
            return False
    
    def _customize_dsn_file(self, dsn_path: str, options: Dict[str, Any]):
        """
        Customize DSN file with specific routing options.
        
        Args:
            dsn_path: Path to DSN file
            options: Routing configuration options
        """
        try:
            with open(dsn_path, 'r') as f:
                content = f.read()
            
            # Add routing directives to DSN file
            # This is a simplified implementation - real DSN modification would be more complex
            modifications = []
            
            if "via_costs" in options:
                modifications.append(f"(via_costs {options['via_costs']})")
            
            if "max_iterations" in options:
                modifications.append(f"(max_iterations {options['max_iterations']})")
            
            # Insert modifications before the closing parenthesis
            if modifications:
                insertion_point = content.rfind(')')
                if insertion_point != -1:
                    modified_content = (
                        content[:insertion_point] +
                        '\n'.join(modifications) + '\n' +
                        content[insertion_point:]
                    )
                    
                    with open(dsn_path, 'w') as f:
                        f.write(modified_content)
                    
                    logger.info(f"DSN file customized with {len(modifications)} options")
        
        except Exception as e:
            logger.warning(f"Error customizing DSN file: {e}")
    
    def run_freerouting(
        self,
        dsn_path: str,
        output_directory: str,
        routing_config: Optional[Dict[str, Any]] = None
    ) -> Tuple[bool, Optional[str]]:
        """
        Run FreeRouting autorouter on DSN file.
        
        Args:
            dsn_path: Path to input DSN file
            output_directory: Directory for output files
            routing_config: Optional routing configuration
            
        Returns:
            Tuple of (success, output_ses_path)
        """
        if not self.freerouting_jar_path:
            raise FreeRoutingError("FreeRouting JAR path not configured")
        
        config = {**self.routing_config, **(routing_config or {})}
        
        try:
            # Prepare FreeRouting command
            cmd = [
                self.java_executable,
                "-jar", self.freerouting_jar_path,
                "-de", dsn_path,  # Input DSN file
                "-do", output_directory,  # Output directory
            ]
            
            # Add routing parameters
            if config.get("automatic_layer_dimming", True):
                cmd.extend(["-ld", "true"])
            
            if config.get("automatic_neckdown", True):
                cmd.extend(["-nd", "true"])
            
            if config.get("postroute_optimization", True):
                cmd.extend(["-opt", "true"])
            
            logger.info(f"Running FreeRouting: {' '.join(cmd)}")
            
            # Run FreeRouting
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=300,  # 5 minute timeout
                cwd=output_directory
            )
            
            if result.returncode == 0:
                # Find output SES file
                ses_files = list(Path(output_directory).glob("*.ses"))
                if ses_files:
                    ses_path = str(ses_files[0])
                    logger.info(f"FreeRouting completed successfully: {ses_path}")
                    return True, ses_path
                else:
                    logger.error("FreeRouting completed but no SES file found")
                    return False, None
            else:
                logger.error(f"FreeRouting failed: {result.stderr}")
                return False, None
        
        except subprocess.TimeoutExpired:
            logger.error("FreeRouting timed out")
            return False, None
        except Exception as e:
            logger.error(f"Error running FreeRouting: {e}")
            return False, None
    
    def import_ses_to_kicad(
        self,
        board_path: str,
        ses_path: str,
        backup_original: bool = True
    ) -> bool:
        """
        Import SES routing results back into KiCad board.
        
        Args:
            board_path: Path to .kicad_pcb file
            ses_path: Path to SES file with routing results
            backup_original: Whether to backup original board file
            
        Returns:
            True if import successful
        """
        try:
            # Backup original board if requested
            if backup_original:
                backup_path = f"{board_path}.backup.{int(time.time())}"
                import shutil
                shutil.copy2(board_path, backup_path)
                logger.info(f"Original board backed up to: {backup_path}")
            
            # Use KiCad CLI to import SES file
            cmd = [
                "kicad-cli", "pcb", "import", "specctra-ses",
                "--output", board_path,
                ses_path
            ]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0:
                logger.info(f"SES imported successfully to: {board_path}")
                return True
            else:
                logger.error(f"SES import failed: {result.stderr}")
                return False
        
        except subprocess.TimeoutExpired:
            logger.error("SES import timed out")
            return False
        except Exception as e:
            logger.error(f"Error importing SES: {e}")
            return False
    
    def route_board_complete(
        self,
        board_path: str,
        routing_config: Optional[Dict[str, Any]] = None,
        preserve_existing: bool = False
    ) -> Dict[str, Any]:
        """
        Complete automated routing workflow for a KiCad board.
        
        Args:
            board_path: Path to .kicad_pcb file
            routing_config: Optional routing configuration
            preserve_existing: Whether to preserve existing routing
            
        Returns:
            Dictionary with routing results and statistics
        """
        config = {**self.routing_config, **(routing_config or {})}
        
        # Create temporary directory for routing files
        with tempfile.TemporaryDirectory(prefix="freerouting_") as temp_dir:
            try:
                # Prepare file paths
                dsn_path = os.path.join(temp_dir, "board.dsn")
                
                # Step 1: Export DSN from KiCad
                logger.info("Step 1: Exporting DSN file from KiCad")
                if not self.export_dsn_from_kicad(board_path, dsn_path, config):
                    return {
                        "success": False,
                        "error": "Failed to export DSN file from KiCad",
                        "step": "dsn_export"
                    }
                
                # Step 2: Get pre-routing statistics
                pre_stats = self._analyze_board_connectivity(board_path)
                
                # Step 3: Run FreeRouting
                logger.info("Step 2: Running FreeRouting autorouter")
                success, ses_path = self.run_freerouting(dsn_path, temp_dir, config)
                if not success or not ses_path:
                    return {
                        "success": False,
                        "error": "FreeRouting execution failed",
                        "step": "freerouting",
                        "pre_routing_stats": pre_stats
                    }
                
                # Step 4: Import results back to KiCad
                logger.info("Step 3: Importing routing results back to KiCad")
                if not self.import_ses_to_kicad(board_path, ses_path):
                    return {
                        "success": False,
                        "error": "Failed to import SES file to KiCad",
                        "step": "ses_import",
                        "pre_routing_stats": pre_stats
                    }
                
                # Step 5: Get post-routing statistics
                post_stats = self._analyze_board_connectivity(board_path)
                
                # Step 6: Generate routing report
                routing_report = self._generate_routing_report(pre_stats, post_stats, config)
                
                return {
                    "success": True,
                    "message": "Automated routing completed successfully",
                    "pre_routing_stats": pre_stats,
                    "post_routing_stats": post_stats,
                    "routing_report": routing_report,
                    "config_used": config
                }
            
            except Exception as e:
                logger.error(f"Error during automated routing: {e}")
                return {
                    "success": False,
                    "error": str(e),
                    "step": "general_error"
                }
    
    def _analyze_board_connectivity(self, board_path: str) -> Dict[str, Any]:
        """
        Analyze board connectivity status.
        
        Args:
            board_path: Path to board file
            
        Returns:
            Connectivity statistics
        """
        try:
            with kicad_ipc_session(board_path=board_path) as client:
                return client.check_connectivity()
        except Exception as e:
            logger.warning(f"Could not analyze connectivity via IPC: {e}")
            return {"error": str(e)}
    
    def _generate_routing_report(
        self,
        pre_stats: Dict[str, Any],
        post_stats: Dict[str, Any],
        config: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Generate routing completion report.
        
        Args:
            pre_stats: Pre-routing statistics
            post_stats: Post-routing statistics
            config: Routing configuration used
            
        Returns:
            Routing report
        """
        report = {
            "routing_improvement": {},
            "completion_metrics": {},
            "recommendations": []
        }
        
        if "routing_completion" in pre_stats and "routing_completion" in post_stats:
            pre_completion = pre_stats["routing_completion"]
            post_completion = post_stats["routing_completion"]
            improvement = post_completion - pre_completion
            
            report["routing_improvement"] = {
                "pre_completion_percent": pre_completion,
                "post_completion_percent": post_completion,
                "improvement_percent": improvement
            }
        
        if "unrouted_nets" in post_stats:
            unrouted = post_stats["unrouted_nets"]
            if unrouted > 0:
                report["recommendations"].append(
                    f"Manual routing may be needed for {unrouted} remaining unrouted nets"
                )
            else:
                report["recommendations"].append("All nets successfully routed!")
        
        if "total_nets" in post_stats:
            total = post_stats["total_nets"]
            routed = post_stats.get("routed_nets", 0)
            
            report["completion_metrics"] = {
                "total_nets": total,
                "routed_nets": routed,
                "routing_success_rate": round(routed / max(total, 1) * 100, 1)
            }
        
        return report
    
    def optimize_routing_parameters(
        self,
        board_path: str,
        target_completion: float = 95.0
    ) -> Dict[str, Any]:
        """
        Optimize routing parameters for best results on a specific board.
        
        Args:
            board_path: Path to board file
            target_completion: Target routing completion percentage
            
        Returns:
            Optimized parameters and results
        """
        parameter_sets = [
            # Conservative approach
            {
                "via_costs": 30,
                "start_ripup_costs": 50,
                "max_iterations": 500,
                "approach": "conservative"
            },
            # Balanced approach
            {
                "via_costs": 50,
                "start_ripup_costs": 100,
                "max_iterations": 1000,
                "approach": "balanced"
            },
            # Aggressive approach
            {
                "via_costs": 80,
                "start_ripup_costs": 200,
                "max_iterations": 2000,
                "approach": "aggressive"
            }
        ]
        
        best_result = None
        best_completion = 0
        
        for i, params in enumerate(parameter_sets):
            logger.info(f"Testing parameter set {i+1}/3: {params['approach']}")
            
            # Create backup before testing
            backup_path = f"{board_path}.param_test_{i}"
            import shutil
            shutil.copy2(board_path, backup_path)
            
            try:
                result = self.route_board_complete(board_path, params)
                
                if result["success"]:
                    completion = result["post_routing_stats"].get("routing_completion", 0)
                    
                    if completion > best_completion:
                        best_completion = completion
                        best_result = {
                            "parameters": params,
                            "result": result,
                            "completion": completion
                        }
                    
                    if completion >= target_completion:
                        logger.info(f"Target completion {target_completion}% achieved!")
                        break
                
                # Restore backup for next test
                shutil.copy2(backup_path, board_path)
                
            except Exception as e:
                logger.error(f"Error testing parameter set {i+1}: {e}")
                # Restore backup
                shutil.copy2(backup_path, board_path)
            
            finally:
                # Clean up backup
                if os.path.exists(backup_path):
                    os.remove(backup_path)
        
        if best_result:
            # Apply best parameters one final time
            final_result = self.route_board_complete(board_path, best_result["parameters"])
            
            return {
                "success": True,
                "best_parameters": best_result["parameters"],
                "best_completion": best_completion,
                "final_result": final_result,
                "optimization_summary": f"Best approach: {best_result['parameters']['approach']} "
                                      f"(completion: {best_completion:.1f}%)"
            }
        else:
            return {
                "success": False,
                "error": "No successful routing configuration found",
                "tested_parameters": parameter_sets
            }


def check_routing_prerequisites() -> Dict[str, Any]:
    """
    Check if all prerequisites for automated routing are available.
    
    Returns:
        Dictionary with prerequisite status
    """
    status = {
        "overall_ready": False,
        "components": {}
    }
    
    # Check KiCad IPC API
    try:
        from kicad_mcp.utils.ipc_client import check_kicad_availability
        kicad_status = check_kicad_availability()
        status["components"]["kicad_ipc"] = kicad_status
    except Exception as e:
        status["components"]["kicad_ipc"] = {
            "available": False,
            "error": str(e)
        }
    
    # Check FreeRouting
    engine = FreeRoutingEngine()
    freerouting_status = engine.check_freerouting_availability()
    status["components"]["freerouting"] = freerouting_status
    
    # Check KiCad CLI
    try:
        result = subprocess.run(
            ["kicad-cli", "--version"],
            capture_output=True,
            text=True,
            timeout=10
        )
        status["components"]["kicad_cli"] = {
            "available": result.returncode == 0,
            "version": result.stdout.strip() if result.returncode == 0 else None,
            "error": result.stderr if result.returncode != 0 else None
        }
    except Exception as e:
        status["components"]["kicad_cli"] = {
            "available": False,
            "error": str(e)
        }
    
    # Determine overall readiness
    all_components_ready = all(
        comp.get("available", False) for comp in status["components"].values()
    )
    
    status["overall_ready"] = all_components_ready
    status["message"] = (
        "All routing prerequisites are available" if all_components_ready
        else "Some routing prerequisites are missing or not working"
    )
    
    return status