"""
Internal matplotlib hook implementation.

This module contains the machinery for hooking into matplotlib functions
to capture plotting data. Users should not need to interact with this directly.
"""

import matplotlib.pyplot as plt
import matplotlib.axes
import functools
import inspect
from typing import Callable

# Optional xarray import
try:
    import xarray as xr
    HAS_XARRAY = True
except ImportError:
    HAS_XARRAY = False
    xr = None


# List of functions and methods to hook
HOOKED_PYPLOT_FUNCTIONS = [
    "plot",
    "scatter",
    "bar",
    "hist",
    "imshow",
    "contour",
    "contourf",
    "pcolormesh",
    "fill_between",
    "errorbar",
    "boxplot",
    "violin",
    "pie",
    "polar",
    "loglog",
    "semilogx",
    "semilogy",
]

HOOKED_AXES_METHODS = [
    "plot",
    "scatter",
    "bar",
    "hist",
    "imshow",
    "contour",
    "contourf",
    "pcolormesh",
    "fill_between",
    "errorbar",
    "boxplot",
    "violin",
]

# Global state
HOOKS_INSTALLED = False
ACTIVE_FIGUREDATA = None


def set_active_figuredata(figuredata):
    """Set the currently active FigureData instance."""
    global ACTIVE_FIGUREDATA
    ACTIVE_FIGUREDATA = figuredata


def get_active_figuredata():
    """Get the currently active FigureData instance."""
    return ACTIVE_FIGUREDATA


def store_call_data(
    *args, name: str, is_method: bool, convert_xarray_to_numpy: bool = True
) -> None:
    """Store function call data in the active FigureData."""
    if ACTIVE_FIGUREDATA is None:
        return

    if name in ACTIVE_FIGUREDATA.figure_data:
        raise ValueError(f"Artist name '{name}' is already present in this figure's data")

    # If this is a method, drop the first arg (self)
    if is_method:
        args = args[1:]

    # Convert xarray objects to numpy arrays if requested
    if convert_xarray_to_numpy and HAS_XARRAY:
        args = [
            (
                x.values
                if (isinstance(x, xr.DataArray) or isinstance(x, xr.Dataset))
                else x
            )
            for x in args
        ]

    ACTIVE_FIGUREDATA.figure_data[name] = tuple(args)
    ACTIVE_FIGUREDATA.print(f"Stored data for artist {name}")


def create_hooked_fn(original_func: Callable, is_method: bool) -> Callable:
    """
    Create a wrapper around a matplotlib function to intercept and store data.

    Args:
        original_func: The original matplotlib function to wrap
        is_method: Whether this is an Axes method (vs. a pyplot function)

    Returns:
        Wrapped function that captures data when 'name' parameter is provided
    """
    # Get the names of keyword arguments specific to our hook
    hook_kwargs = [
        x for x in inspect.signature(store_call_data).parameters.keys() if x != "args"
    ]

    @functools.wraps(original_func)
    def wrapper(*args, **kwargs):
        # Split kwargs into those for the original function and those for the hook
        original_func_kwargs = {k: v for k, v in kwargs.items() if k not in hook_kwargs}
        store_call_data_kwargs = {k: v for k, v in kwargs.items() if k in hook_kwargs}

        # If 'name' parameter is provided, capture the data
        if "name" in store_call_data_kwargs:
            store_call_data(*args, **store_call_data_kwargs, is_method=is_method)

        # Call the original function (without our custom hook kwargs)
        return original_func(*args, **original_func_kwargs)

    return wrapper


def install_pyplot_hooks() -> None:
    """Install hooks on pyplot functions."""
    for func_name in HOOKED_PYPLOT_FUNCTIONS:
        if hasattr(plt, func_name):
            original_func = getattr(plt, func_name)
            hooked_func = create_hooked_fn(original_func, is_method=False)
            setattr(plt, func_name, hooked_func)


def install_axes_hooks() -> None:
    """Install hooks on Axes methods."""
    for func_name in HOOKED_AXES_METHODS:
        if hasattr(matplotlib.axes.Axes, func_name):
            original_func = getattr(matplotlib.axes.Axes, func_name)
            hooked_func = create_hooked_fn(original_func, is_method=True)
            setattr(matplotlib.axes.Axes, func_name, hooked_func)


def install_all_hooks() -> None:
    """Install hooks on both pyplot functions and Axes methods."""
    global HOOKS_INSTALLED
    if HOOKS_INSTALLED:
        return

    install_pyplot_hooks()
    install_axes_hooks()
    HOOKS_INSTALLED = True


def hooks_installed() -> bool:
    """Check if hooks have been installed."""
    return HOOKS_INSTALLED
