"""Plot rendering mixin for processed radar result collections."""

from __future__ import annotations

from PyQt6.QtCore import QRectF, Qt
import numpy as np
import pyqtgraph as pg

from python_app.gui.plotting.bscan_history import (
    build_bscan_signature,
    pick_bscan_display_key,
    rebuild_bscan_history_from_results,
)
from python_app.gui.plotting.bscan_math import (
    bscan_levels,
    bscan_lookup_table,
    build_lut,
)
from python_app.models.dataset_model import ResultCollection, TraceData


class AppWindowPlotMixin:
    """Renders result collections on the main pyqtgraph plot."""

    def _draw_preferred_collection(
        self,
        *,
        result_latest: ResultCollection | None,
    ) -> None:
        """Draw latest available result collection if present."""
        if result_latest is None:
            return
        self._draw_results(result_latest)

    def _draw_results(self, collection: ResultCollection) -> bool:
        """Draw collection based on currently selected processing mode."""
        if self._processing_mode.currentText() == "bscan":
            return self._draw_bscan_heatmap(collection)
        return self._draw_trace_lines(collection)

    def _show_magnitude_curves(self) -> bool:
        """Return whether magnitude curves should be rendered."""
        return self._show_magnitude_checkbox.isChecked()

    def _show_phase_curves(self) -> bool:
        """Return whether phase curves should be rendered."""
        return self._show_phase_checkbox.isChecked()

    def _on_trace_visibility_changed(self, *_args) -> None:
        """Redraw pass-through traces when magnitude/phase toggles changed."""
        if self._processing_mode.currentText() == "bscan":
            return
        if self._result_history:
            self._draw_results(self._result_history[-1])
            return
        self._clear_trace_plots()

    def _clear_trace_plots(self) -> None:
        """Clear pass-through magnitude and phase plots."""
        self._trace_magnitude_plot.clear()
        self._trace_phase_plot.clear()
        self._clear_trace_legends()
        self._trace_magnitude_curves.clear()
        self._trace_phase_curves.clear()

    def _clear_trace_legends(self) -> None:
        """Remove trace plot legends to avoid stale combo-color mappings."""
        mag_legend = self._trace_magnitude_legend
        if mag_legend is not None:
            try:
                self._trace_magnitude_plot.getPlotItem().removeItem(mag_legend)
            except Exception:  # noqa: BLE001
                pass
            self._trace_magnitude_legend = None
        self._trace_magnitude_legend_combo_keys.clear()

        phase_legend = self._trace_phase_legend
        if phase_legend is not None:
            try:
                self._trace_phase_plot.getPlotItem().removeItem(phase_legend)
            except Exception:  # noqa: BLE001
                pass
            self._trace_phase_legend = None
        self._trace_phase_legend_combo_keys.clear()

    def _draw_trace_lines(self, collection: ResultCollection) -> bool:
        """Draw result payload traces as stacked magnitude/phase plots."""
        show_magnitude = self._show_magnitude_curves()
        show_phase = self._show_phase_curves()
        magnitude_plot = self._trace_magnitude_plot
        phase_plot = self._trace_phase_plot

        magnitude_plot.setVisible(show_magnitude)
        phase_plot.setVisible(show_phase)
        if not show_magnitude and not show_phase:
            self._clear_trace_plots()
            return False

        if show_magnitude:
            mag_item = magnitude_plot.getPlotItem()
            magnitude_plot.getViewBox().invertY(False)
            magnitude_plot.getViewBox().enableAutoRange(x=True, y=True)
            mag_item.showAxis("left", show=True)
            mag_item.showAxis("bottom", show=not show_phase)
            magnitude_plot.setLabel("left", "Magnitude", units="dB")
            if not show_phase:
                magnitude_plot.setLabel("bottom", "Frequency", units="Hz")

        if show_phase:
            phase_item = phase_plot.getPlotItem()
            phase_plot.getViewBox().invertY(False)
            phase_plot.getViewBox().enableAutoRange(x=True, y=False)
            phase_item.showAxis("left", show=True)
            phase_item.showAxis("bottom", show=True)
            phase_plot.setLabel("left", "Phase", units="deg")
            phase_plot.setLabel("bottom", "Frequency", units="Hz")

        palette = [
            "#4cc9f0",
            "#f72585",
            "#b8f2e6",
            "#ffd166",
            "#90be6d",
            "#ff595e",
            "#6a4c93",
            "#1982c4",
        ]

        combo_colors: dict[tuple[int, int], str] = {}
        legend_source_magnitude: dict[tuple[int, int], pg.PlotCurveItem] = {}
        legend_source_phase: dict[tuple[int, int], pg.PlotCurveItem] = {}
        active_magnitude_keys: set[tuple[int, int, int, str]] = set()
        active_phase_keys: set[tuple[int, int, int, str]] = set()
        has_data = False
        x_min = np.inf
        x_max = -np.inf
        for block in collection.blocks:
            combo_key = (int(block.combo.input_pos), int(block.combo.output_pos))
            if combo_key not in combo_colors:
                combo_colors[combo_key] = palette[len(combo_colors) % len(palette)]
            color = combo_colors[combo_key]
            combo_label = f"in{combo_key[0]}/out{combo_key[1]}"

            for payload_index, payload in enumerate(block.payloads):
                if payload.kind != 1 or payload.trace.size == 0:
                    continue
                if payload.frequency_hz.size == 0 or payload.frequency_hz.size != payload.trace.size:
                    continue
                curve_key = (
                    combo_key[0],
                    combo_key[1],
                    int(payload_index),
                    str(payload.processing_name),
                )

                local_x_min = float(np.min(payload.frequency_hz))
                local_x_max = float(np.max(payload.frequency_hz))
                x_min = min(x_min, local_x_min)
                x_max = max(x_max, local_x_max)

                if show_magnitude:
                    magnitude_values = 20.0 * np.log10(np.maximum(np.abs(payload.trace), 1e-12))
                    active_magnitude_keys.add(curve_key)
                    magnitude_curve = self._trace_magnitude_curves.get(curve_key)
                    if magnitude_curve is None:
                        magnitude_curve = pg.PlotCurveItem(pen=pg.mkPen(color, width=1.4))
                        self._trace_magnitude_curves[curve_key] = magnitude_curve
                        magnitude_plot.addItem(magnitude_curve)
                    else:
                        magnitude_curve.setPen(pg.mkPen(color, width=1.4))
                    magnitude_curve.setData(payload.frequency_hz, magnitude_values)
                    legend_source_magnitude.setdefault(combo_key, magnitude_curve)
                    has_data = True

                if show_phase:
                    active_phase_keys.add(curve_key)
                    phase_curve = self._trace_phase_curves.get(curve_key)
                    if phase_curve is None:
                        phase_curve = pg.PlotCurveItem(pen=pg.mkPen(color, width=1.2))
                        self._trace_phase_curves[curve_key] = phase_curve
                        phase_plot.addItem(phase_curve)
                    else:
                        phase_curve.setPen(pg.mkPen(color, width=1.2))
                    phase_x, phase_values = self._phase_display_arrays(payload.frequency_hz, payload.trace)
                    phase_curve.setData(phase_x, phase_values)
                    legend_source_phase.setdefault(combo_key, phase_curve)
                    has_data = True

        if show_magnitude:
            self._remove_inactive_trace_curves(
                plot=magnitude_plot,
                cache=self._trace_magnitude_curves,
                active_keys=active_magnitude_keys,
            )
        else:
            self._remove_all_trace_curves(plot=magnitude_plot, cache=self._trace_magnitude_curves)

        if show_phase:
            self._remove_inactive_trace_curves(
                plot=phase_plot,
                cache=self._trace_phase_curves,
                active_keys=active_phase_keys,
            )
        else:
            self._remove_all_trace_curves(plot=phase_plot, cache=self._trace_phase_curves)

        self._sync_trace_legends(
            show_magnitude=show_magnitude,
            show_phase=show_phase,
            magnitude_sources=legend_source_magnitude,
            phase_sources=legend_source_phase,
        )

        if has_data:
            if np.isfinite(x_min) and np.isfinite(x_max):
                if show_magnitude:
                    magnitude_plot.setXRange(x_min, x_max, padding=0.02)
                if show_phase:
                    phase_plot.setXRange(x_min, x_max, padding=0.02)
            if show_phase:
                phase_plot.setYRange(-180.0, 180.0, padding=0.02)
        return has_data

    @staticmethod
    def _remove_inactive_trace_curves(
        *,
        plot: pg.PlotWidget,
        cache: dict[tuple[int, int, int, str], pg.PlotCurveItem],
        active_keys: set[tuple[int, int, int, str]],
    ) -> None:
        """Delete curve items no longer present in latest result collection."""
        for key in list(cache.keys()):
            if key in active_keys:
                continue
            curve = cache.pop(key)
            plot.removeItem(curve)

    @staticmethod
    def _remove_all_trace_curves(
        *,
        plot: pg.PlotWidget,
        cache: dict[tuple[int, int, int, str], pg.PlotCurveItem],
    ) -> None:
        """Delete all cached curves from selected plot."""
        for curve in cache.values():
            plot.removeItem(curve)
        cache.clear()

    def _phase_display_arrays(self, frequency_hz: np.ndarray, trace: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Return phase display arrays with decimation for faster rendering."""
        max_points = int(getattr(self, "_trace_phase_render_max_points", 1200))
        if max_points > 0 and trace.size > max_points:
            step = max(1, int(np.ceil(trace.size / max_points)))
            frequency_hz = frequency_hz[::step]
            trace = trace[::step]
        phase_values = np.arctan2(trace.imag, trace.real) * (180.0 / np.pi)
        return frequency_hz, phase_values

    def _sync_trace_legends(
        self,
        *,
        show_magnitude: bool,
        show_phase: bool,
        magnitude_sources: dict[tuple[int, int], pg.PlotCurveItem],
        phase_sources: dict[tuple[int, int], pg.PlotCurveItem],
    ) -> None:
        """Rebuild legends only when active combo set changes."""
        self._sync_single_trace_legend(
            show=show_magnitude,
            plot=self._trace_magnitude_plot,
            legend_attr="_trace_magnitude_legend",
            legend_keys_attr="_trace_magnitude_legend_combo_keys",
            sources=magnitude_sources,
        )
        self._sync_single_trace_legend(
            show=show_phase,
            plot=self._trace_phase_plot,
            legend_attr="_trace_phase_legend",
            legend_keys_attr="_trace_phase_legend_combo_keys",
            sources=phase_sources,
        )

    def _sync_single_trace_legend(
        self,
        *,
        show: bool,
        plot: pg.PlotWidget,
        legend_attr: str,
        legend_keys_attr: str,
        sources: dict[tuple[int, int], pg.PlotCurveItem],
    ) -> None:
        """Rebuild one legend from provided combo->curve mapping when needed."""
        legend = getattr(self, legend_attr)
        existing_keys = getattr(self, legend_keys_attr)
        active_keys = set(sources.keys())
        if not show or not active_keys:
            if legend is not None:
                try:
                    plot.getPlotItem().removeItem(legend)
                except Exception:  # noqa: BLE001
                    pass
            setattr(self, legend_attr, None)
            existing_keys.clear()
            return

        if legend is not None and existing_keys == active_keys:
            return

        if legend is not None:
            try:
                plot.getPlotItem().removeItem(legend)
            except Exception:  # noqa: BLE001
                pass

        legend = plot.addLegend(offset=(8, 8))
        for combo_key in sorted(active_keys):
            curve = sources[combo_key]
            legend.addItem(curve, f"in{combo_key[0]}/out{combo_key[1]}")
        setattr(self, legend_attr, legend)
        existing_keys.clear()
        existing_keys.update(active_keys)

    def _draw_bscan_heatmap(self, _collection: ResultCollection) -> bool:
        """Draw B-scan image rebuilt from processed result history."""
        self._disable_phase_axis()
        self._sync_bscan_history_from_results()
        return self._draw_bscan_heatmap_from_history()

    def _draw_bscan_heatmap_from_history(self) -> bool:
        """Render B-scan heatmap from currently cached history arrays."""
        display_key = self._pick_bscan_display_key()
        if display_key is None:
            return False

        history = self._bscan_history_by_combo.get(display_key)
        depth_axis = self._bscan_depth_axis_by_combo.get(display_key)
        if not history or depth_axis is None:
            return False

        sweeps = np.vstack(history).astype(np.float32, copy=False)
        if sweeps.size == 0:
            return False

        depth_min = float(np.min(depth_axis))
        depth_max = float(np.max(depth_axis))
        depth_span = max(depth_max - depth_min, 1e-6)
        sweep_count = sweeps.shape[0]
        sweep_width = float(max(sweep_count, 1))
        x_min = 0.5
        x_max = x_min + sweep_width

        image_item = pg.ImageItem(axisOrder="row-major")
        image_item.setImage(sweeps.T, autoLevels=False)
        image_item.setRect(QRectF(x_min, depth_min, sweep_width, depth_span))

        axis_mode = self._bscan_axis.currentText()
        image_item.setLookupTable(self._bscan_lookup_table(axis_mode))
        image_item.setLevels(self._bscan_levels(sweeps, axis_mode))

        self._bscan_plot.clear()
        view_box = self._bscan_plot.getViewBox()
        view_box.invertY(True)
        view_box.enableAutoRange(x=False, y=False)
        self._bscan_plot.getPlotItem().showAxis("left", show=True)
        self._bscan_plot.getPlotItem().showAxis("bottom", show=True)
        self._bscan_plot.setLabel("bottom", "Sweep #")
        self._bscan_plot.setLabel("left", "Depth", units="m")
        self._bscan_plot.addItem(image_item)
        self._bscan_plot.setXRange(x_min, x_max, padding=0.02)
        self._bscan_plot.setYRange(depth_min, depth_max, padding=0.02)
        self._bscan_plot.setTitle(f"B-scan in{display_key[0]}/out{display_key[1]} | sweeps={sweep_count}")
        return True

    def _sync_bscan_history_from_results(self) -> None:
        """Rebuild B-scan history cache when live params or inputs changed."""
        self._advance_bscan_floor_to_cpp_window()
        signature = self._bscan_signature()
        if signature == self._bscan_render_signature:
            return
        self._rebuild_bscan_history_from_results()
        self._bscan_render_signature = signature

    def _bscan_signature(self) -> tuple[object, ...]:
        """Build state signature for B-scan history cache invalidation."""
        live_config = self._live_processing_config()
        result_history = list(self._result_history)
        return build_bscan_signature(
            live_config=live_config,
            result_history=result_history,
            history_limit=self._bscan_history_limit,
            floor_collection_id=self._bscan_history_floor_collection_id,
        )

    def _rebuild_bscan_history_from_results(self) -> None:
        """Recompute B-scan history cache from results history buffer."""
        result_history = list(self._result_history)
        history_by_combo, depth_axis_by_combo = rebuild_bscan_history_from_results(
            result_history=result_history,
            history_limit=self._bscan_history_limit,
            floor_collection_id=self._bscan_history_floor_collection_id,
        )
        self._bscan_history_by_combo = history_by_combo
        self._bscan_depth_axis_by_combo = depth_axis_by_combo

    def _pick_bscan_display_key(self) -> tuple[int, int] | None:
        """Choose combo history key to render."""
        return pick_bscan_display_key(self._bscan_history_by_combo)

    def _bscan_lookup_table(self, axis_mode: str) -> np.ndarray:
        """Return lookup table for current B-scan axis mode."""
        return bscan_lookup_table(axis_mode)

    @staticmethod
    def _build_lut(stops: list[str], *, size: int = 256) -> np.ndarray:
        """Backward-compatible wrapper around LUT builder."""
        return build_lut(stops, size=size)

    @staticmethod
    def _bscan_levels(sweeps: np.ndarray, axis_mode: str) -> tuple[float, float]:
        """Return display levels for B-scan image."""
        return bscan_levels(sweeps, axis_mode)

    def _clear_bscan_plot_history(self) -> None:
        """Drop cached B-scan history and invalidate cache signature."""
        self._bscan_history_by_combo.clear()
        self._bscan_depth_axis_by_combo.clear()
        self._bscan_render_signature = None

    def _advance_bscan_floor_to_cpp_window(self) -> None:
        """Clamp B-scan source history to C++ available replay window."""
        if not self._result_history:
            return

        cpp_window_limit = min(
            int(self._defaults_config.rings.preprocessed.capacity),
            int(self._defaults_config.rings.results.capacity),
        )
        cpp_window_limit = max(1, cpp_window_limit)
        latest_collection_id = int(self._result_history[-1].collection_id)
        current_floor = int(self._bscan_history_floor_collection_id)

        # Collection ids restart from 1 on new C++ run; release floor only while
        # acquisition is running, so manual "remove last" behavior in stopped mode
        # remains deterministic.
        if latest_collection_id < current_floor and self._supervisor.is_running():
            self._bscan_history_floor_collection_id = 0
            current_floor = 0

        floor_candidate = max(0, latest_collection_id - cpp_window_limit)
        if floor_candidate > current_floor:
            self._bscan_history_floor_collection_id = floor_candidate

    def _ensure_phase_view_box(self) -> pg.ViewBox:
        """Create or return secondary right-axis ViewBox for phase curves."""
        plot_item = self._bscan_plot.getPlotItem()
        phase_view_box = self._phase_viewbox
        if phase_view_box is None:
            phase_view_box = pg.ViewBox()
            self._phase_viewbox = phase_view_box
            plot_item.scene().addItem(phase_view_box)
            plot_item.getAxis("right").linkToView(phase_view_box)
            phase_view_box.setXLink(plot_item.vb)
            plot_item.vb.sigResized.connect(self._update_phase_view_box_geometry)
        self._update_phase_view_box_geometry()
        return phase_view_box

    def _update_phase_view_box_geometry(self) -> None:
        """Keep right-axis ViewBox geometry in sync with main plot ViewBox."""
        phase_view_box = self._phase_viewbox
        if phase_view_box is None:
            return
        plot_item = self._bscan_plot.getPlotItem()
        phase_view_box.setGeometry(plot_item.vb.sceneBoundingRect())
        phase_view_box.linkedViewChanged(plot_item.vb, phase_view_box.XAxis)

    def _clear_phase_overlay(self) -> None:
        """Remove all phase curves from secondary ViewBox."""
        self._trace_phase_plot.clear()

    def _disable_phase_axis(self) -> None:
        """Hide right axis and clear phase overlay when phase is not rendered."""
        self._clear_phase_overlay()

    def _result_collection_has_trace(self, collection: ResultCollection) -> bool:
        """Return `True` when collection contains at least one trace payload."""
        for block in collection.blocks:
            for payload in block.payloads:
                if payload.kind == 1 and payload.trace.size > 0:
                    return True
        return False

    def _draw_single_trace(self, trace: TraceData, title: str) -> None:
        """Draw one trace on stacked magnitude/phase plots."""
        show_magnitude = self._show_magnitude_curves()
        show_phase = self._show_phase_curves()
        magnitude_plot = self._trace_magnitude_plot
        phase_plot = self._trace_phase_plot

        magnitude_plot.setVisible(show_magnitude)
        phase_plot.setVisible(show_phase)
        self._clear_trace_plots()
        if not show_magnitude and not show_phase:
            return

        if show_magnitude:
            magnitude_plot.getViewBox().invertY(False)
            magnitude_plot.getViewBox().enableAutoRange(x=True, y=True)
            magnitude_plot.getPlotItem().showAxis("bottom", show=not show_phase)
            magnitude_plot.setLabel("left", "Magnitude", units="dB")
            magnitude_plot.setTitle(title)
            if not show_phase:
                magnitude_plot.setLabel("bottom", "Frequency", units="Hz")
        if show_phase:
            phase_plot.getViewBox().invertY(False)
            phase_plot.getViewBox().enableAutoRange(x=True, y=False)
            phase_plot.getPlotItem().showAxis("bottom", show=True)
            phase_plot.setLabel("left", "Phase", units="deg")
            phase_plot.setLabel("bottom", "Frequency", units="Hz")
            phase_plot.setTitle(title)

        if show_magnitude:
            magnitude_db = 20.0 * np.log10(np.maximum(np.abs(trace.s21), 1e-12))
            magnitude_curve = pg.PlotCurveItem(
                trace.frequency_hz,
                magnitude_db,
                pen=pg.mkPen("#ffd166", width=1.8),
            )
            magnitude_plot.addItem(magnitude_curve)
            self._trace_magnitude_curves[
                (int(trace.combo.input_pos), int(trace.combo.output_pos), 0, "__single_trace__")
            ] = magnitude_curve

        if show_phase:
            phase_deg = np.degrees(np.angle(trace.s21))
            phase_curve = pg.PlotCurveItem(
                trace.frequency_hz,
                phase_deg,
                pen=pg.mkPen("#80ed99", width=1.4, style=Qt.PenStyle.DashLine),
            )
            phase_plot.addItem(phase_curve)
            self._trace_phase_curves[
                (int(trace.combo.input_pos), int(trace.combo.output_pos), 0, "__single_trace__")
            ] = phase_curve
            phase_plot.setYRange(-180.0, 180.0, padding=0.02)

        if np.size(trace.frequency_hz) > 1:
            x_min = float(np.min(trace.frequency_hz))
            x_max = float(np.max(trace.frequency_hz))
            if show_magnitude:
                magnitude_plot.setXRange(x_min, x_max, padding=0.02)
            if show_phase:
                phase_plot.setXRange(x_min, x_max, padding=0.02)
