"""Factory functions to create publication-ready figures."""

from typing import Optional, Tuple, Dict, Any, Union, Literal, TYPE_CHECKING

import matplotlib.pyplot as plt
from matplotlib.figure import Figure

from neoplot.style import Style, load_theme
from neoplot.axes.base import PlotAxis

if TYPE_CHECKING:
    from neoplot.axes.timeseries import TimeSeriesAxis
    from neoplot.axes.scatter import ScatterAxis
    from neoplot.axes.map import MapAxis

ThemeName = Literal[
    "publication/paper", "publication/paper_dark",
    "publication/poster", "publication/poster_dark",
    "presentation/slides", "presentation/slides_dark"
]
StyleLike = Union[Style, ThemeName, str]


def _create_figure_with_axis(
    axis_class: type,
    style: Optional[Style] = None,
    figsize: Optional[Tuple[float, float]] = None,
    nrows: int = 1,
    ncols: int = 1,
    **axis_kwargs
) -> Tuple[Figure, Any]:
    """
    Internal helper to create a figure and axis with proper styling.
    
    Args:
        axis_class: The PlotAxis subclass to instantiate.
        style: Style object. If None, uses defaults.
        figsize: Figure size in inches (width, height). If None, uses style default.
        nrows, ncols: Number of subplot rows and columns.
        **axis_kwargs: Additional keyword arguments passed to axis_class.
    
    Returns:
        tuple: (matplotlib Figure, PlotAxis instance)
    """
    if style is None:
        style = Style()
    
    if figsize is None:
        figsize = style.figsize

    # Apply style rcParams before creating figure/axes so defaults propagate.
    plt.rcParams.update(style.to_matplotlib_rc())
    
    # Create figure
    fig = plt.figure(figsize=figsize, dpi=style.dpi)
    
    # Create and setup axis
    plot_axis = axis_class(style=style, **axis_kwargs)
    
    # For a simple grid, this currently just adds one subplot,
    # but could be expanded. For now, matching previous single-axis behavior:
    plot_axis.add_subplot(fig, nrows, ncols, 1)
    
    return fig, plot_axis


def timeseries(
    style: Optional[StyleLike] = None,
    figsize: Optional[Tuple[float, float]] = None,
    **kwargs
) -> Tuple[Figure, "TimeSeriesAxis"]:
    """
    Create a figure with a TimeSeriesAxis (not yet implemented in phase 1).
    
    Args:
        style: Style object or theme name (e.g., 'publication/paper').
                If None, uses default style.
        figsize: Figure size in inches (width, height).
        **kwargs: Additional keyword arguments passed to TimeSeriesAxis.
    
    Returns:
        tuple: (matplotlib Figure, TimeSeriesAxis)
    
    Example:
        >>> from neoplot import timeseries, load_theme
        >>> style = load_theme('publication/paper')
        >>> fig, ax = timeseries(style=style)
        >>> ax.plot(my_data)
        >>> ax.finalize()
    """
    # Handle string theme names
    _style: Style = load_theme(style) if isinstance(style, str) else (style or Style())
    
    # Import here to avoid circular imports
    from neoplot.axes.timeseries import TimeSeriesAxis
    
    fig, plot_axis = _create_figure_with_axis(TimeSeriesAxis, _style, figsize, **kwargs)
    from typing import cast
    return fig, cast("TimeSeriesAxis", plot_axis)


def scatter(
    style: Optional[StyleLike] = None,
    figsize: Optional[Tuple[float, float]] = None,
    **kwargs
) -> Tuple[Figure, "ScatterAxis"]:
    """
    Create a figure with a ScatterAxis (not yet implemented in phase 1).
    
    Args:
        style: Style object or theme name (e.g., 'presentation/slides').
                If None, uses default style.
        figsize: Figure size in inches (width, height).
        **kwargs: Additional keyword arguments passed to ScatterAxis.
    
    Returns:
        tuple: (matplotlib Figure, ScatterAxis)
    """
    # Handle string theme names
    _style: Style = load_theme(style) if isinstance(style, str) else (style or Style())
    
    # Import here to avoid circular imports
    from neoplot.axes.scatter import ScatterAxis
    
    fig, plot_axis = _create_figure_with_axis(ScatterAxis, _style, figsize, **kwargs)
    from typing import cast
    return fig, cast("ScatterAxis", plot_axis)


