import time  # For measuring update time
import tkinter as tk
from tkinter import ttk
from typing import Any, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk

from topovision.gui.i18n import get_translator
from topovision.visualization.plot3d import (
    create_initial_surface_plot,
    update_surface_plot_data,
)


class Plot3DWindow(tk.Toplevel):
    """
    A Toplevel window to display a live 3D surface plot generated by Matplotlib.
    Includes interactive controls for plot customization and performance.
    Uses blitting for faster updates.
    """

    def __init__(self, master: tk.Tk, lang: str = "en") -> None:
        super().__init__(master)
        self.translator = get_translator(lang)
        self._ = self.translator

        self.title(self._("3d_plot_window_title"))
        self.geometry("1000x800")  # Increased size to accommodate controls
        self.protocol("WM_DELETE_WINDOW", self._on_close)

        self.fig: Optional[plt.Figure] = None
        self.ax: Optional[plt.Axes] = None
        self.surface_plot_object: Optional[Any] = (
            None  # Stores the Poly3DCollection object for updates
        )
        self.wireframe_plot_object: Optional[Any] = (
            None  # Stores the LineCollection object for wireframe updates
        )
        self.canvas_widget: Optional[tk.Widget] = None
        self.toolbar: Optional[NavigationToolbar2Tk] = None
        self.canvas_agg: Optional[FigureCanvasTkAgg] = None

        # Store latest data received from main window
        self._latest_x_data: Optional[np.ndarray] = None
        self._latest_y_data: Optional[np.ndarray] = None
        self._latest_z_data: Optional[np.ndarray] = None

        # Store current stride values for updates
        self._current_rstride: int = 1
        self._current_cstride: int = 1
        self._current_strided_z_shape: Optional[Tuple[int, ...]] = (
            None  # Store the shape of the strided Z-data for comparison
        )

        # Live update loop management
        self._update_job: Optional[str] = None
        self._update_interval_ms: int = 200  # Update 3D plot every 200ms (5 FPS)

        self._create_widgets()

    def _create_widgets(self) -> None:
        # Main container frame
        self.main_container = ttk.Frame(self)
        self.main_container.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # Control Panel (left side)
        self.control_panel = ttk.LabelFrame(
            self.main_container, text=self._("plot_controls_title")
        )
        self.control_panel.pack(side=tk.LEFT, fill=tk.Y, padx=5, pady=5)

        # Plot Frame (right side)
        self.plot_frame = ttk.Frame(self.main_container)
        self.plot_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=5, pady=5)

        self.message_label = ttk.Label(
            self.plot_frame, text=self._("no_3d_plot_yet"), anchor="center"
        )
        self.message_label.pack(expand=True)

        self._setup_controls()

    def _setup_controls(self) -> None:
        # Colormap selection
        ttk.Label(self.control_panel, text=self._("colormap_label")).pack(
            pady=(10, 0), padx=5, anchor=tk.W
        )
        self.colormap_var = tk.StringVar(value="viridis")
        self.colormap_options = [
            "viridis",
            "plasma",
            "inferno",
            "magma",
            "cividis",
            "jet",
            "coolwarm",
            "gray",
        ]
        self.colormap_menu = ttk.Combobox(
            self.control_panel,
            textvariable=self.colormap_var,
            values=self.colormap_options,
            state="readonly",
        )
        self.colormap_menu.pack(pady=2, padx=5, fill=tk.X)
        self.colormap_menu.bind(
            "<<ComboboxSelected>>", lambda e: self._redraw_plot_with_settings()
        )

        # Shading checkbox
        self.shade_var = tk.BooleanVar(value=True)
        self.shade_check = ttk.Checkbutton(
            self.control_panel,
            text=self._("shading_label"),
            variable=self.shade_var,
            command=self._redraw_plot_with_settings,
        )
        self.shade_check.pack(pady=5, padx=5, anchor=tk.W)

        # Wireframe checkbox
        self.wireframe_var = tk.BooleanVar(value=False)
        self.wireframe_check = ttk.Checkbutton(
            self.control_panel,
            text=self._("wireframe_label"),
            variable=self.wireframe_var,
            command=self._redraw_plot_with_settings,
        )
        self.wireframe_check.pack(pady=5, padx=5, anchor=tk.W)

        # Resolution control (now controls rstride/cstride)
        ttk.Label(self.control_panel, text=self._("resolution_label")).pack(
            pady=(10, 0), padx=5, anchor=tk.W
        )
        # Changed default resolution to a higher value to prevent memory issues
        # with full resolution plots
        self.resolution_var = tk.IntVar(
            value=10
        )  # 1 means full resolution, higher means coarser mesh
        self.resolution_slider = ttk.Scale(
            self.control_panel,
            from_=1,
            to=20,
            orient=tk.HORIZONTAL,
            variable=self.resolution_var,
            command=self._on_resolution_change,
        )
        self.resolution_slider.pack(pady=2, padx=5, fill=tk.X)
        self.resolution_label = ttk.Label(
            self.control_panel, textvariable=self.resolution_var
        )
        self.resolution_label.pack(pady=2, padx=5, anchor=tk.W)

    def _on_resolution_change(self, value: str) -> None:
        """Callback for resolution slider change."""
        # This will trigger a full redraw to apply new rstride/cstride
        self._redraw_plot_with_settings()

    def set_latest_data(
        self, x_data: np.ndarray, y_data: np.ndarray, z_data: np.ndarray
    ) -> None:
        """Receives the latest data from the main window."""
        self._latest_x_data = x_data
        self._latest_y_data = y_data
        self._latest_z_data = z_data

    def start_live_update(self) -> None:
        """Starts the internal update loop for the 3D plot."""
        if self._update_job is None:
            self._update_plot_loop()

    def stop_live_update(self) -> None:
        """Stops the internal update loop for the 3D plot."""
        if self._update_job is not None:
            self.after_cancel(self._update_job)
            self._update_job = None

    def _update_plot_loop(self) -> None:
        """Internal loop to update the 3D plot."""
        if not self.winfo_exists():  # Check if the Toplevel window still exists
            self.stop_live_update()
            return

        if (
            self._latest_x_data is not None
            and self._latest_y_data is not None
            and self._latest_z_data is not None
        ):
            self.update_live_surface(
                self._latest_x_data, self._latest_y_data, self._latest_z_data
            )
        self._update_job = self.after(self._update_interval_ms, self._update_plot_loop)

    def initialize_live_surface(
        self,
        x_data: np.ndarray,
        y_data: np.ndarray,
        z_data: np.ndarray,
        title: str = "3D Surface Plot",
        xlabel: str = "X",
        ylabel: str = "Y",
        zlabel: str = "Z",
    ) -> None:
        """
        Initializes the 3D surface plot. This should be called once when data
        first becomes available.
        """
        if not self.winfo_exists():  # Check if the Toplevel window still exists
            return

        self._clear_plot()  # Clear any previous plot

        rstride = self.resolution_var.get()
        cstride = self.resolution_var.get()
        self._current_rstride = rstride  # Store for updates
        self._current_cstride = cstride  # Store for updates

        self.fig, self.ax, self.surface_plot_object, self.wireframe_plot_object = (
            create_initial_surface_plot(
                x_data,
                y_data,
                z_data,
                title,
                xlabel,
                ylabel,
                zlabel,
                cmap=self.colormap_var.get(),
                shade=self.shade_var.get(),
                wireframe=self.wireframe_var.get(),  # Pass wireframe state
                rstride=rstride,
                cstride=cstride,  # Pass strides
            )
        )

        # Calculate and store the shape of the strided Z-data for the newly
        # created plot
        strided_z_data = z_data[::rstride, ::cstride]
        self._current_strided_z_shape = strided_z_data.shape

        self._draw_plot()

    def update_live_surface(
        self, x_data: np.ndarray, y_data: np.ndarray, z_data: np.ndarray
    ) -> None:
        """
        Updates the data of the live 3D surface plot.
        This function now uses a targeted redraw approach for Z-value changes
        and a full re-initialization for structural changes.
        """
        if not self.winfo_exists():  # Check if the Toplevel window still exists
            return

        # If plot not initialized, or if the data shape has changed,
        # re-initialize fully
        rstride = self.resolution_var.get()
        cstride = self.resolution_var.get()
        incoming_strided_z_shape = z_data[::rstride, ::cstride].shape

        if (
            self.surface_plot_object is None
            or self._current_strided_z_shape != incoming_strided_z_shape
        ):
            print(
                f"DEBUG: Shape mismatch or uninitialized plot. Expected "
                f"{self._current_strided_z_shape}, got {incoming_strided_z_shape}. "
                "Triggering full redraw."
            )
            self._redraw_plot_with_settings()
            return  # Exit to avoid further processing with mismatched data

        # If shapes match, update the existing plot by removing old artists
        # and adding new ones
        if self.ax is not None:
            self.surface_plot_object, self.wireframe_plot_object = (
                update_surface_plot_data(
                    self.ax,
                    x_data,
                    y_data,
                    z_data,
                    self.surface_plot_object,
                    self.wireframe_plot_object,
                    cmap=self.colormap_var.get(),
                    shade=self.shade_var.get(),
                    wireframe=self.wireframe_var.get(),
                    rstride=rstride,
                    cstride=cstride,
                )
            )

            # Update axis limits if data range changes significantly
            current_z_min, current_z_max = self.ax.get_zlim()
            new_z_min, new_z_max = z_data.min(), z_data.max()
            if not (
                current_z_min <= new_z_min <= current_z_max
                and current_z_min <= new_z_max <= current_z_max
            ):
                self.ax.set_zlim(
                    new_z_min - (new_z_max - new_z_min) * 0.1,
                    new_z_max + (new_z_max - new_z_min) * 0.1,
                )

            # Force a direct draw of the canvas
            if self.fig and self.fig.canvas:
                self.fig.canvas.draw()
                self.fig.canvas.flush_events()  # Process events to keep GUI responsive

    def _redraw_plot_with_settings(self) -> None:
        """
        Redraws the entire plot if settings like colormap, shading, or wireframe change,
        or if the Z-data's shape has changed.
        """
        if not self.winfo_exists():  # Check if the Toplevel window still exists
            return

        if (
            self._latest_x_data is not None
            and self._latest_y_data is not None
            and self._latest_z_data is not None
        ):
            self.initialize_live_surface(
                self._latest_x_data,
                self._latest_y_data,
                self._latest_z_data,
                title=self._("3d_surface_plot_title"),
                xlabel="X",
                ylabel="Y",
                zlabel="Height (Z)",
            )
        else:
            self._clear_plot()
            self.message_label.pack(expand=True)

    def _draw_plot(self) -> None:
        """Draws the matplotlib figure on the Tkinter canvas."""
        if (
            not self.winfo_exists() or not self.plot_frame.winfo_exists()
        ):  # Check if the Toplevel window and plot_frame still exist
            print("DEBUG: _draw_plot - Window or plot_frame does not exist. Aborting.")
            return

        if self.fig:
            self.message_label.pack_forget()  # Hide message if plot is drawn

            # _clear_plot should have already destroyed these, but as a safeguard:
            if self.canvas_widget:
                self.canvas_widget.pack_forget()  # Ensure it's removed from layout
                self.canvas_widget.destroy()
                self.canvas_widget = None
            if self.toolbar:
                self.toolbar.pack_forget()  # Ensure it's removed from layout
                self.toolbar.destroy()
                self.toolbar = None

            self.canvas_agg = FigureCanvasTkAgg(self.fig, master=self.plot_frame)
            self.canvas_widget = self.canvas_agg.get_tk_widget()
            self.canvas_widget.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

            self.toolbar = NavigationToolbar2Tk(self.canvas_agg, self.plot_frame)
            self.toolbar.update()
            self.canvas_agg.draw()
        else:
            self.message_label.pack(expand=True)

    def clear_plot_data(self) -> None:
        """Clears the current plot and resets stored data."""
        self.stop_live_update()  # Stop the update loop
        self._clear_plot()
        self._latest_x_data = None
        self._latest_y_data = None
        self._latest_z_data = None
        self.message_label.pack(expand=True)  # Show message again

    def _clear_plot(self) -> None:
        """Clears the current plot from the window."""
        if self.canvas_widget:
            self.canvas_widget.pack_forget()  # Explicitly remove from layout
            self.canvas_widget.destroy()
            self.canvas_widget = None
        if self.toolbar:
            self.toolbar.pack_forget()  # Explicitly remove from layout
            self.toolbar.destroy()
            self.toolbar = None

        if self.fig:
            plt.close(self.fig)
            self.fig = None
            self.ax = None
            self.surface_plot_object = None
            self.wireframe_plot_object = None  # Clear wireframe object
            self._current_strided_z_shape = None  # Clear stored shape
            # self.background = None  # Blitting is removed, so no background to clear

    def _on_close(self) -> None:
        """Handles window closing event."""
        self.stop_live_update()  # Ensure update loop is stopped
        self._clear_plot()
        self.destroy()


if __name__ == "__main__":
    root = tk.Tk()
    root.withdraw()  # Hide the root window

    plot_window = Plot3DWindow(root)
    plot_window.start_live_update()  # Start the update loop

    # Example: Initial 3D Surface Plot
    def f(x: np.ndarray, y: np.ndarray) -> np.ndarray:
        return np.sin(np.sqrt(x**2 + y**2))

    x = np.linspace(-5, 5, 100)  # Full resolution data
    y = np.linspace(-5, 5, 100)
    X, Y = np.meshgrid(x, y)
    Z = f(X, Y)
    plot_window.set_latest_data(X, Y, Z)  # Set initial data

    # Simulate live update by continuously updating data
    def simulate_data_update() -> None:
        new_Z = (
            f(X, Y + time.time() * 0.1) + np.sin(time.time() * 0.5) * 0.5
        )  # Change Z data over time (full resolution)
        plot_window.set_latest_data(X, Y, new_Z)
        root.after(10, simulate_data_update)  # Update data frequently

    root.after(10, simulate_data_update)

    root.mainloop()
