"""
Main driedfigs API for capturing and saving matplotlib figure data.

This module provides the primary user-facing functions for driedfigs.
"""

from typing import Dict, Optional, TypeAlias, Union
from dataclasses import dataclass, field
from pathlib import Path
import pickle as pkl

from ._hooks import install_all_hooks, set_active_figuredata, get_active_figuredata

PathLike: TypeAlias = Union[str, Path]


@dataclass
class FigureData:
    """
    Container for captured matplotlib plotting data.

    Attributes:
        figure_name: Name identifier for this figure
        figure_data: Dictionary mapping artist names to their data
        verbose: Whether to print informational messages
    """

    figure_name: str
    figure_data: Dict = field(default_factory=dict)
    verbose: Optional[bool] = False

    @property
    def is_active(self):
        """Check if this FigureData instance is currently active."""
        return get_active_figuredata() is self

    def save(self, output_path: PathLike) -> None:
        """
        Save the captured figure data to a pickle file.

        Args:
            output_path: Path where the pickle file should be saved
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with output_path.open("wb") as f:
            pkl.dump(self.figure_data, f)

    def print(self, *args, **kwargs):
        """Print message if verbose mode is enabled."""
        if self.verbose:
            print(*args, **kwargs)

    def __repr__(self) -> str:
        return (
            f"FigureData(figure_name={self.figure_name}, "
            f"artist names={list(self.figure_data.keys())})"
        )


def start_figure(figure_name: str, verbose: bool = False) -> FigureData:
    """
    Start capturing data for a new figure.

    This function must be called before creating any plots that you want to capture.
    It automatically installs matplotlib hooks on the first call.

    Args:
        figure_name: Name for this figure (used for output filenames)
        verbose: If True, print informational messages about captured data

    Returns:
        FigureData instance that will collect the plotting data

    Example:
        >>> import driedfigs as df
        >>> import matplotlib.pyplot as plt
        >>>
        >>> fig_data = df.start_figure("my_plot")
        >>> plt.plot([1, 2, 3], [4, 5, 6], name="line1")
        >>> fig_data.save("output/my_plot.pkl")
    """
    # Install hooks on first use
    install_all_hooks()

    # Create and activate the new FigureData
    figuredata = FigureData(figure_name=figure_name, verbose=verbose)
    set_active_figuredata(figuredata)

    return figuredata


def save_figuredata(
    output_dir: PathLike, figuredata: Optional[FigureData] = None
) -> None:
    """
    Save figure data to a pickle file.

    Args:
        output_path: Path where the pickle file should be saved
        figuredata: FigureData instance to save. If None, uses the currently active one.

    Raises:
        ValueError: If no FigureData is active and none was provided

    Example:
        >>> import driedfigs as df
        >>> import matplotlib.pyplot as plt
        >>>
        >>> df.start_figure("my_plot")
        >>> plt.plot([1, 2, 3], [4, 5, 6], name="data")
        >>> df.save_figuredata("output/my_plot.pkl")
    """
    fig_data = figuredata or get_active_figuredata()

    if fig_data is None:
        raise ValueError(
            "No active FigureData. Call start_figure() first or provide a figuredata"
            " argument."
        )

    fig_data.save(output_dir)
