from typing import List, Optional
import plotly.graph_objects as go
from plotly.subplots import make_subplots


class DashboardRow:

    def __init__(
        self,
        figures,
        theme=None,
        weight: float = 1.0,
    ):
        if not figures:
            raise ValueError("DashboardRow must contain at least one figure.")

        if weight <= 0:
            raise ValueError("Row weight must be > 0.")

        self.figures = figures
        self.theme = theme
        self.weight = weight


class Dashboard:
    """
    Professional Plotly dashboard builder with:

    - Row-based layout
    - Bootstrap-style 12 column grid
    - Mixed subplot types (xy + domain)
    - Title + legend handling
    - Theme inheritance
    """

    def __init__(
        self,
        figures: Optional[List[go.Figure]] = None,
        columns: Optional[int] = None,
        rows: Optional[List[DashboardRow]] = None,
        title: Optional[str] = None,
        theme: Optional[str] = None,
        horizontal_spacing: float = 0.06,
        vertical_spacing: float = 0.10,
        base_row_height: int = 400,
    ):
        if not figures and not rows:
            raise ValueError("Provide either figures or rows.")

        self.figures = figures
        self.columns = columns
        self.rows = rows
        self.title = title
        self.theme = theme
        self.horizontal_spacing = horizontal_spacing
        self.vertical_spacing = vertical_spacing
        self.base_row_height = base_row_height

    # ---------------------------------------------------------
    # Subplot Type Detection
    # ---------------------------------------------------------

    def _get_subplot_type(self, fig: go.Figure) -> str:
        for trace in fig.data:
            if trace.type in ("indicator", "pie"):
                return "domain"
        return "xy"

    # ---------------------------------------------------------
    # Extract Title
    # ---------------------------------------------------------

    def _extract_title(self, fig: go.Figure) -> str:
        if fig.layout.title and fig.layout.title.text:
            return fig.layout.title.text
        return ""

    # ---------------------------------------------------------
    # Apply Theme to Figure
    # ---------------------------------------------------------

    def _apply_theme_to_figure(self, fig: go.Figure, theme: Optional[str]):

        if theme:
            fig.update_layout(template=theme)

    # ---------------------------------------------------------
    # Add traces
    # ---------------------------------------------------------

    def _add_figure_to_subplot(
        self,
        dashboard_fig: go.Figure,
        fig: go.Figure,
        row: int,
        col: int,
    ):
        for trace in fig.data:
            dashboard_fig.add_trace(trace, row=row, col=col)

    # ---------------------------------------------------------
    # Professional Row Layout (12 column grid)
    # ---------------------------------------------------------

    def _build_rows(self):

        TOTAL_COLUMNS = 12
        num_rows = len(self.rows)

        # -----------------------------------------------------
        # Calculate proportional row heights
        # -----------------------------------------------------

        total_weight = sum(row.weight for row in self.rows)

        row_heights = [
            row.weight / total_weight
            for row in self.rows
        ]

        specs = []
        subplot_titles = []

        for row in self.rows:

            row_theme = row.theme or self.theme

            for fig in row.figures:
                self._apply_theme_to_figure(fig, row_theme)

            num_figs = len(row.figures)
            col_span = TOTAL_COLUMNS // num_figs

            row_specs = []

            for fig in row.figures:
                row_specs.append({
                    "type": self._get_subplot_type(fig),
                    "colspan": col_span
                })
                subplot_titles.append(self._extract_title(fig))

                for _ in range(col_span - 1):
                    row_specs.append(None)

            while len(row_specs) < TOTAL_COLUMNS:
                row_specs.append(None)

            specs.append(row_specs)

        dashboard_fig = make_subplots(
            rows=num_rows,
            cols=TOTAL_COLUMNS,
            specs=specs,
            subplot_titles=subplot_titles,
            horizontal_spacing=self.horizontal_spacing,
            vertical_spacing=self.vertical_spacing,
            row_heights=row_heights,   # 🔥 KEY ADDITION
        )

        # Add traces
        for row_idx, row in enumerate(self.rows, start=1):

            col_position = 1
            num_figs = len(row.figures)
            col_span = TOTAL_COLUMNS // num_figs

            for fig in row.figures:
                self._add_figure_to_subplot(
                    dashboard_fig,
                    fig,
                    row=row_idx,
                    col=col_position
                )
                col_position += col_span

        dashboard_fig.update_layout(
            height=self.base_row_height * total_weight,
            title=self.title,
            template=self.theme,
            showlegend=True,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=-0.1,
                xanchor="center",
                x=0.5
            ),
            margin=dict(t=80, l=40, r=40, b=80),
        )

        return dashboard_fig


    # ---------------------------------------------------------
    # Grid Layout (Optional Simple Mode)
    # ---------------------------------------------------------

    def _build_grid(self):

        total = len(self.figures)
        cols = self.columns or 2
        rows = (total + cols - 1) // cols

        specs = []
        subplot_titles = []

        for r in range(rows):
            row_specs = []
            for c in range(cols):
                idx = r * cols + c
                if idx < total:
                    fig = self.figures[idx]
                    self._apply_theme_to_figure(fig, self.theme)
                    row_specs.append({"type": self._get_subplot_type(fig)})
                    subplot_titles.append(self._extract_title(fig))
                else:
                    row_specs.append(None)
                    subplot_titles.append("")
            specs.append(row_specs)

        dashboard_fig = make_subplots(
            rows=rows,
            cols=cols,
            specs=specs,
            subplot_titles=subplot_titles,
            horizontal_spacing=self.horizontal_spacing,
            vertical_spacing=self.vertical_spacing,
        )

        for idx, fig in enumerate(self.figures):
            row = (idx // cols) + 1
            col = (idx % cols) + 1
            self._add_figure_to_subplot(dashboard_fig, fig, row, col)

        dashboard_fig.update_layout(
            height=self.base_row_height * rows,
            title=self.title,
            template=self.theme,
            showlegend=True,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=-0.05,
                xanchor="center",
                x=0.5
            ),
            margin=dict(t=80, l=40, r=40, b=80),
        )

        return dashboard_fig

    # ---------------------------------------------------------
    # Public API
    # ---------------------------------------------------------

    def build(self):
        if self.rows:
            return self._build_rows()
        return self._build_grid()

    def show(self):
        fig = self.build()
        return fig

    def save(self, filename: str):
        fig = self.build()
        fig.write_html(filename)
