#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Core reproduction logic for figure reproduction."""

import shutil
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes

from .._recorder import CallRecord, FigureRecord
from .._serializer import load_recipe
from .._utils._bundle import resolve_recipe_path


def reproduce(
    path: Union[str, Path],
    calls: Optional[List[str]] = None,
    skip_decorations: bool = False,
    apply_overrides: bool = True,
):
    """Reproduce a figure from a recipe file.

    Parameters
    ----------
    path : str or Path
        Path to recipe. Supports multiple formats:
        - .yaml/.yml file: Direct recipe file
        - .png/.jpg/etc: Image with associated .yaml
        - Directory: Bundle containing recipe.yaml
        - .zip: ZIP archive containing recipe.yaml
    calls : list of str, optional
        If provided, only reproduce these specific call IDs.
    skip_decorations : bool
        If True, skip decoration calls (labels, legends, etc.).
    apply_overrides : bool
        If True (default), apply .overrides.json if it exists.
        This preserves manual GUI editor changes.

    Returns
    -------
    fig : RecordingFigure
        Reproduced figure (same type as subplots() returns).
    axes : RecordingAxes or ndarray of RecordingAxes
        Reproduced axes (single if 1x1, otherwise numpy array).

    Examples
    --------
    >>> import figrecipe as fr
    >>> fig, ax = fr.reproduce("experiment_001.yaml")
    >>> fig, ax = fr.reproduce("experiment_001.png")  # Also works
    >>> fig, ax = fr.reproduce("figure_bundle/")      # Directory bundle
    >>> fig, ax = fr.reproduce("figure.zip")          # ZIP bundle
    >>> plt.show()
    """
    # Resolve path to actual recipe YAML (handles directories, ZIPs, images)
    path, temp_dir = resolve_recipe_path(path)

    try:
        record = load_recipe(path)

        # Check for override file and merge if exists
        if apply_overrides:
            overrides_path = path.with_suffix(".overrides.json")
            if overrides_path.exists():
                import json

                with open(overrides_path) as f:
                    data = json.load(f)

                # Apply style overrides
                manual_overrides = data.get("manual_overrides", {})
                if manual_overrides:
                    if record.style is None:
                        record.style = {}
                    record.style.update(manual_overrides)

                # Apply call overrides (kwargs changes from editor)
                call_overrides = data.get("call_overrides", {})
                if call_overrides:
                    for ax_key, ax_record in record.axes.items():
                        for call in ax_record.calls:
                            if call.id in call_overrides:
                                call.kwargs.update(call_overrides[call.id])

        return reproduce_from_record(
            record, calls=calls, skip_decorations=skip_decorations
        )
    finally:
        if temp_dir is not None and temp_dir.exists():
            shutil.rmtree(temp_dir, ignore_errors=True)