def map(
    style: Optional[StyleLike] = None,
    projection: str = "Mercator",
    projection_kwargs: Optional[Dict[str, Any]] = None,
    extent: Optional[Tuple[float, float, float, float]] = None,
    extent_crs: Optional[Any] = None,
    default_data_crs: Optional[Any] = None,
    gridline_left_labels: bool = True,
    gridline_right_labels: bool = False,
    gridline_top_labels: bool = False,
    gridline_bottom_labels: bool = True,
    gridline_linewidth: Optional[float] = None,
    gridline_linestyle: Optional[Any] = None,
    coastline_resolution: str = "50m",
    figsize: Optional[Tuple[float, float]] = None,
    **kwargs
) -> Tuple[Figure, "MapAxis"]:
    """
    Create a figure with a MapAxis (Cartopy integration).
    
    Args:
        style: Style object or theme name (e.g., 'publication/paper').
                If None, uses default style.
        projection: Cartographic projection name (e.g., 'Mercator', 'RotatedPole', 'Orthographic').
        projection_kwargs: Optional dict of keyword arguments for projection constructor
                          (e.g., {'pole_latitude': 90.0, 'pole_longitude': 0.0} for RotatedPole).
        extent: Optional map extent as (lon_min, lon_max, lat_min, lat_max).
                If provided, extent_crs must also be provided.
        extent_crs: Cartopy CRS instance for interpreting extent coordinates.
                   Required if extent is provided.
        default_data_crs: Default Cartopy CRS for data in ax.plot() if transform not specified.
                        If None, defaults to PlateCarree().
        gridline_left_labels: Whether to show latitude labels on left side (default True).
        gridline_right_labels: Whether to show latitude labels on right side (default False).
        gridline_top_labels: Whether to show longitude labels on top side (default False).
        gridline_bottom_labels: Whether to show longitude labels on bottom side (default True).
        gridline_linewidth: Optional map gridline width override.
        gridline_linestyle: Optional map gridline style override.
        coastline_resolution: Natural Earth coastline/border resolution
                 ("110m", "50m", "10m"), default "50m".
        figsize: Figure size in inches (width, height).
        **kwargs: Additional keyword arguments passed to MapAxis.
    
    Returns:
        tuple: (matplotlib Figure, MapAxis)
    
    Example:
        >>> from neoplot import map, load_theme
        >>> style = load_theme('publication/paper')
        >>> fig, ax = map(style=style, projection='Mercator')
        >>> ax.plot(coastline_data)
        >>> ax.finalize()
    
    Rotated-pole example (EU CORDEX 3km):
        Note: When using RotatedPole with PlateCarree extent, the factory 
        automatically transforms the extent to rotated-pole coordinates.
        
        >>> import cartopy.crs as ccrs
        >>> fig, ax = map(
        ...     projection='RotatedPole',
        ...     projection_kwargs={'pole_latitude': 39.25, 'pole_longitude': -162.0},
        ...     extent=[-8.9, 40.1, 24.0, 67.6],  # In PlateCarree
        ...     extent_crs=ccrs.PlateCarree(),  # Automatically transformed!
        ...     default_data_crs=ccrs.RotatedPole(pole_latitude=39.25, pole_longitude=-162.0)
        ... )
    """
    # Handle string theme names
    _style: Style = load_theme(style) if isinstance(style, str) else (style or Style())

    if figsize is None:
        figsize = _style.figsize

    try:
        import cartopy.crs as ccrs
    except ImportError as exc:
        raise ImportError(
            "Cartopy is required for mapping. Install with: pip install cartopy"
        ) from exc

    projection_class = getattr(ccrs, projection, None)
    if projection_class is None:
        raise ValueError(
            f"Unknown projection: {projection}. Must be a valid cartopy projection."
        )

    # Build projection instance with kwargs if provided
    if projection_kwargs is None:
        projection_kwargs = {}

    # Validate extent and extent_crs are both provided or both absent
    if (extent is None) != (extent_crs is None):
        raise ValueError(
            "Both 'extent' and 'extent_crs' must be provided together, or both None."
        )

    # Apply style before figure creation so defaults propagate.
    plt.rcParams.update(_style.to_matplotlib_rc())

    fig = plt.figure(figsize=figsize, dpi=_style.dpi)

    # Import here to avoid circular imports
    from neoplot.axes.map import MapAxis

    plot_axis = MapAxis(
        style=_style,
        projection=projection,
        projection_kwargs=projection_kwargs,
        extent=extent,
        extent_crs=extent_crs,
        default_data_crs=default_data_crs,
        gridline_left_labels=gridline_left_labels,
        gridline_right_labels=gridline_right_labels,
        gridline_top_labels=gridline_top_labels,
        gridline_bottom_labels=gridline_bottom_labels,
        gridline_linewidth=gridline_linewidth,
        gridline_linestyle=gridline_linestyle,
        coastline_resolution=coastline_resolution,
        **kwargs
    )
    plot_axis.add_subplot(fig, 1, 1, 1)

    return fig, plot_axis


__all__ = ["timeseries", "scatter", "map"]
