"""
Figure-level helpers shared by every renderer: saving or showing, and the
interactive cursor readout.
"""

import os

import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs

from .coords import is_lat, is_lon, centers_to_edges

# Extensions that mean "write a movie" rather than "write a still". Shared with
# the 3D globe, which writes its own videos through VTK.
VIDEO_SUFFIXES = ('.mp4', '.gif', '.avi', '.mov', '.mkv', '.webm')


def is_video(output_path):
    """
    Whether an output path names a movie rather than a still image.
    """
    return bool(output_path) and output_path.lower().endswith(VIDEO_SUFFIXES)


def still_path(output_path):
    """
    A path matplotlib can actually write as a single image.

    `--animate ... -o movie.mp4` names a movie. The renderers that can produce
    one write it themselves; a figure that cannot be animated - a section, a
    profile - takes the same stem as a PNG rather than failing on an extension
    matplotlib has never supported.
    """
    if not output_path:
        return output_path
    stem, ext = os.path.splitext(output_path)
    return f"{stem}.png" if ext.lower() in VIDEO_SUFFIXES else output_path


def axes_for(ctx, figsize, **subplot_kw):
    """
    The figure and axes one renderer should draw into.

    Every plot kind keeps its own default size - a profile is portrait and a map
    is landscape for a reason, and `--figsize` overrides that rather than
    flattening the lot to one shape.

    This is also the seam a multi-panel layout needs: a renderer that asks for
    its axes instead of calling `plt.subplots` itself can be pointed at one cell
    of a grid without knowing it is in one.
    """
    panel = getattr(ctx, 'panel', None)
    if panel is not None:
        return panel.next_axes(**subplot_kw)
    return plt.subplots(figsize=getattr(ctx, 'figsize', None) or figsize,
                        subplot_kw=subplot_kw or None)


def ensure_directory(path):
    """
    Make the directory an output path names. True when it is there afterwards.

    `-o runs/2024/tsurf.png` inside a shell loop is a reasonable thing to write,
    and it used to end in a FileNotFoundError traceback where every other
    refusal in this tool is one printed line.
    """
    folder = os.path.dirname(os.path.abspath(path))
    try:
        os.makedirs(folder, exist_ok=True)
    except OSError as err:
        print(f"Error: cannot create '{folder}': {err}")
        return False
    return True


def save_figure(fig, path, dpi=None, **kwargs):
    """
    Write one figure, closing it afterwards. Returns a process exit status.

    Everything that can go wrong here - a directory that is not there, a
    read-only one, an extension matplotlib has never supported - used to come
    out as a traceback. The tool's own convention is a sentence and a status,
    so that a sweep over many variables reports the one that failed and carries
    on.
    """
    annotate_blank_cells(fig)
    if not ensure_directory(path):
        return 1
    try:
        # A map is saved whole rather than trimmed. `bbox_inches='tight'` asks
        # each axes for its own extent, and a cartopy GeoAxes answers with the
        # map alone - not its title, and not anything drawn beside it. Trimming
        # to that answer threw the title away entirely, and, where a colour bar
        # had taken its room out of the axes, cropped an eight-inch figure down
        # to seventy-four pixels of colour bar with no map at all. The figure is
        # already sized for what it holds, so there is nothing to trim.
        if not _has_geoaxes(fig):
            kwargs.setdefault('bbox_inches', 'tight')
        if dpi:
            kwargs['dpi'] = dpi
        fig.savefig(path, **kwargs)
    except (OSError, ValueError) as err:
        print(f"Error: cannot write '{path}': {err}")
        return 1
    finally:
        plt.close(fig)
    print(f"Saved to {path}")
    return 0


def finish_figure(fig, output_path, dpi=None):
    """
    Save a figure when an output path is given, otherwise display it.
    Returns a process exit status.

    `dpi` applies to the file only. Setting it on the figure would resize the
    window too, so a --dpi meant for a printed page would shrink or blow up the
    one on screen.

    A saved figure is closed; a shown one is not, because that one is the point.
    The interactive loop keeps every figure it has drawn, so that two variables
    can be compared side by side, and closes them on request instead.
    """
    # A figure still being filled in - one panel of several - is neither saved
    # nor shown, because the renderer drawing this panel cannot know there are
    # more coming. The grid clears the mark as it hands out its last cell.
    if getattr(fig, '_dispnc_incomplete', False):
        return 0

    if output_path:
        return save_figure(fig, still_path(output_path), dpi=dpi)
    display_figure(fig)
    return 0


