"""Main entry point for mesh visualization with backend selection."""

from typing import Literal, TYPE_CHECKING
import torch

if TYPE_CHECKING:
    from torchmesh.mesh import Mesh


def draw_mesh(
    mesh: "Mesh",
    backend: Literal["matplotlib", "pyvista", "auto"] = "auto",
    show: bool = True,
    point_scalars: None | torch.Tensor | str | tuple[str, ...] = None,
    cell_scalars: None | torch.Tensor | str | tuple[str, ...] = None,
    cmap: str = "viridis",
    vmin: float | None = None,
    vmax: float | None = None,
    alpha_points: float = 1.0,
    alpha_cells: float = 1.0,
    alpha_edges: float = 1.0,
    show_edges: bool = True,
    ax=None,
    **kwargs,
):
    """Draw a mesh using matplotlib or PyVista backend.

    This is the main visualization function for Mesh objects. It automatically
    selects the appropriate backend based on spatial dimensions, or allows
    explicit backend specification.

    Args:
        mesh: Mesh object to visualize
        backend: Visualization backend to use:
            - "auto": Automatically select based on n_spatial_dims
              (matplotlib for 0D/1D/2D, PyVista for 3D)
            - "matplotlib": Force matplotlib backend (supports 3D via mplot3d)
            - "pyvista": Force PyVista backend (requires n_spatial_dims <= 3)
        show: Whether to display the plot immediately (calls plt.show() or
            plotter.show()). If False, returns the plotter/axes for further
            customization before display.
        point_scalars: Scalar data to color points. Mutually exclusive with
            cell_scalars. Can be:
            - None: Points use neutral color (black)
            - torch.Tensor: Direct scalar values, shape (n_points,) or
              (n_points, ...) where trailing dimensions are L2-normed
            - str or tuple[str, ...]: Key to lookup in mesh.point_data
        cell_scalars: Scalar data to color cells. Mutually exclusive with
            point_scalars. Can be:
            - None: Cells use neutral color (lightblue if no scalars,
              lightgray if point_scalars active)
            - torch.Tensor: Direct scalar values, shape (n_cells,) or
              (n_cells, ...) where trailing dimensions are L2-normed
            - str or tuple[str, ...]: Key to lookup in mesh.cell_data
        cmap: Colormap name for scalar visualization (default: "viridis")
        vmin: Minimum value for colormap normalization. If None, uses data min.
        vmax: Maximum value for colormap normalization. If None, uses data max.
        alpha_points: Opacity for points, range [0, 1] (default: 1.0)
        alpha_cells: Opacity for cells/faces, range [0, 1] (default: 0.3)
        alpha_edges: Opacity for cell edges, range [0, 1] (default: 0.7)
        show_edges: Whether to draw cell edges (default: True)
        ax: (matplotlib only) Existing matplotlib axes to plot on. If None,
            creates new figure and axes.
        **kwargs: Additional backend-specific keyword arguments

    Returns:
        - matplotlib backend: matplotlib.axes.Axes object
        - PyVista backend: pyvista.Plotter object

    Raises:
        ValueError: If both point_scalars and cell_scalars are specified,
            or if n_spatial_dims is not supported by the chosen backend,
            or if backend selection fails.

    Example:
        >>> # Draw mesh with automatic backend selection
        >>> mesh.draw()
        >>>
        >>> # Color cells by pressure data
        >>> mesh.draw(cell_scalars="pressure", cmap="coolwarm")
        >>>
        >>> # Color points by velocity magnitude (computing norm of vector field)
        >>> mesh.draw(point_scalars="velocity")  # velocity is (n_points, 3)
        >>>
        >>> # Use nested TensorDict key
        >>> mesh.draw(cell_scalars=("flow", "temperature"))
        >>>
        >>> # Customize and display later
        >>> ax = mesh.draw(show=False, backend="matplotlib")
        >>> ax.set_title("My Mesh")
        >>> plt.show()
    """
    ### Validate and process scalar data
    from torchmesh.visualization._scalar_utils import validate_and_process_scalars

    point_scalar_values, cell_scalar_values, active_scalar_source = (
        validate_and_process_scalars(
            point_scalars=point_scalars,
            cell_scalars=cell_scalars,
            point_data=mesh.point_data,
            cell_data=mesh.cell_data,
            n_points=mesh.n_points,
            n_cells=mesh.n_cells,
        )
    )

    ### Determine backend
    if backend == "auto":
        if mesh.n_spatial_dims in {0, 1, 2}:
            backend = "matplotlib"
        elif mesh.n_spatial_dims == 3:
            backend = "pyvista"
        else:
            raise ValueError(
                f"Cannot automatically select backend for {mesh.n_spatial_dims=}.\n"
                f"Supported spatial dimensions: 0, 1, 2, 3.\n"
                f"Please specify backend explicitly."
            )

    ### Validate backend compatibility
    if backend == "pyvista" and mesh.n_spatial_dims > 3:
        raise ValueError(
            f"PyVista backend does not support {mesh.n_spatial_dims=}.\n"
            f"Maximum spatial dimensions for PyVista: 3."
        )

    if backend == "matplotlib" and mesh.n_spatial_dims > 3:
        raise ValueError(
            f"Matplotlib backend does not support {mesh.n_spatial_dims=}.\n"
            f"Maximum spatial dimensions for matplotlib: 3."
        )

    ### Dispatch to appropriate backend
    if backend == "matplotlib":
        from torchmesh.visualization._matplotlib_impl import draw_mesh_matplotlib

        return draw_mesh_matplotlib(
            mesh=mesh,
            point_scalar_values=point_scalar_values,
            cell_scalar_values=cell_scalar_values,
            active_scalar_source=active_scalar_source,
            show=show,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            alpha_points=alpha_points,
            alpha_cells=alpha_cells,
            alpha_edges=alpha_edges,
            show_edges=show_edges,
            ax=ax,
            **kwargs,
        )

    elif backend == "pyvista":
        from torchmesh.visualization._pyvista_impl import draw_mesh_pyvista

        if ax is not None:
            raise ValueError(
                "The 'ax' parameter is only supported for matplotlib backend.\n"
                "PyVista backend creates its own plotter."
            )

        return draw_mesh_pyvista(
            mesh=mesh,
            point_scalar_values=point_scalar_values,
            cell_scalar_values=cell_scalar_values,
            active_scalar_source=active_scalar_source,
            show=show,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            alpha_points=alpha_points,
            alpha_cells=alpha_cells,
            show_edges=show_edges,
            **kwargs,
        )

    else:
        raise ValueError(
            f"Unknown backend: {backend!r}. "
            f"Supported backends: 'matplotlib', 'pyvista', 'auto'."
        )