def reproduce_from_record(
    record: FigureRecord,
    calls: Optional[List[str]] = None,
    skip_decorations: bool = False,
):
    """Reproduce a figure from a FigureRecord.

    Parameters
    ----------
    record : FigureRecord
        The figure record to reproduce.
    calls : list of str, optional
        If provided, only reproduce these specific call IDs.
    skip_decorations : bool
        If True, skip decoration calls.

    Returns
    -------
    fig : RecordingFigure
        Reproduced figure (wrapped).
    axes : RecordingAxes or ndarray of RecordingAxes
        Reproduced axes (wrapped, numpy array for multi-axes).
    """
    from .._recorder import Recorder
    from .._wrappers import RecordingAxes, RecordingFigure

    # Determine grid size from axes positions
    max_row = 0
    max_col = 0
    for ax_key in record.axes.keys():
        parts = ax_key.split("_")
        if len(parts) >= 3:
            max_row = max(max_row, int(parts[1]))
            max_col = max(max_col, int(parts[2]))

    nrows = max_row + 1
    ncols = max_col + 1

    # Create figure
    fig, mpl_axes = plt.subplots(
        nrows,
        ncols,
        figsize=record.figsize,
        dpi=record.dpi,
        constrained_layout=record.constrained_layout,
    )

    # Apply layout if recorded (skip if constrained_layout is used)
    if record.layout is not None and not record.constrained_layout:
        fig.subplots_adjust(**record.layout)

    # Ensure axes is 2D array
    if nrows == 1 and ncols == 1:
        axes_2d = np.array([[mpl_axes]])
    else:
        axes_2d = np.atleast_2d(mpl_axes)
        if nrows == 1:
            axes_2d = axes_2d.reshape(1, -1)
        elif ncols == 1:
            axes_2d = axes_2d.reshape(-1, 1)

    # Apply style BEFORE replaying calls (to match original order:
    # style is applied during subplots(), then user creates plots/decorations)
    if record.style is not None:
        from ..styles._internal import apply_style_mm

        for row in range(nrows):
            for col in range(ncols):
                apply_style_mm(axes_2d[row, col], record.style)

    # Result cache for resolving references (e.g., clabel needs ContourSet from contour)
    result_cache: Dict[str, Any] = {}

    # Replay calls on each axes
    for ax_key, ax_record in record.axes.items():
        parts = ax_key.split("_")
        if len(parts) >= 3:
            row, col = int(parts[1]), int(parts[2])
        else:
            row, col = 0, 0

        ax = axes_2d[row, col]

        # Replay plotting calls
        for call in ax_record.calls:
            if calls is not None and call.id not in calls:
                continue
            result = _replay_call(ax, call, result_cache)
            if result is not None:
                result_cache[call.id] = result

        # Replay decorations
        if not skip_decorations:
            for call in ax_record.decorations:
                if calls is not None and call.id not in calls:
                    continue
                result = _replay_call(ax, call, result_cache)
                if result is not None:
                    result_cache[call.id] = result

        # Apply panel visibility
        if not getattr(ax_record, "visible", True):
            ax.set_visible(False)

    # Replay recorded colorbars (exact reproduction)
    _replay_colorbars(fig, axes_2d, record, result_cache)

    # Finalize tick configuration and special plot types
    from ..styles._style_applier import finalize_special_plots, finalize_ticks
    from ._line_styles import apply_line_styles

    for row in range(nrows):
        for col in range(ncols):
            finalize_ticks(axes_2d[row, col])
            finalize_special_plots(axes_2d[row, col], record.style or {})

    # Apply trace linewidth to all Line2D objects created during replay
    apply_line_styles(axes_2d, record.style or {})

    # Apply figure-level labels if recorded
    if record.suptitle is not None:
        text = record.suptitle.get("text", "")
        kwargs = record.suptitle.get("kwargs", {}).copy()
        # Only add y=1.02 if not using constrained_layout (which handles positioning)
        if "y" not in kwargs and not record.constrained_layout:
            kwargs["y"] = 1.02
        fig.suptitle(text, **kwargs)

    if record.supxlabel is not None:
        text = record.supxlabel.get("text", "")
        kwargs = record.supxlabel.get("kwargs", {})
        fig.supxlabel(text, **kwargs)

    if record.supylabel is not None:
        text = record.supylabel.get("text", "")
        kwargs = record.supylabel.get("kwargs", {})
        fig.supylabel(text, **kwargs)

    # Wrap in Recording types (same as subplots() returns)
    recorder = Recorder()
    recorder._figure_record = record

    # Wrap axes in RecordingAxes
    wrapped_axes = np.empty((nrows, ncols), dtype=object)
    for i in range(nrows):
        for j in range(ncols):
            wrapped_axes[i, j] = RecordingAxes(axes_2d[i, j], recorder, position=(i, j))

    # Create RecordingFigure
    wrapped_fig = RecordingFigure(fig, recorder, wrapped_axes.tolist())

    # Restore mm_layout for consistent cropping during save
    if record.mm_layout is not None:
        wrapped_fig._mm_layout = record.mm_layout

    # Reproduce panel labels if recorded
    if record.panel_labels is not None:
        labels = record.panel_labels.get("labels")
        loc = record.panel_labels.get("loc", "upper left")
        offset = tuple(record.panel_labels.get("offset", (-0.1, 1.05)))
        fontsize = record.panel_labels.get("fontsize")
        fontweight = record.panel_labels.get("fontweight", "bold")
        color = record.panel_labels.get("color")
        extra_kwargs = record.panel_labels.get("kwargs", {})
        if color is not None:
            extra_kwargs["color"] = color
        wrapped_fig.add_panel_labels(
            labels=labels,
            loc=loc,
            offset=offset,
            fontsize=fontsize,
            fontweight=fontweight,
            **extra_kwargs,
        )

    # Return in appropriate format (matching subplots() behavior)
    if nrows == 1 and ncols == 1:
        return wrapped_fig, wrapped_axes[0, 0]
    elif nrows == 1:
        return wrapped_fig, np.array(wrapped_axes[0], dtype=object)
    elif ncols == 1:
        return wrapped_fig, np.array(wrapped_axes[:, 0], dtype=object)
    else:
        return wrapped_fig, wrapped_axes


