import json

import dash
import dash_bootstrap_components as dbc
import pandas as pd
import plotly.express as px
from dash import Input, Output, State, dash_table, dcc, html, no_update
from dash_bootstrap_templates import load_figure_template

from dot_metrics.models import BatchResult, EvalResult

load_figure_template("darkly")


def _to_dataframe(data: BatchResult | EvalResult) -> pd.DataFrame:
    if isinstance(data, EvalResult):
        data = BatchResult({"_single": data})

    rows = []
    keys_iter = iter(data)
    first_key = next(keys_iter, None)
    if first_key is None:
        return pd.DataFrame()

    tuple_keys = isinstance(first_key, tuple)
    key_names = data.key_names

    if tuple_keys and key_names is None:
        key_names = [f"key[{i}]" for i in range(len(first_key))]
    elif not tuple_keys and key_names is None:
        key_names = ["key"]

    for key in [first_key, *keys_iter]:
        eval_result = data[key]
        row: dict = {}

        if tuple_keys:
            for i, part in enumerate(key):
                row[key_names[i]] = part
        else:
            row[key_names[0]] = key

        for name, metric in eval_result.metrics.items():
            row[f"metric:{name}"] = metric.value
            if metric.debug:
                row[f"debug:{name}"] = json.dumps(metric.debug)

        for name, constraint in eval_result.constraints.items():
            row[f"constraint:{name}"] = constraint.value
            row[f"constraint_passed:{name}"] = constraint.passed
            if constraint.debug:
                row[f"debug:{name}"] = json.dumps(constraint.debug)

        if eval_result.constraints:
            row["constraints_ok"] = eval_result.constraints_ok

        rows.append(row)

    return pd.DataFrame(rows)


def _render_tree(obj, depth: int = 0) -> html.Div:
    indent = {"marginLeft": f"{depth * 16}px"}
    if isinstance(obj, dict):
        if not obj:
            return html.Span("{}", style={"color": "#6c757d"})
        items = []
        for k, v in obj.items():
            items.append(html.Div([
                html.Span(str(k) + ": ", style={"color": "#adb5bd", "fontWeight": "600"}),
                _render_tree(v, depth + 1) if isinstance(v, (dict, list)) else _render_tree(v),
            ], style={"marginBottom": "2px", **indent}))
        return html.Div(items, style={"borderLeft": "2px solid #343a40", "paddingLeft": "8px", **indent} if depth > 0 else {})
    elif isinstance(obj, list):
        if not obj:
            return html.Span("[]", style={"color": "#6c757d"})
        items = [
            html.Div([
                html.Span(f"[{i}] ", style={"color": "#6c757d", "fontSize": "11px"}),
                _render_tree(v, depth + 1) if isinstance(v, (dict, list)) else _render_tree(v),
            ], style={"marginBottom": "2px"})
            for i, v in enumerate(obj)
        ]
        return html.Div(items, style={"borderLeft": "2px solid #343a40", "paddingLeft": "8px", **indent})
    elif isinstance(obj, bool):
        return html.Span(str(obj), style={"color": "#4ade80" if obj else "#f87171"})
    elif isinstance(obj, (int, float)):
        return html.Span(str(obj), style={"color": "#93c5fd"})
    elif obj is None:
        return html.Span("null", style={"color": "#6c757d"})
    else:
        return html.Span(str(obj), style={"color": "#f8f9fa"})


