"""SAX circuit simulation tool handler."""

from __future__ import annotations

import json
import logging
import textwrap
from typing import TYPE_CHECKING, Any

from mcp.types import ImageContent, TextContent, Tool

from ..config import MCPConfig
from ..notify import notify_open_simulation
from ..plot import HAS_MATPLOTLIB, render_simulation_plot
from .base import EndpointMapping, ToolHandler, add_project_param

if TYPE_CHECKING:
    from ..client import FastAPIClient

logger = logging.getLogger(__name__)

__all__ = ["SimulateComponentHandler"]


class SimulateComponentHandler(ToolHandler):
    """Handler for running SAX circuit simulations."""

    @property
    def name(self) -> str:
        return "simulate_component"

    @property
    def definition(self) -> Tool:
        return Tool(
            name="simulate_component",
            description=(
                "Run a SAX circuit simulation on a photonic component by name. "
                "This simulates the optical behavior of the component using its "
                "SAX model. Returns S-parameters showing how light propagates "
                "through the component's ports. Use this to analyze component "
                "performance before fabrication. Supports customizing both layout "
                "parameters (e.g., {'length_mmi': 12, 'gap_mmi': 0.3}) and SAX "
                "model simulation settings (e.g., {'wl': [1.5, 1.55, 1.6], 'loss': 0.2})."
            ),
            inputSchema=add_project_param(
                {
                    "type": "object",
                    "properties": {
                        "name": {
                            "type": "string",
                            "description": (
                                "Name of the component/cell to simulate. The component "
                                "must have a SAX model defined."
                            ),
                        },
                        "layout": {
                            "type": "object",
                            "description": (
                                "Optional component-specific geometric parameters. "
                                "Use this to customize the component's physical dimensions "
                                "before simulation. Example: {'length_mmi': 12, 'gap_mmi': 0.3}. "
                                "Default is empty (use default layout parameters)."
                            ),
                        },
                        "model": {
                            "type": "object",
                            "description": (
                                "Optional SAX simulation settings. Use this to configure "
                                "simulation parameters like wavelength sweeps or loss values. "
                                "Example: {'wl': [1.5, 1.55, 1.6], 'loss': 0.2}. "
                                "Default is empty (use default model parameters)."
                            ),
                        },
                        "how": {
                            "type": "string",
                            "enum": ["from_layout", "from_netlist"],
                            "description": (
                                "Simulation method. 'from_layout' simulates from the physical "
                                "layout geometry (default). 'from_netlist' simulates from the "
                                "hierarchical netlist representation."
                            ),
                        },
                        "visualize": {
                            "type": "boolean",
                            "description": (
                                "If true, render S-parameter results as a spectral plot "
                                "(|S|² in dB vs wavelength) and return it as a PNG image "
                                "alongside the text summary. Requires matplotlib."
                            ),
                        },
                    },
                    "required": ["name"],
                }
            ),
        )

    @property
    def mapping(self) -> EndpointMapping:
        return EndpointMapping(method="POST", path="/api/sax/{name}/simulate")

    def transform_request(self, args: dict[str, Any]) -> dict[str, Any]:
        """Transform simulate_component MCP args to FastAPI request body."""
        name = args["name"]
        layout = args.get("layout", {})
        model = args.get("model", {})
        how = args.get("how", "from_layout")

        return {
            "path": f"/api/sax/{name}/simulate",
            "json_data": {
                "layout": layout,
                "model": model,
                "from_netlist": how == "from_netlist",
            },
        }

    def transform_response(self, response: Any) -> dict[str, Any]:
        """Transform simulate_component response for LLM consumption.

        Returns only the file path and token estimate instead of the full
        S-parameter data, which can be extremely large (500k+ tokens).
        """
        if not isinstance(response, dict):
            return {"error": f"Unexpected response type: {type(response).__name__}"}

        if "detail" in response and "file_path" not in response:
            return {"error": response["detail"]}

        json_str = json.dumps(response)
        estimated_tokens = len(json_str) // 4

        file_path = response.get("file_path")

        return {
            "file_path": file_path,
            "estimated_tokens": estimated_tokens,
            "message": f"Simulation complete. Results saved to: {file_path}",
            "note": textwrap.dedent("""If needed you can further analyse the simulation results
            by loading the S-parameter data from the file path provided.
            The structure of the json file is as follows: result: {"o1": {"o3": {"real": [],
            "imag": []},"o4": {"real": [],"imag": []}}, ...}, "wavelengths": [],"""),
            "simulation_id": response.get("simulation_id"),
        }

    async def handle(
        self,
        arguments: dict[str, Any],
        client: FastAPIClient,
    ) -> list[TextContent | ImageContent]:
        """Run simulation, optionally render S-parameter plot, and notify VS Code extension."""
        try:
            project = arguments.get("project")
            transformed = self.transform_request(arguments)

            response = await client.request(
                method=self.mapping.method,
                path=transformed["path"],
                json_data=transformed.get("json_data"),
                project=project,
            )

            result_data = self.transform_response(response)
            logger.debug("Simulation result: %s", result_data)

            results: list[TextContent | ImageContent] = [
                TextContent(
                    type="text",
                    text=json.dumps(result_data, indent=2),
                )
            ]

            visualize = arguments.get("visualize", False)
            if visualize and HAS_MATPLOTLIB and "error" not in result_data:
                image = render_simulation_plot(response, arguments["name"])
                if image:
                    results.append(image)
                    logger.info("Added S-parameter plot for: %s", arguments["name"])

            # Notify extension to open simulation panel
            if MCPConfig.UI_NOTIFICATIONS and "error" not in result_data:
                try:
                    await notify_open_simulation(
                        arguments.get("project"),
                        arguments["name"],
                        result_data.get("simulation_id"),
                    )
                except Exception:
                    logger.debug(
                        "Failed to send simulation notification", exc_info=True
                    )

            return results

        except Exception as e:
            error_msg = f"Tool execution failed: {e!s}"
            logger.exception(error_msg)
            return [
                TextContent(
                    type="text",
                    text=json.dumps({"error": error_msg}),
                )
            ]
