"""Style system for neoplot: encapsulates colors, fonts, DPI, and output formats."""

import copy
import json
from pathlib import Path
from typing import Any, Dict, Optional, Union

import yaml


class Style:
    """
    Encapsulates plot styling: colors, fonts, line styles, DPI, output formats.
    
    Styles can be created programmatically or loaded from YAML theme files.
    Accepts all properties directly in the constructor — method-chaining via
    ``with_*`` methods is still available but no longer required.

    Two independent font slots are supported:

    * **body font** — used for axis labels, tick labels, and legend text.
      Controlled by ``font_family`` / ``font`` (short alias).
      Default: ``"Montserrat"``.
    * **title font** — used only for the axis/figure title.
      Controlled by ``title_font_family`` / ``title_font`` (short alias).
      Default: ``"Karma"``.

    Both fonts are resolved from the TTF files bundled in ``neoplot/fonts/``
    and registered automatically on package import.

    Attributes:
        name (str): Style identifier (e.g., 'publication/paper').
        colors (dict): Named color palette. Keys are semantic names ('primary', 'secondary', etc.).
        color_overrides (dict, optional): Per-key colour overrides merged on top of ``colors``.
        font_family (str): Body font family (labels, ticks, legend). Default ``"Montserrat"``.
        font (str, optional): Short alias for ``font_family``.
        title_font_family (str): Title font family. Default ``"Karma"``.
        title_font (str, optional): Short alias for ``title_font_family``.
        font_size_base (float): Base font size in points.
        font_size (float, optional): Short alias for ``font_size_base``.
        axis_label_size (float, optional): Size for x/y axis labels. Defaults
            to ``font_size_base + 1``.
        tick_label_size (float, optional): Size for tick labels. Defaults to
            ``font_size_base - 1``.
        line_width_base (float): Base line width in points.
        dpi (int): Output DPI (e.g., 300 for print, 96 for screen).
        figsize (tuple): Default figure size in inches (width, height).
        context (str): Context for this style ('publication' or 'presentation').
        grid_visible (bool): Whether to show gridlines by default.
        grid_alpha (float): Grid line transparency (0-1).
        grid_linestyle (str): Grid line style ('-', '--', '-.', ':', etc.). Default '--' (dashed).
        grid_linewidth (float): Grid line width. Default 0.8.
        zero_line_visible (bool): Whether to draw a horizontal reference line at
            y=0. Default True.
        zero_line_color (str): Color of the y=0 reference line. Default black.
        zero_line_linestyle (str): Line style for the y=0 line. Default '--'.
        zero_line_linewidth (float): Width of y=0 line. Default 1.0.
        zero_line_alpha (float): Opacity of the y=0 line. Default 0.5.
        axis_baselines_visible (bool): Whether to draw custom left/bottom frame
            lines when spines are hidden. Default True.
        axis_baselines_color (str): Color for custom left/bottom frame lines.
            Default black.
        axis_baselines_linestyle (str): Line style for custom frame lines.
            Default '--'.
        axis_baselines_linewidth (float): Line width for custom frame lines.
            Default 1.0.
        axis_baselines_alpha (float): Opacity for custom frame lines. Default 0.5.
        spines_visible (bool): Whether to show axis spines (borders). Default False.
    """
    
    def __init__(
        self,
        name: str = "default",
        colors: Optional[Dict[str, str]] = None,
        color_overrides: Optional[Dict[str, str]] = None,
        font_family: str = "Montserrat",
        font: Optional[str] = None,
        title_font_family: str = "Karma",
        title_font: Optional[str] = None,
        title_font_weight: str = "bold",
        font_size_base: float = 10.0,
        font_size: Optional[float] = None,
        axis_label_size: Optional[float] = None,
        tick_label_size: Optional[float] = None,
        line_width_base: float = 1.0,
        dpi: int = 96,
        figsize: tuple = (8, 6),
        context: str = "presentation",
        grid_visible: bool = True,
        grid_alpha: float = 0.5,
        grid_linestyle: str = "--",
        grid_linewidth: float = 0.8,
        zero_line_visible: bool = True,
        zero_line_color: str = "black",
        zero_line_linestyle: str = "--",
        zero_line_linewidth: float = 1.0,
        zero_line_alpha: float = 0.5,
        axis_baselines_visible: bool = True,
        axis_baselines_color: str = "black",
        axis_baselines_linestyle: str = "--",
        axis_baselines_linewidth: float = 1.0,
        axis_baselines_alpha: float = 0.5,
        spines_visible: bool = False,
        **kwargs
    ):
        """Initialize a Style.

        ``colors`` replaces the entire default palette.  ``color_overrides``
        merges on top of the default palette so you can change individual
        colours without repeating the rest.  ``font`` / ``font_size`` are
        short aliases for ``font_family`` / ``font_size_base``.  ``title_font``
        is a short alias for ``title_font_family``.
        """
        self.name = name

        # colors= replaces the full palette; color_overrides= merges on top of
        # the default palette so individual keys can be changed without
        # repeating every other colour.
        if colors is not None:
            merged = dict(colors)
            if color_overrides:
                merged.update(color_overrides)
        else:
            merged = self._default_colors()
            if color_overrides:
                merged.update(color_overrides)
        self.colors = merged

        # Short aliases take precedence when provided.
        self.font_family = font if font is not None else font_family
        self.title_font_family = title_font if title_font is not None else title_font_family
        self.title_font_weight = title_font_weight
        self.font_size_base = font_size if font_size is not None else font_size_base
        self.axis_label_size = (
            axis_label_size if axis_label_size is not None else self.font_size_base + 1.0
        )
        self.tick_label_size = (
            tick_label_size if tick_label_size is not None else max(1.0, self.font_size_base - 1.0)
        )
        self.line_width_base = line_width_base
        self.dpi = dpi
        self.figsize = figsize
        self.context = context
        self.grid_visible = grid_visible
        self.grid_alpha = grid_alpha
        self.grid_linestyle = grid_linestyle
        self.grid_linewidth = grid_linewidth
        self.zero_line_visible = zero_line_visible
        self.zero_line_color = zero_line_color
        self.zero_line_linestyle = zero_line_linestyle
        self.zero_line_linewidth = zero_line_linewidth
        self.zero_line_alpha = zero_line_alpha
        self.axis_baselines_visible = axis_baselines_visible
        self.axis_baselines_color = axis_baselines_color
        self.axis_baselines_linestyle = axis_baselines_linestyle
        self.axis_baselines_linewidth = axis_baselines_linewidth
        self.axis_baselines_alpha = axis_baselines_alpha
        self.spines_visible = spines_visible
        # Store any additional attributes from kwargs
        self._custom_attrs = kwargs
    
    @staticmethod
    def _default_colors() -> Dict[str, str]:
        """Return a default perceptually-uniform color palette."""
        return {
            "primary": "#1f77b4",
            "secondary": "#ff7f0e",
            "tertiary": "#2ca02c",
            "accent": "#d62728",
            "background": "#ffffff",
            "text": "#000000",
            "grid": "#cccccc",
            "neutral": "#7f7f7f",
        }
    
    def with_color(self, name: str, value: str) -> "Style":
        """
        Override a color and return self for chaining.
        
        Args:
            name: Color key (e.g., 'primary').
            value: Hex color code or named color.
        
        Returns:
            self (Style): Allows method chaining.
        """
        self.colors[name] = value
        return self
    
    def with_font(self, family: str, size: Optional[float] = None) -> "Style":
        """Override body font family (labels, ticks, legend) and optionally size.

        Args:
            family: Font family name.
            size: Optional new base font size.

        Returns:
            self (Style): Allows method chaining.
        """
        self.font_family = family
        if size is not None:
            self.font_size_base = size
        return self

    def with_title_font(
        self,
        family: str,
        size: Optional[float] = None,
        weight: Optional[str] = None,
    ) -> "Style":
        """Override title font family and optionally size / weight.

        Args:
            family: Font family name used for plot/axis titles.
            size: Optional new base font size (shared with body).
            weight: Optional title weight (e.g. "bold", "normal").

        Returns:
            self (Style): Allows method chaining.
        """
        self.title_font_family = family
        if size is not None:
            self.font_size_base = size
        if weight is not None:
            self.title_font_weight = weight
        return self

    def with_axis_label_size(self, size: float) -> "Style":
        """Override axis label size (x/y labels)."""
        self.axis_label_size = size
        return self

    def with_tick_label_size(self, size: float) -> "Style":
        """Override tick label size."""
        self.tick_label_size = size
        return self
    
    def with_dpi(self, dpi: int) -> "Style":
        """Override DPI."""
        self.dpi = dpi
        return self
    
    def with_figsize(self, width: float, height: float) -> "Style":
        """Override figure size in inches."""
        self.figsize = (width, height)
        return self

    def with_zero_line(
        self,
        visible: bool = True,
        color: Optional[str] = None,
        linestyle: Optional[str] = None,
        linewidth: Optional[float] = None,
        alpha: Optional[float] = None,
    ) -> "Style":
        """Configure the optional horizontal y=0 reference line."""
        self.zero_line_visible = visible
        if color is not None:
            self.zero_line_color = color
        if linestyle is not None:
            self.zero_line_linestyle = linestyle
        if linewidth is not None:
            self.zero_line_linewidth = linewidth
        if alpha is not None:
            self.zero_line_alpha = alpha
        return self

    def with_axis_baselines(
        self,
        visible: bool = True,
        color: Optional[str] = None,
        linestyle: Optional[str] = None,
        linewidth: Optional[float] = None,
        alpha: Optional[float] = None,
    ) -> "Style":
        """Configure custom left/bottom frame lines used when spines are hidden."""
        self.axis_baselines_visible = visible
        if color is not None:
            self.axis_baselines_color = color
        if linestyle is not None:
            self.axis_baselines_linestyle = linestyle
        if linewidth is not None:
            self.axis_baselines_linewidth = linewidth
        if alpha is not None:
            self.axis_baselines_alpha = alpha
        return self
    
    def to_matplotlib_rc(self) -> Dict[str, Any]:
        """
        Convert Style to matplotlib rcParams dictionary.
        
        Returns:
            dict: rcParams that can be passed to plt.rcdefaults() or set via plt.rc().
        """
        return {
            "font.family": self.font_family,
            "font.size": self.font_size_base,
            "axes.labelsize": self.axis_label_size,
            "xtick.labelsize": self.tick_label_size,
            "ytick.labelsize": self.tick_label_size,
            "lines.linewidth": self.line_width_base,
            "axes.edgecolor": self.colors.get("text", "black"),
            "axes.labelcolor": self.colors.get("text", "black"),
            "text.color": self.colors.get("text", "black"),
            "xtick.color": self.colors.get("text", "black"),
            "ytick.color": self.colors.get("text", "black"),
            "grid.color": self.colors.get("grid", "#cccccc"),
            "grid.alpha": self.grid_alpha,
            "grid.linewidth": self.grid_linewidth,
            "axes.grid": self.grid_visible,
            "figure.figsize": self.figsize,
            "figure.dpi": self.dpi,
            "savefig.dpi": self.dpi,
        }
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert Style to dictionary for serialization."""
        return {
            "name": self.name,
            "colors": self.colors,
            "font_family": self.font_family,
            "title_font_family": self.title_font_family,
            "title_font_weight": self.title_font_weight,
            "font_size_base": self.font_size_base,
            "axis_label_size": self.axis_label_size,
            "tick_label_size": self.tick_label_size,
            "line_width_base": self.line_width_base,
            "dpi": self.dpi,
            "figsize": self.figsize,
            "context": self.context,
            "grid_visible": self.grid_visible,
            "grid_alpha": self.grid_alpha,
            "grid_linestyle": self.grid_linestyle,
            "grid_linewidth": self.grid_linewidth,
            "zero_line_visible": self.zero_line_visible,
            "zero_line_color": self.zero_line_color,
            "zero_line_linestyle": self.zero_line_linestyle,
            "zero_line_linewidth": self.zero_line_linewidth,
            "zero_line_alpha": self.zero_line_alpha,
            "axis_baselines_visible": self.axis_baselines_visible,
            "axis_baselines_color": self.axis_baselines_color,
            "axis_baselines_linestyle": self.axis_baselines_linestyle,
            "axis_baselines_linewidth": self.axis_baselines_linewidth,
            "axis_baselines_alpha": self.axis_baselines_alpha,
            "spines_visible": self.spines_visible,
            **self._custom_attrs,
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Style":
        """Create a Style from a dictionary (e.g., loaded from YAML)."""
        return cls(**data)
    
    def __repr__(self) -> str:
        return f"Style(name='{self.name}', context='{self.context}', dpi={self.dpi})"


class ThemeLoader:
    """Load and cache Style objects from theme pack YAML files."""
    
    def __init__(self, theme_dir: Optional[Path] = None):
        """
        Initialize ThemeLoader.
        
        Args:
            theme_dir: Path to themes directory. If None, use bundled themes.
        """
        if theme_dir is None:
            # Use bundled themes in package
            theme_dir = Path(__file__).parent / "themes"
        self.theme_dir = Path(theme_dir)
        self._cache: Dict[str, Style] = {}
    
    def load(self, theme_path: str) -> Style:
        """
        Load a theme by path (e.g., 'publication/paper').
        
        Args:
            theme_path: Path to theme within theme packs (e.g., 'publication/paper').
        
        Returns:
            Style: Loaded and cached Style object.
        
        Raises:
            FileNotFoundError: If theme file does not exist.
            yaml.YAMLError: If theme YAML is malformed.
        """
        if theme_path in self._cache:
            return copy.deepcopy(self._cache[theme_path])
        
        # Construct full file path
        file_path = self.theme_dir / f"{theme_path}.yaml"
        
        if not file_path.exists():
            raise FileNotFoundError(
                f"Theme not found: {theme_path}\n"
                f"Searched in: {self.theme_dir}"
            )
        
        # Load YAML
        with open(file_path, "r") as f:
            data = yaml.safe_load(f) or {}
        
        # Ensure name is set
        if "name" not in data:
            data["name"] = theme_path
        
        # Create Style and cache
        style = Style.from_dict(data)
        self._cache[theme_path] = style
        
        return copy.deepcopy(style)
    
    def list_themes(self) -> Dict[str, list]:
        """
        List all available themes organized by pack.
        
        Returns:
            dict: Keys are pack names, values are lists of style names.
        """
        themes = {}
        if not self.theme_dir.exists():
            return themes
        
        for pack_dir in self.theme_dir.iterdir():
            if pack_dir.is_dir():
                pack_name = pack_dir.name
                style_names = [
                    f.stem for f in pack_dir.glob("*.yaml")
                ]
                if style_names:
                    themes[pack_name] = sorted(style_names)
        
        return themes
    
    def clear_cache(self):
        """Clear the theme cache."""
        self._cache.clear()


# Create a global theme loader instance
_default_loader = ThemeLoader()


def load_theme(theme_path: str) -> Style:
    """
    Load a theme by path (e.g., 'publication/paper').
    
    Args:
        theme_path: Path to theme within theme packs.
    
    Returns:
        Style: Loaded Style object.
    
    Raises:
        FileNotFoundError: If theme file does not exist.
    """
    return _default_loader.load(theme_path)


def list_themes() -> Dict[str, list]:
    """List all available themes organized by pack."""
    return _default_loader.list_themes()
