mcp-arduino/src/mcp_arduino_server/components/wireviz_manager.py

"""WireViz circuit diagram generation component"""
import base64
import datetime
import logging
import os
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, Optional, Any

from fastmcp.utilities.types import Image
from pydantic import BaseModel, Field

log = logging.getLogger(__name__)


class WireVizRequest(BaseModel):
    """Request model for WireViz operations"""
    yaml_content: Optional[str] = Field(None, description="WireViz YAML content")
    description: Optional[str] = Field(None, description="Natural language circuit description")
    sketch_name: str = Field("circuit", description="Name for output files")
    output_base: str = Field("circuit", description="Base name for output files")


class WireVizManager:
    """Manages WireViz circuit diagram generation"""

    def __init__(self, config, mcp_context=None):
        self.config = config
        self.wireviz_path = config.wireviz_path
        self.mcp_context = mcp_context  # For accessing sampling

    def get_instructions(self) -> str:
        """Get WireViz usage instructions"""
        return """
# WireViz Circuit Diagram Instructions

WireViz is a tool for generating circuit wiring diagrams from YAML descriptions.

## Basic YAML Structure:

```yaml
connectors:
  Arduino:
    type: Arduino Uno
    pins: [GND, 5V, D2, D3, A0]

  LED:
    type: LED
    pins: [cathode, anode]

cables:
  power:
    colors: [BK, RD]  # Black, Red
    gauge: 22 AWG

connections:
  - Arduino: [GND]
    cable: [1]
    LED: [cathode]
  - Arduino: [D2]
    cable: [2]
    LED: [anode]
```

## Color Codes:
- BK: Black, RD: Red, BL: Blue, GN: Green, YE: Yellow
- OR: Orange, VT: Violet, GY: Gray, WH: White, BN: Brown

## Tips:
1. Define all connectors first
2. Specify cable properties (colors, gauge)
3. Map connections clearly
4. Use descriptive names

For AI-powered generation from descriptions, use the
`generate_circuit_diagram_from_description` tool.
"""

    async def generate_from_yaml(self, yaml_content: str, output_base: str = "circuit") -> Dict[str, Any]:
        """Generate circuit diagram from WireViz YAML"""
        try:
            # Create timestamped output directory
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            output_dir = self.config.sketches_base_dir / f"wireviz_{timestamp}"
            output_dir.mkdir(parents=True, exist_ok=True)

            # Write YAML to temporary file
            yaml_path = output_dir / f"{output_base}.yaml"
            yaml_path.write_text(yaml_content)

            # Run WireViz
            cmd = [self.wireviz_path, str(yaml_path), "-o", str(output_dir)]
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=self.config.command_timeout
            )

            if result.returncode != 0:
                error_msg = f"WireViz failed: {result.stderr}"
                log.error(error_msg)
                return {"error": error_msg}

            # Find generated PNG
            png_files = list(output_dir.glob("*.png"))
            if not png_files:
                return {"error": "No PNG file generated"}

            png_path = png_files[0]

            # Read and encode image
            with open(png_path, "rb") as f:
                image_data = f.read()
                encoded_image = base64.b64encode(image_data).decode("utf-8")

            # Open image in default viewer
            self._open_file(png_path)

            return {
                "success": True,
                "message": f"Circuit diagram generated: {png_path}",
                "image": Image(data=encoded_image, format="png"),
                "paths": {
                    "yaml": str(yaml_path),
                    "png": str(png_path),
                    "directory": str(output_dir)
                }
            }

        except subprocess.TimeoutExpired:
            return {"error": f"WireViz timed out after {self.config.command_timeout} seconds"}
        except Exception as e:
            log.exception("WireViz generation failed")
            return {"error": str(e)}

    async def generate_from_description(
        self,
        description: str,
        sketch_name: str = "",
        output_base: str = "circuit"
    ) -> Dict[str, Any]:
        """Generate circuit diagram from natural language description using client's LLM"""

        if not self.mcp_context:
            return {
                "error": "MCP context not available. Client sampling is required for AI generation.",
                "hint": "The MCP client must support sampling for this feature to work."
            }

        try:
            # Create prompt for YAML generation
            prompt = self._create_wireviz_prompt(description, sketch_name)

            # Use FastMCP sampling to request completion from the client
            from mcp.types import SamplingMessage

            messages = [
                SamplingMessage(
                    role="system",
                    content="You are an expert at creating WireViz YAML circuit diagrams. Return ONLY the YAML content, no explanations or markdown."
                ),
                SamplingMessage(
                    role="user",
                    content=prompt
                )
            ]

            # Request completion from the client
            result = await self.mcp_context.sample(
                messages=messages,
                max_tokens=2000,
                temperature=0.3,
                stop_sequences=["```"]
            )

            if not result or not result.content:
                return {"error": "No response from client LLM"}

            yaml_content = result.content

            # Clean up the YAML (remove markdown if present)
            yaml_content = self._clean_yaml_content(yaml_content)

            # Generate diagram from YAML
            diagram_result = await self.generate_from_yaml(yaml_content, output_base)

            if "error" not in diagram_result:
                diagram_result["yaml_generated"] = yaml_content
                diagram_result["generated_by"] = "client_llm"

            return diagram_result

        except ImportError:
            return {
                "error": "Client sampling not available. Your MCP client may not support this feature.",
                "fallback": "You can still create diagrams by writing WireViz YAML manually."
            }
        except Exception as e:
            log.exception("Client-based WireViz generation failed")
            return {"error": f"Generation failed: {str(e)}"}

    def _create_wireviz_prompt(self, description: str, sketch_name: str) -> str:
        """Create prompt for AI to generate WireViz YAML"""
        base_prompt = """Generate a WireViz YAML circuit diagram for the following description:

Description: {description}

Requirements:
1. Use proper WireViz YAML syntax
2. Include all necessary connectors, cables, and connections
3. Use appropriate wire colors and gauges
4. Add descriptive labels
5. Follow electrical safety standards

Return ONLY the YAML content, no explanations."""

        if sketch_name:
            base_prompt += f"\n\nThis is for an Arduino sketch named: {sketch_name}"

        return base_prompt.format(description=description)

    def _clean_yaml_content(self, content: str) -> str:
        """Remove markdown code blocks if present"""
        lines = content.strip().split('\n')

        # Remove markdown code fence if present
        if lines[0].startswith('```'):
            lines = lines[1:]
        if lines[-1].startswith('```'):
            lines = lines[:-1]

        return '\n'.join(lines)

    def _open_file(self, file_path: Path) -> None:
        """Open file in default system application"""
        try:
            if os.name == 'posix':  # macOS and Linux
                subprocess.run(['open' if os.uname().sysname == 'Darwin' else 'xdg-open', str(file_path)])
            elif os.name == 'nt':  # Windows
                os.startfile(str(file_path))
        except Exception as e:
            log.warning(f"Could not open file automatically: {e}")