"""Plot rendering callbacks.

Handles:
- Updating plot figures when channels/transforms change
- Cursor line rendering
- Layout switching
"""

from __future__ import annotations

from collections import defaultdict
from typing import Any

import dash
import numpy as np
import plotly.graph_objects as go
from dash import ALL, Input, Output, State, ctx
from dash.exceptions import PreventUpdate

from impakt.channel.model import Channel
from impakt.plot.engine import DEFAULT_COLORS
from impakt.transform.align import XAlign, YAlign
from impakt.transform.cfc import CFCFilter
from impakt.transform.resultant import resultant_from_channels
from impakt.web.state import AppState


def _resolve_channels(
    selected_keys: list[str],
    app_state: AppState,
    cfc_value: str,
    y_align: bool,
    x_align_method: str,
    x_align_value: float | None,
    show_resultant: bool,
) -> list[tuple[str, Channel]]:
    """Resolve selected channel keys to Channel objects with transforms applied."""
    channels: list[tuple[str, Channel]] = []

    for key in selected_keys:
        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 None:
            continue

        # Apply CFC filter
        if cfc_value != "none":
            try:
                ch = CFCFilter(cfc_class=int(cfc_value)).apply(ch)
            except (ValueError, Exception):
                pass

        # Apply Y-align
        if y_align:
            ch = YAlign().apply(ch)

        # Apply X-align
        if x_align_method == "manual" and x_align_value is not None:
            ch = XAlign(method="manual", reference_time=x_align_value).apply(ch)
        elif x_align_method == "threshold" and x_align_value is not None:
            ch = XAlign(method="threshold", threshold_value=x_align_value).apply(ch)

        # Build label
        multi_test = len(app_state.tests) > 1
        label = ch.code.short_label if ch.code.is_valid else ch.name
        if multi_test:
            label = f"[{test_id}] {label}"

        channels.append((label, ch))

    # Compute resultants if requested
    if show_resultant and len(channels) >= 2:
        groups: dict[str, list[tuple[str, Channel]]] = defaultdict(list)
        for label, ch in channels:
            if ch.code.is_valid and ch.code.is_component():
                groups[ch.code.group_key()].append((label, ch))

        for gkey, group_channels in groups.items():
            if len(group_channels) >= 2:
                try:
                    comps = [ch for _, ch in group_channels]
                    res = resultant_from_channels(*comps)
                    res_label = (
                        f"{res.code.location_label} Resultant" if res.code.is_valid else "Resultant"
                    )
                    if len(app_state.tests) > 1:
                        res_label = f"[{comps[0].source_test_id}] {res_label}"
                    channels.append((res_label, res))
                except Exception:
                    pass

    return channels


def _build_figure(
    channels: list[tuple[str, Channel]],
    cursor_x1: float | None,
    cursor_x2: float | None,
    cfc_value: str,
) -> go.Figure:
    """Build a Plotly figure from resolved channels."""
    fig = go.Figure()

    if not channels:
        fig.update_layout(
            template="plotly_white",
            annotations=[
                {
                    "text": "Select channels from the tree to plot",
                    "xref": "paper",
                    "yref": "paper",
                    "x": 0.5,
                    "y": 0.5,
                    "showarrow": False,
                    "font": {"size": 14, "color": "#bbb"},
                }
            ],
            xaxis={"visible": False},
            yaxis={"visible": False},
            margin={"l": 40, "r": 20, "t": 30, "b": 40},
        )
        return fig

    # Add traces
    for i, (label, ch) in enumerate(channels):
        color = DEFAULT_COLORS[i % len(DEFAULT_COLORS)]
        fig.add_trace(
            go.Scatter(
                x=ch.time.tolist(),
                y=ch.data.tolist(),
                mode="lines",
                name=label,
                line=dict(color=color, width=1.5),
                hovertemplate=f"{label}<br>t=%{{x:.6f}}s<br>%{{y:.4f}} {ch.unit}<extra></extra>",
            )
        )

    # Add cursor lines
    if cursor_x1 is not None and cursor_x2 is not None:
        for x_val, lbl in [(cursor_x1, "x1"), (cursor_x2, "x2")]:
            fig.add_vline(
                x=x_val,
                line_dash="dash",
                line_color="rgba(100,100,100,0.5)",
                line_width=1,
                annotation_text=f"{lbl}={x_val:.4f}s",
                annotation_font_size=10,
            )

    # Layout
    y_label = channels[0][1].unit if channels else ""
    cfc_label = f" (CFC {cfc_value})" if cfc_value != "none" else ""

    fig.update_layout(
        xaxis_title="Time (s)",
        yaxis_title=y_label,
        template="plotly_white",
        hovermode="x unified",
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=-0.2,
            xanchor="center",
            x=0.5,
            font=dict(size=10),
        ),
        margin=dict(l=55, r=15, t=10, b=60),
    )

    return fig


def register_plot_callbacks(app: dash.Dash, app_state: AppState) -> None:
    """Register all plot-related callbacks."""

    @app.callback(
        Output({"type": "plot-graph", "index": 0}, "figure"),
        [
            Input("selected-channels-store", "data"),
            Input("cfc-select", "value"),
            Input("show-resultant", "value"),
            Input("y-align-check", "value"),
            Input("x-align-method", "value"),
            Input("cursor-update-btn", "n_clicks"),
        ],
        [
            State("cursor-x1", "value"),
            State("cursor-x2", "value"),
            State("x-align-value", "value"),
        ],
    )
    def update_main_plot(
        selected_keys: list[str] | None,
        cfc_value: str,
        show_resultant: bool,
        y_align: bool,
        x_align_method: str,
        n_clicks: int | None,
        cursor_x1: float | None,
        cursor_x2: float | None,
        x_align_value: float | None,
    ) -> go.Figure:
        if not selected_keys:
            selected_keys = []

        channels = _resolve_channels(
            selected_keys,
            app_state,
            cfc_value,
            y_align,
            x_align_method or "none",
            x_align_value,
            show_resultant,
        )

        return _build_figure(channels, cursor_x1, cursor_x2, cfc_value)