def _build_app(df: pd.DataFrame, higher_is_better: dict[str, bool] | None = None) -> dash.Dash:
    dbc_css = (
        "https://cdn.jsdelivr.net/gh/AnnMarieW/dash-bootstrap-templates/dbc.min.css"
    )
    app = dash.Dash(__name__, external_stylesheets=[dbc.themes.DARKLY, dbc_css])

    columns = list(df.columns)
    key_cols = [c for c in columns if c.startswith("key")]
    metric_cols = [c for c in columns if c.startswith("metric:")]
    numeric_cols = list(df.select_dtypes(include="number").columns)

    def _label(col: str) -> str:
        for prefix, suffix in (
            ("metric:", " (metric)"),
            ("constraint:", " (constraint)"),
            ("constraint_passed:", " passed"),
            ("debug:", ": debug"),
        ):
            if col.startswith(prefix):
                return col[len(prefix) :] + suffix
        if col in key_cols or col == "key":
            return col + " (key)"
        return col

    sorted_records = df.sort_values(metric_cols[0], ascending=False).to_dict("records") if metric_cols else df.to_dict("records")

    default_x = key_cols[0] if key_cols else columns[0]
    default_y = metric_cols[0] if metric_cols else columns[-1]
    default_color = "constraints_ok" if "constraints_ok" in columns else None
    none_option = [{"label": "None", "value": "__none__"}]
    # dbc.Select requires string values; label strips internal prefixes
    col_options = [{"label": _label(c), "value": c} for c in columns]
    numeric_options = [{"label": _label(c), "value": c} for c in numeric_cols]

    def _select(id, options, value):
        return dbc.Select(id=id, options=options, value=str(value))

    def _control(label, select):
        return dbc.Col(
            [
                dbc.Label(label, class_name="text-uppercase text-muted small"),
                select,
            ],
            xs=12,
            sm=6,
            md=4,
            lg=2,
            class_name="mb-2",
        )

    total_runs = len(df)
    passing_runs = int(df["constraints_ok"].sum()) if "constraints_ok" in columns else None

    def _stat_card(label, value):
        return dbc.Col(
            dbc.Card(dbc.CardBody([
                html.Div(label, className="text-muted small text-uppercase mb-1"),
                html.Div(str(value), className="fw-bold fs-5"),
            ], style={"padding": "10px 14px"})),
            xs=6, sm=4, md=2, class_name="mb-2",
        )

    summary_cards = [_stat_card("Total runs", total_runs)]
    if passing_runs is not None:
        summary_cards.append(_stat_card("Passing", f"{passing_runs} / {total_runs}"))
    for col in metric_cols:
        hib = (higher_is_better or {}).get(col, True)
        best = df[col].max() if hib else df[col].min()
        summary_cards.append(_stat_card(f"Best {col[len('metric:'):]}", f"{best:.3g}"))

    categorical_cols = [c for c in columns if not pd.api.types.is_numeric_dtype(df[c])]
    categorical_options = [{"label": _label(c), "value": c} for c in categorical_cols]

    agg_form = dbc.Card([
        dbc.CardBody([
            dbc.Row([
                dbc.Col([
                    dbc.Label("Group by", class_name="text-uppercase text-muted small"),
                    dcc.Dropdown(
                        id="agg-groupby",
                        options=categorical_options,
                        multi=True,
                        placeholder="Select categories...",
                        className="dbc",
                    ),
                ], xs=12, md=4, class_name="mb-2"),
                dbc.Col([
                    dbc.Label("Quantity", class_name="text-uppercase text-muted small"),
                    dbc.Select(id="agg-quantity", options=numeric_options, value=numeric_cols[0] if numeric_cols else None),
                ], xs=6, md=3, class_name="mb-2"),
                dbc.Col([
                    dbc.Label("Function", class_name="text-uppercase text-muted small"),
                    dbc.Select(
                        id="agg-func",
                        options=[
                            {"label": "Mean", "value": "mean"},
                            {"label": "Median", "value": "median"},
                            {"label": "Min", "value": "min"},
                            {"label": "Max", "value": "max"},
                        ],
                        value="mean",
                    ),
                ], xs=6, md=2, class_name="mb-2"),
                dbc.Col([
                    dbc.Label("\u00a0", class_name="small"),
                    dbc.Button("Add", id="agg-add", color="primary", size="sm", class_name="d-block"),
                ], xs=12, md=1, class_name="mb-2 d-flex flex-column"),
            ]),
            html.Div(id="agg-tags", className="d-flex flex-wrap gap-1"),
        ])
    ], class_name="mb-3")

    app.layout = dbc.Container(
        [
            dbc.Row(dbc.Col(html.H5("dot-metrics explorer", className="mt-3 mb-2 text-muted"))),
            dbc.Row(summary_cards, class_name="mb-4"),
            agg_form,
            dcc.Store(id="agg-store", data=[]),
            dbc.Card([
                dbc.CardBody([
                    dbc.Row(
                        [
                            _control(
                                "Chart type",
                                _select(
                                    "chart-type",
                                    [
                                        {"label": "Scatter", "value": "scatter"},
                                        {"label": "Bar", "value": "bar"},
                                        {"label": "Heatmap", "value": "heatmap"},
                                    ],
                                    "scatter",
                                ),
                            ),
                            _control("X", _select("x-col", col_options, default_x)),
                            _control("Y", _select("y-col", col_options, default_y)),
                            _control(
                                "Color",
                                _select(
                                    "color-col",
                                    none_option + col_options,
                                    default_color or "__none__",
                                ),
                            ),
                            _control(
                                "Size",
                                _select("size-col", none_option + numeric_options, "__none__"),
                            ),
                        ],
                        class_name="mb-0",
                    ),
                    dbc.Row(dbc.Col(dcc.Graph(id="chart"))),
                ])
            ], class_name="mb-4"),
            dbc.Row(
                dbc.Col(
                    html.Div([
                        dbc.Button("Show debug columns", id="toggle-debug", color="secondary", size="sm", class_name="me-2"),
                        dbc.Input(id="csv-filename", value="dot_metrics_results", type="text", size="sm", style={"width": "220px", "display": "inline-block"}, class_name="me-1"),
                        dbc.InputGroupText(".csv", class_name="me-2", style={"fontSize": "13px"}),
                        dbc.Button("Export CSV", id="export-csv", color="secondary", size="sm"),
                        dcc.Download(id="download-csv"),
                    ], className="d-flex align-items-center"),
                    class_name="mb-2",
                )
            ),
            dbc.Row(
                dbc.Col(
                    dash_table.DataTable(
                        id="table",
                        columns=[
                            {"name": _label(c), "id": c}
                            for c in columns
                            if not c.startswith("debug:")
                            and not c.startswith("constraint_passed:")
                        ],
                        data=sorted_records,
                        tooltip_data=[
                            {c: {"value": str(row.get(c, "")), "type": "markdown"} for c in columns}
                            for row in sorted_records
                        ],
                        tooltip_header={c: {"value": _label(c), "type": "markdown"} for c in columns},
                        tooltip_delay=500,
                        tooltip_duration=None,
                        page_size=100,
                        fixed_rows={"headers": True},
                        style_table={"overflowX": "auto"},
                        style_cell={
                            "padding": "8px 12px",
                            "fontFamily": "monospace",
                            "fontSize": "13px",
                            "textAlign": "left",
                            "maxWidth": "300px",
                            "overflow": "hidden",
                            "textOverflow": "ellipsis",
                        },
                        style_data_conditional=(
                            [
                                {
                                    "if": {"row_index": i},
                                    "backgroundColor": "#5c1a1a",
                                    "color": "#ff9999",
                                    "borderLeft": "3px solid #ff4444",
                                }
                                for i, row in enumerate(sorted_records)
                                if not row.get("constraints_ok", True)
                            ]
                            if "constraints_ok" in columns
                            else []
                        ) + [
                            {
                                "if": {"column_id": c},
                                "width": "120px",
                                "maxWidth": "120px",
                                "color": "#6ea8fe",
                                "textDecoration": "underline",
                                "cursor": "pointer",
                            }
                            for c in columns if c.startswith("debug:")
                        ],
                    ),
                    class_name="mb-4",
                )
            ),
            dbc.Modal([
                dbc.ModalHeader(dbc.ModalTitle(id="debug-modal-title")),
                dbc.ModalBody([
                    dbc.Switch(id="debug-raw-toggle", label="Raw JSON", value=False, class_name="mb-3"),
                    html.Div(id="debug-modal-body", style={"fontFamily": "monospace", "fontSize": "13px"}),
                ]),
            ], id="debug-modal", size="lg", scrollable=True),
        ],
        fluid=True,
        class_name="dbc dbc-ag-grid",
    )

    hidden_prefixes = ("debug:", "constraint_passed:")
    visible_cols = [c for c in columns if not c.startswith(hidden_prefixes)]
    all_cols = columns

    @app.callback(
        Output("table", "columns"),
        Output("toggle-debug", "children"),
        Input("toggle-debug", "n_clicks"),
    )
    def toggle_debug_cols(n_clicks):
        showing_all = (n_clicks or 0) % 2 == 1
        cols = all_cols if showing_all else visible_cols
        label = "Hide debug columns" if showing_all else "Show debug columns"
        return [{"name": _label(c), "id": c} for c in cols], label

    @app.callback(
        Output("download-csv", "data"),
        Input("export-csv", "n_clicks"),
        State("csv-filename", "value"),
        prevent_initial_call=True,
    )
    def export_csv(_, filename):
        name = (filename or "dot_metrics_results").strip()
        return dcc.send_data_frame(df.to_csv, f"{name}.csv", index=False)

    @app.callback(
        Output("debug-modal", "is_open"),
        Output("debug-modal-title", "children"),
        Output("debug-modal-body", "children"),
        Input("table", "active_cell"),
        Input("debug-raw-toggle", "value"),
        State("table", "derived_virtual_data"),
        State("table", "columns"),
        prevent_initial_call=True,
    )
    def open_debug_modal(active_cell, show_raw, virtual_data, table_cols):
        if not active_cell or not virtual_data or not table_cols:
            return no_update, no_update, no_update
        col_id = table_cols[active_cell["column"]]["id"]
        if not col_id.startswith("debug:"):
            return no_update, no_update, no_update
        raw = (virtual_data[active_cell["row"]] or {}).get(col_id)
        if not raw:
            return no_update, no_update, no_update
        data = json.loads(raw)
        title = col_id[len("debug:"):]
        if show_raw:
            body = html.Pre(json.dumps(data, indent=2), style={"color": "#f8f9fa", "fontSize": "12px"})
        else:
            body = _render_tree(data)
        return True, title, body

    @app.callback(
        Output("agg-store", "data"),
        Output("agg-tags", "children"),
        Input("agg-add", "n_clicks"),
        Input({"type": "agg-remove", "index": dash.ALL}, "n_clicks"),
        State("agg-groupby", "value"),
        State("agg-quantity", "value"),
        State("agg-func", "value"),
        State("agg-store", "data"),
        prevent_initial_call=True,
    )
    def manage_aggregations(add_click, remove_clicks, groupby, quantity, func, store):
        triggered = dash.ctx.triggered_id
        if triggered == "agg-add":
            if not groupby or not quantity or not func:
                return no_update, no_update
            col_name = f"{func}({_label(quantity)})"
            entry = {"groupby": groupby, "quantity": quantity, "func": func, "col_name": col_name}
            # avoid duplicates
            if any(s["col_name"] == col_name for s in store):
                return no_update, no_update
            store = store + [entry]
        elif isinstance(triggered, dict) and triggered.get("type") == "agg-remove":
            idx = triggered["index"]
            store = [s for i, s in enumerate(store) if i != idx]

        tags = [
            dbc.Badge([
                s["col_name"], " ",
                html.Span(
                    "\u00d7",
                    id={"type": "agg-remove", "index": i},
                    n_clicks=0,
                    style={"cursor": "pointer", "marginLeft": "4px"},
                ),
            ], color="info", class_name="me-1", style={"fontSize": "13px"})
            for i, s in enumerate(store)
        ]
        return store, tags

    def _build_agg_df(store):
        """Compute aggregated columns and merge them onto the base df."""
        extended = df.copy()
        for entry in store:
            groupby = entry["groupby"]
            quantity = entry["quantity"]
            func = entry["func"]
            col_name = entry["col_name"]
            agg_df = extended.groupby(groupby, as_index=False).agg(**{col_name: (quantity, func)})
            extended = extended.merge(agg_df, on=groupby, how="left")
        return extended

    @app.callback(
        Output("x-col", "options"),
        Output("y-col", "options"),
        Output("color-col", "options"),
        Output("size-col", "options"),
        Input("agg-store", "data"),
    )
    def update_dropdown_options(store):
        agg_col_names = [s["col_name"] for s in store]
        agg_options = [{"label": c, "value": c} for c in agg_col_names]
        all_options = col_options + agg_options
        all_numeric_options = numeric_options + agg_options
        return all_options, all_options, none_option + all_options, none_option + all_numeric_options

    @app.callback(
        Output("chart", "figure"),
        Input("chart-type", "value"),
        Input("x-col", "value"),
        Input("y-col", "value"),
        Input("color-col", "value"),
        Input("size-col", "value"),
        Input("agg-store", "data"),
    )
    def update_chart(chart_type, x, y, color, size, store):
        color = None if color == "__none__" else color
        size = None if size == "__none__" else size

        plot_df = _build_agg_df(store) if store else df.copy()

        if color and color in plot_df.columns and not pd.api.types.is_numeric_dtype(plot_df[color]):
            plot_df[color] = plot_df[color].astype(str)

        labels = {c: _label(c) for c in columns}
        kwargs = {"data_frame": plot_df, "x": x, "y": y, "labels": labels}

        if chart_type == "scatter":
            if color:
                kwargs["color"] = color
            if size:
                kwargs["size"] = size
            kwargs["hover_data"] = [c for c in plot_df.columns if not c.startswith("debug:")]
            fig = px.scatter(**kwargs)
        elif chart_type == "bar":
            if color:
                kwargs["color"] = color
            fig = px.bar(**kwargs)
        else:
            if color:
                kwargs["z"] = color
            fig = px.density_heatmap(**kwargs)

        return fig

    return app


def serve(
    data: BatchResult | EvalResult,
    *,
    host: str = "127.0.0.1",
    port: int = 8050,
    debug: bool = False,
) -> None:
    df = _to_dataframe(data)
    eval_result = data if isinstance(data, EvalResult) else next(iter(data.items()))[1]
    higher_is_better = {f"metric:{k}": v.higher_is_better for k, v in eval_result.metrics.items()}
    app = _build_app(df, higher_is_better=higher_is_better)
    app.run(host=host, port=port, debug=debug)
