from mcp.server.fastmcp import FastMCP


def register_solver_run_layer(app: FastMCP):
    import os

    from pathlib import Path

    from am.mcp.types import ToolSuccess, ToolError
    from am.mcp.utils import tool_success, tool_error

    @app.tool(
        title="Run Layer with Solver",
        description="Runs solver on a segments file (segments file should be one layer) and saves the generated meshes.",
        structured_output=True,
    )
    def solver_run_layer(
        workspace: str,
        segments_foldername: str,
        layer_number: int,
        build_parameters_filename: str = "default.json",
        material_filename: str = "default.json",
        mesh_config_filename: str = "default.json",
        model_name: str = "eagar-tsai",
        run_name: str | None = None,
    ) -> ToolSuccess[Path] | ToolError:
        """
        Runs solver for segments at a specified layer number.
        Args:
            workspace: Folder name of existing workspace
            segments_foldername: Folder name of where segments are expected to be found.
            layer_number: Layer number to run solver on, typically starts from 1. For testing out, try skipping the first several layers as those sometimes don't include part geometry.
            distance_xy_max: Maximum segment length when parsing (defaults to 1.0 mm).
            build_parameters_filename: build parameters file to use with solver.
            material_filename: material config file to use with solver.
            mesh_config_filename: mesh config file to use with solver.
            run_name: Name of folder to save generated meshes at, typically autogenerated.
        """
        import json

        from am.config import BuildParameters, Material, MeshParameters
        from am.simulator.models import SolverSegment
        from am.solver.layer import SolverLayer
        from wa.cli.utils import get_workspace_path

        try:
            workspace_path = get_workspace_path(workspace)

            solver = SolverLayer()

            # Segments
            segments_path = workspace_path / "segments" / segments_foldername / "layers"

            # Uses number of files in segments path as total layers for zfill.
            total_layers = len(os.listdir(segments_path))
            z_fill = len(f"{total_layers}")
            layer_number_string = f"{layer_number}".zfill(z_fill)
            segments_file_path = segments_path / f"{layer_number_string}.json"

            # TODO: Settle on a better way to handle loading of lists of a particular schema.
            with open(segments_file_path, "r") as f:
                segments_data = json.load(f)
            segments = [SolverSegment(**seg_data) for seg_data in segments_data]
            # segments = Segment.load(segments_file_path)

            # Configs
            build_parameters = BuildParameters.load(
                workspace_path
                / "configs"
                / "build_parameters"
                / build_parameters_filename
            )
            material = Material.load(
                workspace_path / "configs" / "materials" / material_filename
            )
            mesh_parameters = MeshParameters.load(
                workspace_path / "configs" / "mesh_parameters" / mesh_config_filename
            )

            run_out_path = solver.run(
                segments,
                build_parameters,
                material,
                mesh_parameters,
                workspace_path,
                model_name,
                run_name,
            )

            return tool_success(run_out_path)

        except PermissionError as e:
            return tool_error(
                "Permission denied when running solver",
                "PERMISSION_DENIED",
                workspace_name=workspace,
                exception_type=type(e).__name__,
            )

        except Exception as e:
            return tool_error(
                "Failed to run solver",
                "SOLVER_RUN_FAILED",
                workspace_name=workspace,
                exception_type=type(e).__name__,
                exception_message=str(e),
            )

        _ = solver_run_layer
