import os

from mcp.server import FastMCP
from pathlib import Path

from typing import Union


def register_solver(app: FastMCP):
    from am.mcp.types import ToolSuccess, ToolError
    from am.mcp.utils import tool_success, tool_error
    from am.solver.__main__ import SolverOutputFolder

    @app.tool(
        title="Initialize Solver",
        description="Creates new solver to generate meshes",
        structured_output=True,
    )
    def solver_initialize(
        workspace: str,
    ) -> Union[ToolSuccess[None], ToolError]:
        """Initializes solver with defaults inside workspace folder."""
        from am.solver import Solver
        from am.cli.utils import get_workspace_path

        try:
            workspace_path = get_workspace_path(workspace)

            solver = Solver()
            solver.create_solver_config(solver_path=workspace_path / "solver")
            solver.create_default_configs()

            return tool_success(None)

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

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

    # @app.tool(
    #     title="Measure Melt Pool Dimensions",
    #     description="Generate the expected melt pool dimensional measurements for a given build configuration and material configuration",
    #     structured_output=True,
    # )
    # def solver_measure_melt_pool_dimensions(
    #     workspace: str,
    #     build_parameters_filename: str = "default.json",
    #     material_filename: str = "default.json",
    #     run_name: str | None = None,
    # ) -> Union[ToolSuccess[Path], ToolError]:
    #     """
    #     Runs solver for segments at a specified layer number.
    #     Args:
    #         workspace: Folder name of existing workspace
    #         distance_xy_max: Maximum segment length when parsing (defaults to 1.0 mm).
    #         build_parameters_filename: build config 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.
    #     """
    #     from am.cli.utils import get_workspace_path
    #     from am.schema import BuildParameters, Material
    #     from am.solver import Solver
    #
    #     try:
    #         workspace_path = get_workspace_path(workspace)
    #
    #         solver = Solver()
    #
    #         # Configs
    #         solver_configs_path = workspace_path / "solver" / "config"
    #         build_parameters = BuildParameters.load(
    #             solver_configs_path / "build_parameters" / build_parameters_filename
    #         )
    #         material = Material.load(
    #             solver_configs_path / "materials" / material_filename
    #         )
    #
    #         out_path = solver.measure_melt_pool_dimensions(build_parameters, material, workspace_path, run_name)
    #
    #         return tool_success(out_path)
    #
    #     except PermissionError as e:
    #         return tool_error(
    #             "Permission denied when measuring melt pool dimensions",
    #             "PERMISSION_DENIED",
    #             workspace_name=workspace,
    #             exception_type=type(e).__name__,
    #         )
    #
    #     except Exception as e:
    #         return tool_error(
    #             "Failed to measure melt pool dimensions",
    #             "SOLVER_MEASURE_MELT_POOL_DIMENSIONS_FAILED",
    #             workspace_name=workspace,
    #             exception_type=type(e).__name__,
    #             exception_message=str(e)
    #         )

    @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",
        run_name: str | None = None,
    ) -> Union[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.
        """
        from am.cli.utils import get_workspace_path
        from am.schema import BuildParameters, Material
        from am.segmenter.types import Segment
        from am.solver import Solver
        from am.solver.types import MeshConfig

        try:
            workspace_path = get_workspace_path(workspace)

            solver = Solver()

            # 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"

            segments = Segment.load(segments_file_path)

            # Configs
            solver_configs_path = workspace_path / "solver" / "config"
            build_parameters = BuildParameters.load(
                workspace_path / "build_parameters" / build_parameters_filename
            )
            material = Material.load(workspace_path / "materials" / material_filename)
            mesh_config = MeshConfig.load(
                solver_configs_path / "mesh" / mesh_config_filename
            )

            model_name = "eagar-tsai"

            run_out_path = solver.run_layer(
                segments,
                build_parameters,
                material,
                mesh_config,
                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),
            )

    @app.tool(
        title="Visualize Solver",
        description="Generates the visualizations for a specific mesh that was generated and saved by a solver.",
        structured_output=True,
    )
    def solver_visualize_2D(
        workspace: str,
        run_name: str | None = None,
        output_folder: SolverOutputFolder = SolverOutputFolder.meshes,
        frame_format: str = "png",
        include_axis: bool = True,
        transparent: bool = False,
        units: str = "mm",
    ) -> Union[ToolSuccess[Path], ToolError]:
        """Generates visualizations of mesh"""
        from am.cli.utils import get_workspace_path
        from am.solver import Solver

        try:
            workspace_path = get_workspace_path(workspace)

            runs_folder = workspace_path / "meshes"
            if run_name is None:
                # Get list of subdirectories sorted by modification time (newest first)
                run_dirs = sorted(
                    [d for d in runs_folder.iterdir() if d.is_dir()],
                    key=lambda d: d.stat().st_mtime,
                    reverse=True,
                )

                if not run_dirs:
                    raise FileNotFoundError(
                        f"❌ No run directories found in {runs_folder}"
                    )

                run_name = run_dirs[0].name

            animation_out_path = Solver.visualize_2D(
                workspace_path,
                run_name,
                output_folder,
                frame_format=frame_format,
                include_axis=include_axis,
                transparent=transparent,
                units=units,
            )

            return tool_success(animation_out_path)

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

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

    _ = (
        solver_initialize,
        solver_run_layer,
        solver_visualize_2D,
        # solver_measure_melt_pool_dimensions,
    )
