"""
Plotting module.

This module handles visualization of backtest results:
- Equity curve plot
"""

from pathlib import Path

import matplotlib.pyplot as plt

from jbqlab.types import BacktestResult

# =============================================================================
# Public API
# =============================================================================


def plot_equity_curve(
    result: BacktestResult,
    output_path: str | Path | None = None,
    title: str | None = None,
    figsize: tuple[int, int] = (12, 6),
    show: bool = False,
) -> plt.Figure | None:
    """Plot equity curve from backtest results.

    Args:
        result: BacktestResult from run_backtest.
        output_path: Path to save the plot. If None, returns figure without saving.
        title: Plot title. If None, auto-generates from metadata.
        figsize: Figure size in inches (width, height).
        show: Whether to display the plot interactively.

    Returns:
        matplotlib Figure object if output_path is None, else None.
    """
    fig, ax = plt.subplots(figsize=figsize)

    # Plot equity curve
    equity = result.equity_curve
    ax.plot(equity.index, equity.values, linewidth=1.5, color="#2E86AB", label="Portfolio Value")

    # Add drawdown shading
    running_max = equity.cummax()
    ax.fill_between(
        equity.index,
        equity.values,
        running_max.values,
        alpha=0.3,
        color="#E94F37",
        label="Drawdown",
    )

    # Formatting
    if title is None:
        strategy = result.metadata.get("strategy", "Unknown")
        total_return = result.metrics.get("total_return", 0)
        sharpe = result.metrics.get("sharpe", 0)
        title = f"{strategy} | Return: {total_return * 100:.1f}% | Sharpe: {sharpe:.2f}"

    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.set_xlabel("Date", fontsize=12)
    ax.set_ylabel("Portfolio Value", fontsize=12)
    ax.legend(loc="upper left")
    ax.grid(True, alpha=0.3)

    # Format y-axis with comma separators
    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _p: f"${x:,.0f}"))

    # Rotate x-axis labels
    plt.xticks(rotation=45)
    plt.tight_layout()

    if output_path is not None:
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return None

    if show:
        plt.show()
        return None

    return fig


def plot_returns_distribution(
    result: BacktestResult,
    output_path: str | Path | None = None,
    figsize: tuple[int, int] = (10, 6),
    show: bool = False,
) -> plt.Figure | None:
    """Plot histogram of daily returns.

    Args:
        result: BacktestResult from run_backtest.
        output_path: Path to save the plot. If None, returns figure without saving.
        figsize: Figure size in inches (width, height).
        show: Whether to display the plot interactively.

    Returns:
        matplotlib Figure object if output_path is None, else None.
    """
    fig, ax = plt.subplots(figsize=figsize)

    returns = result.returns * 100  # Convert to percentage
    returns = returns[returns != 0]  # Exclude days with no position

    ax.hist(returns, bins=50, edgecolor="black", alpha=0.7, color="#2E86AB")

    # Add vertical line at mean
    mean_return = returns.mean()
    ax.axvline(mean_return, color="#E94F37", linestyle="--", linewidth=2, label=f"Mean: {mean_return:.2f}%")

    ax.set_title("Daily Returns Distribution", fontsize=14, fontweight="bold")
    ax.set_xlabel("Daily Return (%)", fontsize=12)
    ax.set_ylabel("Frequency", fontsize=12)
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.tight_layout()

    if output_path is not None:
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        return None

    if show:
        plt.show()
        return None

    return fig
