"""Plot rendering callbacks.

Builds a PlotSpec from the UI state and delegates to PlotEngine.render().
This is the single rendering path — the same PlotEngine used by the
scripting API renders the web UI plots.

Transform application uses TransformChain, making the pipeline
serializable and reproducible.

The plot figure is wrapped in a ``FigureResampler`` for dynamic LTTB
downsampling on zoom/pan events. A separate ``relayoutData`` callback
triggers resampling when the user interacts with the axes.
"""

from __future__ import annotations

import logging
from collections import defaultdict
from typing import Any

import dash
import numpy as np
import plotly.graph_objects as go
from dash import Input, Output, State, no_update
from plotly_resampler import FigureResampler

from impakt.channel.model import Channel
from impakt.plot.engine import DEFAULT_COLORS, PlotEngine
from impakt.plot.spec import ChannelRef, Corridor, CorridorStyle, PlotSpec, PlotStyle
from impakt.transform.align import XAlign, YAlign
from impakt.transform.base import TransformChain
from impakt.transform.cfc import CFCFilter
from impakt.transform.resultant import resultant_from_channels
from impakt.web.state import AppState

logger = logging.getLogger(__name__)

# Module-level engine instance (stateless, safe to reuse)
_engine = PlotEngine()

# Server-side storage for the current FigureResampler instance.
# Single-user desktop app — no multi-tenant isolation needed.
_current_resampler: FigureResampler | None = None


def _build_transform_chain(
    cfc_value: str,
    y_align: bool,
    x_align_method: str,
    x_align_value: float | None,
    per_channel_cfc: str | None = None,
) -> TransformChain:
    """Build a TransformChain from the current UI state.

    Per-channel CFC override takes precedence over the global CFC setting.
    """
    chain = TransformChain()

    # CFC filter
    effective_cfc = per_channel_cfc if per_channel_cfc else cfc_value
    if effective_cfc and effective_cfc != "none":
        try:
            chain = chain.append(CFCFilter(cfc_class=int(effective_cfc)))
        except (ValueError, Exception):
            pass

    # Y-align
    if y_align:
        chain = chain.append(YAlign())

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

    return chain


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 transformed Channel objects.

    Uses TransformChain for each channel. Per-channel overrides from
    app_state.channel_overrides take precedence over the global CFC.

    Returns list of (label, transformed_channel) tuples.
    """
    channels: list[tuple[str, Channel]] = []
    multi_test = len(app_state.tests) > 1

    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

        # Build per-channel transform chain
        override = app_state.channel_overrides.get(key, {})
        per_ch_cfc = override.get("cfc", "")
        chain = _build_transform_chain(
            cfc_value,
            y_align,
            x_align_method,
            x_align_value,
            per_channel_cfc=per_ch_cfc,
        )

        # Apply the chain
        if len(chain) > 0:
            ch = chain.apply(ch)

        # Build label
        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 multi_test:
                        res_label = f"[{comps[0].source_test_id}] {res_label}"
                    channels.append((res_label, res))
                except Exception:
                    pass

    return channels


def _build_plot_spec(
    channels: list[tuple[str, Channel]],
    cursor_x1: float | None,
    cursor_x2: float | None,
    corridors: list[dict[str, Any]] | None = None,
) -> PlotSpec:
    """Build a PlotSpec from resolved channels and UI state.

    Channels are already transformed — they are wrapped in ChannelRef
    objects with no additional transform chain (transforms were applied
    during resolution).
    """
    # Build ChannelRef objects
    refs: list[ChannelRef] = []
    for i, (label, ch) in enumerate(channels):
        color = DEFAULT_COLORS[i % len(DEFAULT_COLORS)]
        refs.append(
            ChannelRef(
                channel=ch,
                style=PlotStyle(label=label, color=color),
            )
        )

    # Build Corridor objects from raw dicts
    corridor_objs: list[Corridor] = []
    if corridors:
        for c in corridors:
            if not c.get("visible", True):
                continue
            corridor_objs.append(
                Corridor(
                    name=c.get("name", "Corridor"),
                    time=np.array(c["time"]),
                    lower=np.array(c["lower"]),
                    upper=np.array(c["upper"]),
                    style=CorridorStyle(),
                )
            )

    # Cursor positions
    x_cursors = None
    if cursor_x1 is not None and cursor_x2 is not None:
        x_cursors = (cursor_x1, cursor_x2)

    # Y-axis label from first channel
    y_label = channels[0][1].unit if channels else ""

    return PlotSpec(
        channels=refs,
        corridors=corridor_objs,
        x_cursors=x_cursors,
        y_label=y_label,
        compact=True,  # Web UI always uses compact mode
    )


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

    @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-x1", "value"),
            Input("cursor-x2", "value"),
            Input("channel-overrides-store", "data"),
            Input("corridors-store", "data"),
        ],
        [
            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,
        cursor_x1: float | None,
        cursor_x2: float | None,
        channel_overrides: dict[str, Any] | None,
        corridors_data: list[dict[str, Any]] | None,
        x_align_value: float | None,
    ) -> go.Figure:
        global _current_resampler

        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,
        )

        spec = _build_plot_spec(channels, cursor_x1, cursor_x2, corridors_data)
        fig = _engine.render(spec, resample=True)

        # Store the resampler for relayoutData callbacks
        if isinstance(fig, FigureResampler):
            _current_resampler = fig

        return fig

    @app.callback(
        Output({"type": "plot-graph", "index": 0}, "figure", allow_duplicate=True),
        Input({"type": "plot-graph", "index": 0}, "relayoutData"),
        prevent_initial_call=True,
    )
    def resample_on_zoom(relayout_data: dict[str, Any] | None) -> Any:
        """Resample traces when the user zooms or pans.

        FigureResampler intercepts relayoutData events and returns a
        Patch that updates only the visible trace data — no full re-render.
        """
        if _current_resampler is None or relayout_data is None:
            return no_update
        return _current_resampler.construct_update_data_patch(relayout_data)