def _replay_call(
    ax: Axes, call: CallRecord, result_cache: Optional[Dict[str, Any]] = None
) -> Any:
    """Replay a single call on an axes.

    Parameters
    ----------
    ax : Axes
        The matplotlib axes.
    call : CallRecord
        The call to replay.
    result_cache : dict, optional
        Cache mapping call_id -> result for resolving references.

    Returns
    -------
    Any
        Result of the matplotlib call.
    """
    if result_cache is None:
        result_cache = {}

    method_name = call.function

    # Special method handlers
    if method_name.startswith("sns."):
        from ._seaborn import replay_seaborn_call

        return replay_seaborn_call(ax, call)
    if method_name == "boxplot":
        from ._boxplot import replay_boxplot_call

        return replay_boxplot_call(ax, call)
    if method_name == "violinplot":
        from ._violin import replay_violinplot_call

        return replay_violinplot_call(ax, call)
    if method_name == "joyplot":
        from ._custom_plots import replay_joyplot_call

        return replay_joyplot_call(ax, call)
    if method_name == "swarmplot":
        from ._custom_plots import replay_swarmplot_call

        return replay_swarmplot_call(ax, call)
    if method_name == "stat_annotation":
        from .._wrappers._stat_annotation import draw_stat_annotation

        kwargs = call.kwargs.copy()
        x1, x2 = kwargs.pop("x1", 0), kwargs.pop("x2", 1)
        return draw_stat_annotation(ax, x1, x2, **kwargs)
    if method_name == "graph":
        from ._replay_graph import replay_graph_call

        return replay_graph_call(ax, call)
    if method_name == "diagram":
        from ._replay_diagram import replay_schematic_call

        return replay_schematic_call(ax, call)
    if method_name == "legend":
        from ._legend import replay_legend_call

        return replay_legend_call(ax, call, result_cache)
    if method_name == "stem":
        from ._stem import replay_stem_call

        return replay_stem_call(ax, call)

    method = getattr(ax, method_name, None)

    if method is None:
        # Method not found, skip
        return None

    # Reconstruct args
    from ._reconstruct import reconstruct_kwargs, reconstruct_value

    args = []
    for arg_data in call.args:
        value = reconstruct_value(arg_data, result_cache)
        args.append(value)

    # Get kwargs and reconstruct arrays
    kwargs = reconstruct_kwargs(call.kwargs)

    # Handle special transform markers
    if "transform" in kwargs:
        transform_val = kwargs["transform"]
        if transform_val == "axes":
            kwargs["transform"] = ax.transAxes
        elif transform_val == "data":
            kwargs["transform"] = ax.transData
        elif transform_val == "figure":
            kwargs["transform"] = ax.figure.transFigure

    # Fix fill_between: 'color' overrides 'edgecolor', use 'facecolor' instead
    if method_name in ("fill_between", "fill_betweenx"):
        if "color" in kwargs and "edgecolor" in kwargs:
            kwargs["facecolor"] = kwargs.pop("color")

    # Call the method
    try:
        return method(*args, **kwargs)
    except Exception as e:
        # Log warning but continue
        import warnings

        warnings.warn(f"Failed to replay {method_name}: {e}")
        return None


def _replay_colorbars(fig, axes_2d, record, result_cache):
    """Replay recorded colorbars for exact reproduction."""
    # Get recorded colorbars
    colorbars = getattr(record, "colorbars", []) or []

    if not colorbars:
        # No recorded colorbars - nothing to do
        # (Old recipes without colorbar recording won't reproduce colorbars)
        return

    # 2D plot methods that return ScalarMappable (for finding mappable)
    COLORBAR_METHODS = {
        "imshow",
        "pcolormesh",
        "pcolor",
        "contourf",
        "contour",
        "hexbin",
        "hist2d",
        "tripcolor",
        "tricontourf",
        "tricontour",
        "matshow",
        "spy",
        "specgram",
    }

    # Replay each recorded colorbar
    for cbar_info in colorbars:
        ax_key = cbar_info.get("ax_key")
        kwargs = cbar_info.get("kwargs", {})

        if ax_key is None:
            continue

        # Parse ax position
        parts = ax_key.split("_")
        if len(parts) >= 3:
            row, col = int(parts[1]), int(parts[2])
        else:
            continue

        ax = axes_2d[row, col]

        # Find the mappable for this axes from result_cache
        ax_record = record.axes.get(ax_key)
        if ax_record is None:
            continue

        mappable = None
        method_name = None
        for call in ax_record.calls:
            if call.function in COLORBAR_METHODS and call.id in result_cache:
                mappable = result_cache[call.id]
                method_name = call.function
                break

        if mappable is None:
            continue

        # Some methods return tuples - extract the actual mappable
        if isinstance(mappable, tuple):
            if method_name == "hist2d":
                # hist2d returns (counts, xedges, yedges, image)
                mappable = mappable[3]
            elif method_name == "specgram":
                # specgram returns (spectrum, freqs, t, image)
                mappable = mappable[3]

        # Add colorbar with recorded kwargs - use raw fig.colorbar to avoid
        # add_colorbar adding extra styling kwargs that weren't in original
        cbar = fig.colorbar(mappable, ax=ax, **kwargs)

        # Apply styling to match original (style_colorbar is called by add_colorbar)
        from .._utils._colorbar import style_colorbar

        style_colorbar(cbar, record.style)


def get_recipe_info(path: Union[str, Path]) -> Dict[str, Any]:
    """Get recipe metadata (id, figsize, dpi, n_axes, calls) without reproducing."""
    record = load_recipe(path)

    all_calls = []
    for ax_record in record.axes.values():
        for call in ax_record.calls:
            all_calls.append(
                {
                    "id": call.id,
                    "function": call.function,
                    "n_args": len(call.args),
                    "kwargs": list(call.kwargs.keys()),
                }
            )
        for call in ax_record.decorations:
            all_calls.append(
                {
                    "id": call.id,
                    "function": call.function,
                    "type": "decoration",
                }
            )

    return {
        "id": record.id,
        "created": record.created,
        "matplotlib_version": record.matplotlib_version,
        "figsize": record.figsize,
        "dpi": record.dpi,
        "n_axes": len(record.axes),
        "calls": all_calls,
    }


__all__ = [
    "reproduce",
    "reproduce_from_record",
    "get_recipe_info",
]
