import tkinter as tk
import matplotlib
matplotlib.rcParams['agg.path.chunksize'] = 10000
import os
if not os.environ.get('SPHINX_BUILD'):
    matplotlib.use("TkAgg")
from matplotlib import pyplot as plt
from matplotlib.backends._backend_tk import NavigationToolbar2Tk
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from screeninfo import get_monitors
from PIL import Image, ImageTk
from typing import Any, List, Optional, Tuple, Union

from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.GUI.LAYOUT.ScaleManager import ScaleManager



class MyFigure(tk.Toplevel):
    """
    A custom figure window that can display messages, plots, or images.
    
    This class creates a Toplevel window that can be used to display various types of content:
    - Text messages
    - Matplotlib plots
    - Images
    
    Attributes:
        subplots (Tuple[int, int]): Number of subplots in the figure
        title_plot (List[str]): List of plot titles
        xlabel (List[str]): List of x-axis labels
        ylabel (List[str]): List of y-axis labels
        figure (Figure): The matplotlib figure object
        type_data (str): Type of data being displayed ('Message', 'Plot', or 'Image')
        ax (Union[Axes, List[Axes]]): The matplotlib axes object(s)
    """

    def __init__(
        self,
        title: str,
        message: Optional[str] = None,
        type_data: str = "Message",
        personal_size = None,
        image_path: Optional[str] = None
    ) -> None:
        """
        Initialize the MyFigure window.

        Args:
            title: The window title
            message: Text message to display when type_data is "Message"
            type_data: Type of content to display ("Message", "Plot", or "Image")
            personal_size: Optional tuple (width, height) to determine window size
            image_path: Path to the image file when type_data is "Image"
        """
        super().__init__()
        self.title(title)
        self.protocol("WM_DELETE_WINDOW", self.close_figure)
        
        # Initialize attributes
        self.subplots = None
        self.title_plot = None
        self.xlabel = None
        self.ylabel = None
        self.figure = None
        self.type_data = type_data
        self.ax = None

        if type_data == "Message":
            self.label = tk.Label(
                self,
                text=message,
                anchor="w",
                justify="left",
                font=ScaleManager.get().font_label_normal if ScaleManager.get() else ("Sans", 10)
            )
            self.label.pack()

        elif type_data == "Plot":
            # Setup window size if provided
            if personal_size is not None:
                m = get_monitors()
                width = int(m[0].width * personal_size[0])
                height = int(m[0].height * personal_size[1])
                self.geometry(f"{width}x{height}+0+0")
            else:
                self.geometry(ConstantVariables.geometry)

        elif type_data == "Image":
            if image_path:
                try:
                    if personal_size is None:
                        width = 1000
                        height = 800
                    else:
                        width = int(personal_size[0])
                        height = int(personal_size[1])
                    # Open the image and convert it to a PhotoImage
                    image = Image.open(image_path).resize((width, height))
                    photo = ImageTk.PhotoImage(image)
                    self.image_label = tk.Label(self, image=photo)
                    self.image_label.image = photo  # Keep a reference!
                    self.image_label.pack()
                except Exception as e:
                    self.label = tk.Label(
                        self,
                        text=f"Error loading image: {e}",
                        anchor="w",
                        justify="left",
                        font=ScaleManager.get().font_label_normal if ScaleManager.get() else ("Sans", 10)
                    )
                    self.label.pack()
            else:
                self.label = tk.Label(
                    self,
                    text="No image path provided",
                    anchor="w",
                    justify="left",
                    font=ScaleManager.get().font_label_normal if ScaleManager.get() else ("Sans", 10)
                )
                self.label.pack()

        # Add a close button to the window
        self.close_button = tk.Button(self, text="Close", command=self.close_figure)
        self.close_button.pack()

    def return_ax(
        self,
        title_plot: List[str],
        xlabel: List[str],
        ylabel: List[str],
        subplots = (1, 1),
        figsize: Tuple[int, int] = (8, 6),
        dpi: int = 100,
        hspace: float = 0.1,
        wspace: float = 0.1,
        height_ratios: Optional[List[float]] = None,
        sharex: bool = False,
        sharey: bool = False,
    ) -> Any:
        """
        Create and return matplotlib axes for plotting.

        Args:
            title_plot: List of plot titles
            xlabel: List of x-axis labels
            ylabel: List of y-axis labels
            subplots: Tuple of (rows, columns) for subplot layout
            figsize: Figure size in inches (width, height)
            dpi: Dots per inch for the figure
            hspace: Height spacing between subplots
            wspace: Width spacing between subplots
            height_ratios: List of height ratios for subplots
            sharex: Whether to share x-axis between subplots
            sharey: Whether to share y-axis between subplots

        Returns:
            The matplotlib axes object(s)
        """
        if height_ratios is None:
            height_ratios = [1] * subplots[0]
            
        self.title_plot = title_plot
        self.xlabel = xlabel
        self.ylabel = ylabel
        self.subplots = subplots
        
        self.figure, self.ax = plt.subplots(
            subplots[0],
            subplots[1],
            figsize=figsize,
            dpi=dpi,
            constrained_layout=True,
            gridspec_kw={
                'hspace': hspace,
                'wspace': wspace,
                'height_ratios': height_ratios
            },
            sharex=sharex,
            sharey=sharey,
        )
        return self.ax

    def create_figure(self, path_f: str, dpi: int, tight_layout : bool = True) -> None:
        """
        Create and save the figure with proper labels and formatting.

        Args:
            path_f: Path where to save the figure
            dpi: Dots per inch for the saved figure
            tight_layout: Whether to apply tight layout to the figure
        """
        if len(self.xlabel) == 1:
            self.ax.set_xlabel(self.xlabel[0])
            self.ax.set_ylabel(self.ylabel[0])
            self.ax.set_title(self.title_plot[0])
        elif self.subplots[0] == 1 or self.subplots[1] == 1:
            for i in range(max(self.subplots[0], self.subplots[1])):
                if self.ylabel[i] != "":
                    self.ax[i].set_ylabel(self.ylabel[i])
                if self.xlabel[i] != "":
                    self.ax[i].set_xlabel(self.xlabel[i])
                if self.title_plot[i] != "":
                    self.ax[i].set_title(self.title_plot[i])
        else:
            counter = 0
            for i in range(self.subplots[0]):
                for j in range(self.subplots[1]):
                    if counter < len(self.xlabel):
                        if self.ylabel[counter] != "":
                            self.ax[i, j].set_ylabel(self.ylabel[counter])
                        if self.xlabel[counter] != "":
                            self.ax[i, j].set_xlabel(self.xlabel[counter])
                        if self.title_plot[counter] != "":
                            self.ax[i, j].set_title(self.title_plot[counter])
                        plt.setp(self.ax[i, j].get_xticklabels(), rotation=45, fontsize=10)
                        counter += 1
        if tight_layout:
            plt.tight_layout()
        # Save the figure
        self.figure.savefig(path_f, dpi=dpi)
        if "jpg" in path_f:
            self.figure.savefig(str(path_f).replace("jpg", "pdf"), dpi=dpi)
        elif "png" in path_f:
            self.figure.savefig(str(path_f).replace("png", "pdf"), dpi=dpi)
            
        # Create the canvas and toolbar
        chart_type = FigureCanvasTkAgg(self.figure, self)
        chart_type.draw()
        chart_type.get_tk_widget().pack()
        toolbar = NavigationToolbar2Tk(chart_type, self)
        toolbar.update()
        chart_type.get_tk_widget().pack()

    def close_figure(self) -> None:
        """
        Close the figure window and clean up resources.
        """
        try:
            if self.type_data == "Plot":
                plt.close(self.figure)
        except Exception:
            pass
        self.destroy()

