import typing
from typing import Optional, Tuple

import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from matplotlib.figure import Figure, SubFigure

from zyplib.utils.kv import filter_kwargs

# 类型别名，方便代码阅读
FigureLike = Figure | SubFigure
FigAx = Tuple[FigureLike, Axes]

__all__ = ['_prepare_ax', '_finalize_figure']


AUTO_SHOW_PLOT = False


def _prepare_ax(ax: Optional[Axes] = None, **kwargs) -> FigAx:
    """内部函数，用于处理 ax 不存在的情况并创建新图表"""

    subplots_kwargs = filter_kwargs(
        kwargs,
        allowed=[
            'figsize',
            'nrows',
            'ncols',
            'sharex',
            'sharey',
            'squeeze',
            'width_ratios',
            'height_ratios',
            'subplot_kw',
            'gridspec_kw',
        ],
    )
    if ax is None:
        # figsize = kwargs.pop('figsize', None)
        fig, ax = plt.subplots(**subplots_kwargs)
        return typing.cast(FigAx, (fig, ax))
    else:
        fig = ax.figure
    return typing.cast(FigAx, (fig, ax))


def _finalize_figure(
    fig: FigureLike, show: bool, save_path: Optional[str], auto_close: Optional[bool]
):
    """内部函数，处理图表的显示、保存和关闭逻辑"""
    root_figure = fig if isinstance(fig, Figure) else fig.figure

    if save_path:
        import os

        dir_name = os.path.dirname(save_path)
        if dir_name:
            os.makedirs(dir_name, exist_ok=True)
        root_figure.savefig(save_path, bbox_inches='tight')

    if show:
        root_figure.tight_layout()
        plt.show()

    should_auto_close = auto_close
    if auto_close is None:
        # should_auto_close = (not show) and (save_path is None)
        # 如果既没有 show 也没有保存，大概率是传入的 ax 需要给后续继续编辑调整用，所以不能关闭
        should_auto_close = bool(show or save_path)

    if should_auto_close:
        plt.close(root_figure)
