# TODO: check for official matplotlib type stubs (matplotlib-stubs or inline types).
# As of matplotlib 3.10, stubs are incomplete — many pyplot/axes methods have
# partially unknown signatures in pyright. Re-evaluate when matplotlib ships
# PEP 561 inline types or a maintained stubs package reaches full coverage.
import structlog.stdlib
from matplotlib.axes import Axes
from matplotlib.figure import Figure
from numpy import ndarray
from matplotlib import pyplot as plt

from samgis_core.utilities.type_hints import ListStr, MatplotlibBackend

FigAxes = tuple[Figure, Axes]


logger = structlog.stdlib.get_logger(__name__)


def helper_imshow_output_expected(
    img_list: list[ndarray],
    titles_list: ListStr,
    cmap: str = "gist_rainbow",
    plot_size: int = 5,
    show: bool = False,
    debug: bool = False,
    close_after: float = 0.0,
) -> FigAxes:
    """
    Simple way to display a list of images with their titles, color map.
    Should work also in an automate environments, like tests (use a `close_after` argument > 0)

    Args:
        img_list: ndarray images to display
        titles_list: title images
        cmap: color map
        plot_size: figure plot size
        show: fire plt.show() action if needed
        debug: workaround useful in an interactive context, like Pycharm debugger
        close_after: close after give seconds (useful in tests, contrasted to 'debug' option)

    Returns:
        tuple of matplotlib Figure, Axes

    """
    n_img = len(img_list)
    n_titles = len(titles_list)
    if n_titles != n_img:
        raise ValueError(
            f"wrong number of titles ({n_titles}) vs number of img ({n_img})"
        )
    fig, ax = plt.subplot_mosaic([titles_list], figsize=(n_img * plot_size, plot_size))  # type: ignore[arg-type]

    for title, img in zip(titles_list, img_list):
        ax[title].imshow(img, cmap=cmap)
        ax[title].legend(title)
    if show:
        if debug:
            plt.pause(0.01)
            plt.show()
        if close_after > 0:
            plt.pause(close_after)
            plt.show(block=False)
            plt.close("all")
    return fig, ax


def imshow_raster(
    raster: ndarray,
    title: str,
    cmap: str = "gist_rainbow",
    interpolation: str | None = None,
    alpha: float | None = None,
    transform: object | None = None,
    plot_size: int = 5,
    show: bool = False,
    debug: bool = False,
    close_after: float = 0.0,
    backend: MatplotlibBackend | None = None,
) -> FigAxes:
    """
    Displays raster images lists/arrays with titles, legend, alpha transparency, figure sizes
    and geographic transformations, if not none (leveraging rasterio.plot)

    Args:
        raster: image to display
        title: title image
        cmap: color map
        interpolation: interpolation type
        alpha: alpha transparency
        transform: geographic transform, eventually used for map representation by rasterio
        plot_size: figure plot size
        show: fire plt.show() action if needed
        debug: workaround useful in an interactive context, like Pycharm debugger
        close_after: close after give seconds (useful in tests, contrasted to 'debug' option)
        backend: matplotlib backend string

    Returns:
        tuple of matplotlib Figure, Axes

    """
    from rasterio import plot  # type: ignore[import-untyped]

    if not backend:
        backend = MatplotlibBackend(plt.get_backend())
        plt.rcParams["backend"] = backend
    logger.info(f"use {backend} as matplotlib backend...")

    fig, ax = plt.subplots(figsize=(plot_size, plot_size))
    raster_ax = raster[0] if transform is not None else raster
    image_hidden = ax.imshow(
        raster_ax, cmap=cmap, interpolation=interpolation, alpha=alpha
    )
    if transform is not None:
        plot.show(
            raster,
            transform=transform,
            ax=ax,
            cmap=cmap,
            interpolation=interpolation,
            alpha=alpha,
        )
    fig.colorbar(image_hidden, ax=ax)
    ax.set_title(title)
    if show:
        if debug:
            plt.pause(0.01)
            plt.show()
        if close_after > 0:
            plt.pause(close_after)
            plt.show(block=False)
            plt.close("all")
    return fig, ax
