"""Math expression builder callbacks.

Handles:
- Compute button -> evaluate expression with bound channels
- Add result as a derived channel to AppState
- Plot the derived channel
"""

from __future__ import annotations

from typing import Any

import dash
from dash import Input, Output, State, html
from dash.exceptions import PreventUpdate

from impakt.transform.math_expr import math_expr
from impakt.web.state import AppState


def register_math_callbacks(app: dash.Dash, app_state: AppState) -> None:
    """Register math expression builder callbacks."""

    @app.callback(
        [
            Output("math-status", "children"),
            Output("selected-channels-store", "data", allow_duplicate=True),
        ],
        [Input("math-compute-btn", "n_clicks")],
        [
            State("math-expression", "value"),
            State("math-var-a", "value"),
            State("math-var-b", "value"),
            State("math-var-c", "value"),
            State("math-result-name", "value"),
            State("math-result-unit", "value"),
            State("selected-channels-store", "data"),
        ],
        prevent_initial_call=True,
    )
    def compute_math_expression(
        n_clicks: int | None,
        expression: str | None,
        var_a_key: str | None,
        var_b_key: str | None,
        var_c_key: str | None,
        result_name: str | None,
        result_unit: str | None,
        current_selected: list[str] | None,
    ) -> tuple[Any, list[str]]:
        if not n_clicks or not expression:
            raise PreventUpdate

        expression = expression.strip()
        if not expression:
            return html.Div(
                "Enter an expression", className="text-warning small"
            ), current_selected or []

        # Resolve variable channels
        channels: dict[str, Any] = {}
        var_map = {"a": var_a_key, "b": var_b_key, "c": var_c_key}

        for var_name, key in var_map.items():
            if not key:
                continue
            if "::" in key:
                test_id, ch_name = key.split("::", 1)
            elif app_state.primary_test:
                test_id = app_state.primary_test.test_id
                ch_name = key
            else:
                continue

            ch = app_state.get_channel(test_id, ch_name)
            if ch is not None:
                channels[var_name] = ch

        if not channels:
            return (
                html.Div("Bind at least one variable to a channel", className="text-warning small"),
                current_selected or [],
            )

        # Compute
        name = result_name or "derived"
        unit = result_unit or ""

        try:
            result_ch = math_expr(
                expression=expression,
                channels=channels,
                name=name,
                unit=unit,
            )

            # Store the derived channel in the primary test's data
            primary = app_state.primary_test
            if primary:
                primary.data._channels[name] = result_ch

                # Add to selected channels
                result_key = f"{primary.test_id}::{name}"
                updated_selected = list(current_selected or [])
                if result_key not in updated_selected:
                    updated_selected.append(result_key)

                return (
                    html.Div(
                        f"Computed '{name}': peak={result_ch.peak:.4f} {unit}",
                        className="text-success small",
                    ),
                    updated_selected,
                )
            else:
                return (
                    html.Div("No test loaded", className="text-danger small"),
                    current_selected or [],
                )

        except Exception as e:
            return (
                html.Div(f"Error: {e}", className="text-danger small"),
                current_selected or [],
            )
