from typing import TypedDict

import pandas as pd
from plotly.graph_objects import Figure

class FigureWithMaxY(TypedDict):
    """
    A helper container class that gets used by report pages to encapsulate
    Plotly figures along with the maximum y value represented on the figures.

    This aids in axis synchronization when two of the same report are shown side by side
    """
    figure:Figure
    max_y:float

def find_fig_max_y_and_generate_wrapper(fig: Figure) -> 'FigureWithMaxY':
    """
    Helper to automatically extract the max Y value from a Plotly figure
    and wrap it in the FigureWithMaxY utility class.
    """
    max_y = 0.0
    # Iterate through all traces (bars, lines, etc.) in the figure
    for trace in fig.data:
        if hasattr(trace, 'y') and trace.y is not None:
            # Get the max of this trace, ignoring NaNs
            current_max = max(trace.y) if len(trace.y) > 0 else 0
            max_y = max(max_y, float(current_max))

    return FigureWithMaxY(figure=fig, max_y=max_y)

def extract_figure_data(figure:Figure) -> pd.DataFrame:
    data_list = []
    for trace in figure.data:
        df_trace = pd.DataFrame({
            'x': trace.x,
            'y': trace.y
        })
        data_list.append(df_trace)

    if len(data_list) == 0:
        return pd.DataFrame()

    full_df = pd.concat(data_list)
    return full_df