def display_figure(fig):
    """
    Put one finished figure on screen.

    `plt.show()` raises every open window, which is right when there is one and
    wrong once the interactive loop has a screenful: plotting a fifth variable
    would shuffle the four already arranged for comparison. Under `plt.ion()`
    the figure shows itself instead and the rest stay where the user put them.
    """
    annotate_blank_cells(fig)
    if plt.isinteractive():
        name_window(fig)
        fig.show()
        fig.canvas.flush_events()
    else:
        plt.show()


def name_window(fig, title=None):
    """
    Put the figure's own title on its window, so a screen of them can be told
    apart in a task switcher or a window list.

    Falls back to whatever the figure already says about itself, and does
    nothing at all under a backend with no windows to name.
    """
    manager = getattr(getattr(fig, 'canvas', None), 'manager', None)
    if manager is None:
        return
    if title is None:
        suptitle = getattr(fig, '_suptitle', None)
        title = (suptitle.get_text() if suptitle is not None else '') or \
            next((ax.get_title() for ax in fig.axes if ax.get_title()), '')
    if not title:
        return
    try:
        manager.set_window_title(' '.join(title.split())[:80])
    except Exception:
        pass


def annotate_blank_cells(fig):
    """
    Write the note a log scale's greyed cells came with under the picture.

    `colors.log_colormap` hangs the sentence off the colormap rather than
    passing it down through every renderer's signature, and this reads it back
    off whatever was drawn with it. Called from both exits - saving and showing
    - so a figure that never goes through `finish_figure`, as the polar pair
    does not, is annotated too; the mark stops the nested save from writing it
    twice.
    """
    if getattr(fig, '_dispnc_blank_note_drawn', False):
        return
    for ax in fig.axes:
        for artist in list(ax.collections) + list(ax.images):
            note = getattr(getattr(artist, 'cmap', None), '_dispnc_blank_note', None)
            if note:
                fig.text(0.5, 0.01, note, ha='center', fontsize=8, color='0.35')
                fig._dispnc_blank_note_drawn = True
                return


def _has_geoaxes(fig):
    """
    Whether any axes on the figure is a cartopy GeoAxes.

    Asked by duck-typing rather than by import, so a figure of ordinary axes
    never pays for cartopy being installed.
    """
    return any(hasattr(ax, 'projection') for ax in fig.axes)


def attach_format_coord(ax, mat, x, y, x_dim, y_dim, varname, data_crs=ccrs.PlateCarree()):
    """
    Attach a format_coord function to the axes to display x, y, and value at cursor.
    Assumes a pcolormesh-style grid whose x/y are 1D cell-center coordinates.
    """
    # Determine dimensions
    if mat.ndim == 2:
        ny, nx = mat.shape
    elif mat.ndim == 3 and mat.shape[2] in (3, 4):
        ny, nx, nc = mat.shape
    else:
        raise ValueError(f"Unsupported mat shape {mat.shape}")

    xarr = np.asarray(x, dtype=float).ravel()
    yarr = np.asarray(y, dtype=float).ravel()

    xedges, x_ascending = centers_to_edges(xarr)
    yedges, y_ascending = centers_to_edges(yarr)

    # Detect if ax is a GeoAxes with a projection we can invert
    proj = getattr(ax, 'projection', None)
    x_is_lon, y_is_lon = is_lon(x_dim), is_lon(y_dim)
    x_is_lat, y_is_lat = is_lat(x_dim), is_lat(y_dim)
    geo_axes = isinstance(proj, ccrs.Projection) and ((x_is_lon and y_is_lat) or (x_is_lat and y_is_lon))

    def format_coord(xp, yp):
        # Geographic transform if appropriate
        if geo_axes:
            try:
                lonp, latp = data_crs.transform_point(xp, yp, src_crs=proj)
            except Exception:
                lonp, latp = xp, yp
            if x_is_lon and y_is_lat:
                xi, yi = lonp, latp
            else:
                xi, yi = latp, lonp
        else:
            xi, yi = xp, yp

        # Map to matrix indices
        if x_ascending:
            col = np.searchsorted(xedges, xi, side='right') - 1
        else:
            col = (nx - 1) - (np.searchsorted(xedges, xi, side='right') - 1)

        if y_ascending:
            row = np.searchsorted(yedges, yi, side='right') - 1
        else:
            row = (ny - 1) - (np.searchsorted(yedges, yi, side='right') - 1)

        # Build the label
        label_xy = f"{x_dim}={xi:.3g}, {y_dim}={yi:.3g}"
        if 0 <= row < ny and 0 <= col < nx:
            if mat.ndim == 2:
                v = mat[row, col]
                return f"{label_xy}, {varname}={v:.3g}"
            else:
                vals = mat[row, col]
                txt = ", ".join(f"{vv:.3g}" for vv in vals[:3])
                return f"{label_xy}, {varname}=({txt})"
        else:
            return label_xy

    ax.format_coord = format_coord