import numpy as np
import sympy as sp
import plotly.graph_objects as go

def plot_3d(exprs, var,
            # --- Plot Labels ---
            title="3D Plot", xlabel="x", ylabel="y", zlabel="Value", labels=None,
            # --- Styling ---
            colormap=None, show_colorbar=True,
            # --- Limits ---
            xlim=None, ylim=None, zlim=None,
            # --- Resolution ---
            resolution=100,
            # --- Display ---
            show=True,
            # --- Figure Size ---
            width=800, height=700):
    
    """
    Plots 3D surfaces from symbolic expressions using Plotly with internal helpers.

    Parameters:
        exprs      : SymPy expression or list of expressions
        var        : (x_sym, x_range, y_sym, y_range)
        title      : Plot title (default="3D Plot")
        labels     : List of labels for surfaces
        colormap   : None, string, or list (Plotly scale names or single colors)
        resolution : Grid resolution per axis (default=100)
        show       : Whether to display the plot (default=True)

    Returns:
        plotly.graph_objects.Figure
    """

    # Color Processing
    def _get_colorscale(cmap_input):
        SINGLE_COLORS = {
            'red', 'blue', 'green', 'yellow', 'orange', 'purple', 'pink', 
            'cyan', 'magenta', 'brown', 'gray', 'grey', 'black', 'white',
            'navy', 'teal', 'lime', 'olive', 'maroon', 'aqua', 'fuchsia',
            'silver', 'gold', 'indigo', 'violet', 'crimson', 'coral'
        }
        
        if cmap_input is None:
            return "Viridis" # Default
            
        # Check if single solid color
        is_single = (cmap_input.lower() in SINGLE_COLORS or 
                     cmap_input.startswith('#') or 
                     cmap_input.startswith('rgb'))
        
        if is_single:
            # Create a uniform colorscale (start and end are same color)
            return [[0.0, cmap_input], [1.0, cmap_input]]
            
        return cmap_input

    # Surface Evaluation
    def _evaluate_surface(expression, x_sym, y_sym, x_grid, y_grid):
        expression = sp.sympify(expression)
        f = sp.lambdify((x_sym, y_sym), expression, modules="numpy")
        Z_vals = f(x_grid, y_grid)
        
        # Handle constant planes (scalar result)
        if np.isscalar(Z_vals):
            Z_vals = np.full_like(x_grid, Z_vals, dtype=float)
            
        return Z_vals

    # Build Figure
    def _build_figure(surface_data, x_range, y_range):
        fig = go.Figure()
        
        x_min, x_max = float(x_range[0]), float(x_range[1])
        y_min, y_max = float(y_range[0]), float(y_range[1])
        
        x_vals = np.linspace(x_min, x_max, resolution)
        y_vals = np.linspace(y_min, y_max, resolution)
        X, Y = np.meshgrid(x_vals, y_vals)

        # Handle List Inputs
        expr_list = exprs if isinstance(exprs, list) else [exprs]
        
        # Handle Colormap Inputs
        if colormap is None:
            cmaps = [None] * len(expr_list)
        elif isinstance(colormap, (list, tuple)):
            cmaps = list(colormap)
        else:
            cmaps = [colormap] * len(expr_list)

        for i, expr in enumerate(expr_list):
            # Label
            lbl = labels[i] if labels and i < len(labels) else f"Expr {i + 1}"
            
            # Color
            cmap_raw = cmaps[i] if i < len(cmaps) else None
            scale = _get_colorscale(cmap_raw)

            # Evaluate
            Z = _evaluate_surface(expr, var[0], var[2], X, Y)
            
            # Add Trace
            fig.add_trace(go.Surface(
                x=X, y=Y, z=Z,
                name=lbl,
                colorscale=scale,
                showscale=show_colorbar,
                opacity=0.8
            ))
            
        return fig

    # Main Execution
    if not isinstance(var, tuple) or len(var) != 4:
        raise ValueError("`var` must be (x_sym, x_range, y_sym, y_range)")

    # Build
    final_fig = _build_figure(exprs, var[1], var[3])

    # Layout
    final_fig.update_layout(
        title=title,
        autosize=False, width=width, height=height,
        scene=dict(
            xaxis_title=xlabel,
            yaxis_title=ylabel,
            zaxis_title=zlabel,
            xaxis=dict(range=list(xlim) if xlim else None),
            yaxis=dict(range=list(ylim) if ylim else None),
            zaxis=dict(range=list(zlim) if zlim else None),
            aspectratio=dict(x=1, y=1, z=1),
            camera=dict(eye=dict(x=1.87, y=0.88, z=1.5)),
        ),
        margin=dict(l=10, r=10, t=50, b=10),
        legend=dict(x=0, y=1),
    )

    if show:
        final_fig.show()

    return final_fig