from dash_extensions.enrich import html, dcc, Output, Input, State


def TrafficLight(component_id, size=60):
    """
    Creates a reusable traffic light component with adjustable size.

    Args:
        component_id (str): The unique identifier for the component.
        size (int): The size of the traffic lights (width and height in pixels).

    Returns:
        html.Div: A Dash component representing the traffic light.
    """
    size_px = f"{size}px"
    font_size = round(16*size/60)
    label_style = dict(color="white", whiteSpace="nowrap", fontSize=font_size)

    return html.Div(
        id=component_id,
        style={
            "width": f"{size * 4}px",  # Adjust container width based on size
            "backgroundColor": "black",
            "borderRadius": "10px",
            "padding": "10px",
            "display": "flex",
            "flexDirection": "column",
            "alignItems": "flex-start",
        },
        children=[
            dcc.Store(id=f"{component_id}-active", data=0),  # Store to track active state
            # Green Light with Label
            html.Div(
                style={"display": "flex", "alignItems": "center", "marginBottom": "10px"},
                children=[
                    html.Div(
                        id=f"{component_id}-green-light",
                        style={
                            "width": size_px,
                            "height": size_px,
                            "backgroundColor": "gray",
                            "borderRadius": "50%",
                            "marginRight": "10px",
                        },
                    ),
                    html.Div("Schedule Optimal", style=label_style),
                ],
            ),
            # Yellow Light with Label
            html.Div(
                style={"display": "flex", "alignItems": "center", "marginBottom": "10px"},
                children=[
                    html.Div(
                        id=f"{component_id}-yellow-light",
                        style={
                            "width": size_px,
                            "height": size_px,
                            "backgroundColor": "gray",
                            "borderRadius": "50%",
                            "marginRight": "10px",
                        },
                    ),
                    html.Div("Schedule Feasible", style=label_style),
                ],
            ),
            # Orange Light with Label
            html.Div(
                style={"display": "flex", "alignItems": "center", "marginBottom": "10px"},
                children=[
                    html.Div(
                        id=f"{component_id}-orange-light",
                        style={
                            "width": size_px,
                            "height": size_px,
                            "backgroundColor": "gray",
                            "borderRadius": "50%",
                            "marginRight": "10px",
                        },
                    ),
                    html.Div(
                        ["Schedule Executable ", html.Br(), "(but Infeasible)"],
                        style=label_style,
                    ),
                ],
            ),
            # Red Light with Label
            html.Div(
                style={"display": "flex", "alignItems": "center"},
                children=[
                    html.Div(
                        id=f"{component_id}-red-light",
                        style={
                            "width": size_px,
                            "height": size_px,
                            "backgroundColor": "gray",
                            "borderRadius": "50%",
                            "marginRight": "10px",
                        },
                    ),
                    html.Div("No Schedule Found", style=label_style),
                ],
            ),
        ],
    )



# Callback to update traffic lights
def register_traffic_light_callbacks(app, component_id):
    """
    Registers the callback for the traffic light component.

    Args:
        app (Dash): The Dash app instance.
        component_id (str): The unique identifier of the traffic light.
    """

    @app.callback(
        [
            Output(f"{component_id}-green-light", "style"),
            Output(f"{component_id}-yellow-light", "style"),
            Output(f"{component_id}-orange-light", "style"),
            Output(f"{component_id}-red-light", "style"),
        ],
        [
            Input(f"{component_id}-active", "data")
        ],
        [
            State(f"{component_id}-green-light", "style"),
            State(f"{component_id}-yellow-light", "style"),
            State(f"{component_id}-orange-light", "style"),
            State(f"{component_id}-red-light", "style"),
        ]
    )
    def update_traffic_light(active, green_style, yellow_style, orange_style, red_style):
        """
        Updates the traffic light's active light based on the `active` property,
        while preserving existing style entries.

        Args:
            active (int): The index of the active light (0 for green, 1 for yellow, 2 for orange, 3 for red).
            green_style, yellow_style, orange_style, red_style (dict): Current styles of each light.

        Returns:
            tuple: Updated styles for each light.
        """

        def update_style(current_style, color, is_active):
            # Create a new dictionary to preserve existing styles while updating relevant keys
            updated_style = current_style.copy()  # Copy the current style
            updated_style.update({
                "backgroundColor": color if is_active else "gray",
                "boxShadow": f"0 0 20px {color}" if is_active else "",
            })
            return updated_style

        return (
            update_style(green_style or {}, "green", active == 0),
            update_style(yellow_style or {}, "yellow", active == 1),
            update_style(orange_style or {}, "orange", active == 2),
            update_style(red_style or {}, "red", active == 3),
        )
