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

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


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 .yaml or .png recipe file. If .png is provided,
        the corresponding .yaml file will be loaded.
    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 ps
    >>> fig, ax = ps.reproduce("experiment_001.yaml")
    >>> fig, ax = ps.reproduce("experiment_001.png")  # Also works
    >>> plt.show()
    """
    path = Path(path)

    # Accept both .png and .yaml - find the yaml file
    if path.suffix.lower() in (".png", ".jpg", ".jpeg", ".pdf", ".svg"):
        yaml_path = path.with_suffix(".yaml")
        if not yaml_path.exists():
            raise FileNotFoundError(
                f"Recipe file not found: {yaml_path}. "
                f"Expected .yaml file alongside {path}"
            )
        path = yaml_path

    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:
                # Merge overrides into record style
                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:
                            # Merge call kwargs overrides
                            call.kwargs.update(call_overrides[call.id])

    return reproduce_from_record(
        record,
        calls=calls,
        skip_decorations=skip_decorations,
    )


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 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

    # Finalize tick configuration and special plot types (avoids categorical axis interference)
    from ..styles._style_applier import finalize_special_plots, finalize_ticks

    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 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())

    # 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

    # Check if it's a seaborn call
    if method_name.startswith("sns."):
        from ._seaborn import replay_seaborn_call

        return replay_seaborn_call(ax, call)

    # Handle violinplot with inner option specially
    if method_name == "violinplot":
        from ._violin import replay_violinplot_call

        return replay_violinplot_call(ax, call)

    # Handle joyplot specially (custom method)
    if method_name == "joyplot":
        from ._custom_plots import replay_joyplot_call

        return replay_joyplot_call(ax, call)

    # Handle swarmplot specially (custom method)
    if method_name == "swarmplot":
        from ._custom_plots import replay_swarmplot_call

        return replay_swarmplot_call(ax, call)

    method = getattr(ax, method_name, None)

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

    # Reconstruct args
    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
        # If it's already a Transform object or something else, leave it

    # 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 _reconstruct_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
    """Reconstruct kwargs, converting 2D lists back to numpy arrays.

    Parameters
    ----------
    kwargs : dict
        Raw kwargs from call record.

    Returns
    -------
    dict
        Kwargs with arrays properly reconstructed.
    """
    result = {}
    for key, value in kwargs.items():
        # Handle 'colors' parameter specially - must be a list for pie/bar/etc.
        # A single color string like 'red' would be interpreted as ['r','e','d']
        if key == "colors" and isinstance(value, str):
            result[key] = [value]
        elif isinstance(value, list) and len(value) > 0:
            # Check if it's a 2D list (list of lists) - should be numpy array
            if isinstance(value[0], list):
                result[key] = np.array(value)
            else:
                # 1D list - could be array or just list, try to preserve
                result[key] = value
        else:
            result[key] = value
    return result


def _reconstruct_value(
    arg_data: Dict[str, Any], result_cache: Optional[Dict[str, Any]] = None
) -> Any:
    """Reconstruct a value from serialized arg data.

    Parameters
    ----------
    arg_data : dict
        Serialized argument data.
    result_cache : dict, optional
        Cache mapping call_id -> result for resolving references.

    Returns
    -------
    Any
        Reconstructed value.
    """
    if result_cache is None:
        result_cache = {}

    # Check if we have a pre-loaded array
    if "_loaded_array" in arg_data:
        return arg_data["_loaded_array"]

    data = arg_data.get("data")

    # Check if it's a reference to another call's result (e.g., ContourSet for clabel)
    if isinstance(data, dict) and "__ref__" in data:
        ref_id = data["__ref__"]
        if ref_id in result_cache:
            return result_cache[ref_id]
        else:
            import warnings

            warnings.warn(f"Could not resolve reference to {ref_id}")
            return None

    # Check if it's a list of arrays (e.g., boxplot, violinplot)
    if arg_data.get("_is_array_list") and isinstance(data, list):
        dtype = arg_data.get("dtype")
        # Convert each inner list to numpy array
        return [
            np.array(arr_data, dtype=dtype if isinstance(dtype, str) else None)
            for arr_data in data
        ]

    # If data is a list, convert to numpy array
    if isinstance(data, list):
        dtype = arg_data.get("dtype")
        try:
            return np.array(data, dtype=dtype if dtype else None)
        except (TypeError, ValueError):
            return np.array(data)

    return data


def get_recipe_info(path: Union[str, Path]) -> Dict[str, Any]:
    """Get information about a recipe without reproducing.

    Parameters
    ----------
    path : str or Path
        Path to .yaml recipe file.

    Returns
    -------
    dict
        Recipe information including:
        - id: Figure ID
        - created: Creation timestamp
        - matplotlib_version: Version used
        - figsize: Figure size
        - n_axes: Number of axes
        - calls: List of call IDs
    """
    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",
]
