"""Tests for reusable chart components."""

from __future__ import annotations

import polars as pl
import pytest
from tradepose_client.analysis.mae_mfe.config import DEFAULT_CONFIG
from tradepose_client.analysis.trades.components import (
    get_color_values,
    resolve_axis,
)


class TestResolveAxis:
    """Test axis resolution for string and pl.Expr."""

    def test_resolve_string_column(self, sample_trades_df):
        """Test resolving a simple column name."""
        label, values = resolve_axis(sample_trades_df, "pnl")

        assert label == "pnl"
        assert isinstance(values, pl.Series)
        assert len(values) == len(sample_trades_df)

    def test_resolve_polars_expression(self, sample_trades_df):
        """Test resolving a Polars expression."""
        expr = pl.col("mae") / pl.col("mae_volatility")
        label, values = resolve_axis(sample_trades_df, expr)

        assert isinstance(values, pl.Series)
        assert len(values) == len(sample_trades_df)

        # Values should match manual calculation
        expected = sample_trades_df["mae"] / sample_trades_df["mae_volatility"]
        assert values.to_list() == pytest.approx(expected.to_list(), rel=0.01)

    def test_resolve_aliased_expression(self, sample_trades_df):
        """Test resolving an aliased expression."""
        expr = (pl.col("mae") / pl.col("mae_volatility")).alias("mae_atr")
        label, values = resolve_axis(sample_trades_df, expr)

        assert label == "mae_atr"
        assert len(values) == len(sample_trades_df)

    def test_resolve_complex_expression(self, sample_trades_df):
        """Test resolving a complex expression."""
        expr = (pl.col("mfe") - pl.col("mae")) / pl.col("entry_price")
        label, values = resolve_axis(sample_trades_df, expr)

        assert isinstance(values, pl.Series)
        assert len(values) == len(sample_trades_df)


class TestGetColorValues:
    """Test color value generation for different color_by options."""

    def test_color_by_outcome(self, sample_trades_df):
        """Test coloring by outcome (Win/Lose)."""
        color_series, color_map, legend_title = get_color_values(
            sample_trades_df, "outcome", DEFAULT_CONFIG
        )

        assert legend_title == "Outcome"
        assert "Win" in color_map
        assert "Lose" in color_map
        assert len(color_series) == len(sample_trades_df)

    def test_color_by_direction(self, sample_trades_df):
        """Test coloring by direction (Long/Short)."""
        color_series, color_map, legend_title = get_color_values(
            sample_trades_df, "direction", DEFAULT_CONFIG
        )

        assert legend_title == "Direction"
        assert "Long" in color_map
        assert "Short" in color_map

    def test_color_by_exit_reason(self, sample_trades_df):
        """Test coloring by exit reason."""
        color_series, color_map, legend_title = get_color_values(
            sample_trades_df, "exit_reason", DEFAULT_CONFIG
        )

        assert legend_title == "Exit Reason"
        # Should have labels from EXIT_REASON_LABELS
        assert any("Stop Loss" in str(k) or "Take Profit" in str(k) for k in color_map)

    def test_color_by_entry_reason(self, sample_trades_df):
        """Test coloring by entry reason."""
        color_series, color_map, legend_title = get_color_values(
            sample_trades_df, "entry_reason", DEFAULT_CONFIG
        )

        assert legend_title == "Entry Reason"

    def test_color_by_strategy(self, sample_trades_df):
        """Test coloring by strategy name."""
        color_series, color_map, legend_title = get_color_values(
            sample_trades_df, "strategy", DEFAULT_CONFIG
        )

        assert legend_title == "Strategy"
        assert "Strategy_A" in color_map
        assert "Strategy_B" in color_map

    def test_color_by_blueprint(self, sample_trades_df):
        """Test coloring by blueprint name."""
        color_series, color_map, legend_title = get_color_values(
            sample_trades_df, "blueprint", DEFAULT_CONFIG
        )

        assert legend_title == "Blueprint"
        assert "Blueprint_1" in color_map or "Blueprint_2" in color_map

    def test_color_by_missing_column_error(self, minimal_trades_df):
        """Test error when required column is missing."""
        with pytest.raises(ValueError, match="exit_reason column not found"):
            get_color_values(minimal_trades_df, "exit_reason", DEFAULT_CONFIG)

    def test_color_by_unknown_option_error(self, sample_trades_df):
        """Test error for unknown color_by option."""
        with pytest.raises(ValueError, match="Unknown color_by option"):
            get_color_values(sample_trades_df, "unknown", DEFAULT_CONFIG)
