from datetime import date, datetime

import polars as pl
import pytest

from ldnode.calculator.operators import (
    stk_event_backtest,
    stk_screen,
    stk_vector_backtest,
)
from ldnode.calculator.schemas import (
    BenchmarkType,
    HoldRangeType,
    HoldWeight,
    RebalanceType,
    RotationType,
    SortOrder,
    StopType,
)


def _stock_factor_frame(factor_values: list[float | None] | None = None) -> pl.DataFrame:
    trade_dates = [date(2026, 6, 8), date(2026, 6, 9)]
    codes = ["A", "B", "C", "D", "E", "F"]
    factors = factor_values or [1.0, 2.0, 3.0, 0.0, 0.0, 0.0]
    rows = []
    for trade_date in trade_dates:
        for idx, code in enumerate(codes):
            rows.append(
                {
                    "trade_date": trade_date,
                    "time": datetime.combine(trade_date, datetime.min.time()),
                    "code": code,
                    "name": f"stock-{code}",
                    "close": 10.0 + idx,
                    "pct_chg": 0.01,
                    "return": 0.01,
                    "amount": 1000.0 + idx * 100,
                    "limit": 0,
                    "suspension": False,
                    "last_trade_day": False,
                    "stk_div": None,
                    "cash_div": None,
                    "rank_alpha": 100.0 - idx,
                    "factor_alpha": factors[idx],
                }
            )
    return pl.DataFrame(rows)


def _stock_parameters(hold_weight: HoldWeight, factor_values: list[float | None] | None = None):
    hold_method = {
        "hold_range_type": HoldRangeType.NUM,
        "hold_range_value": {"start": 1, "end": 6},
        "hold_threshold": 0,
        "hold_weight": hold_weight,
        "max_hold_weight": 1,
    }
    if hold_weight == HoldWeight.FACTOR:
        hold_method["hold_weight_factor"] = "factor_alpha"
    return {
        "computed_factors": [],
        "filter_factors": None,
        "rank_factors": [
            {
                "factor": "rank_alpha",
                "ascending": False,
                "weight": 1.0,
                "neutralization_factor": None,
                "process_missing_data_method": None,
            }
        ],
        "more_by": [("code", SortOrder.ASCENDING)],
        "hold_method": hold_method,
        "date_range": {"start_date": "2026-06-08", "end_date": "2026-06-09"},
        "capital": 1_000_000.0,
        "commission": 0.001,
        "benchmark": BenchmarkType.EQUALWEIGHT,
        "rebalance_method": {
            "rebalance_type": RebalanceType.TRADEDAY,
            "rebalance_value": 1,
        },
        "stop_method": {"stop_type": StopType.NOTSTOP, "stop_value": 0.0},
        "rotation_type": RotationType.BCSC,
    }


def test_factor_weight_screen_vector_event_paths_use_same_positive_factor_input():
    frame = _stock_factor_frame()
    parameters = _stock_parameters(HoldWeight.FACTOR)

    screen = stk_screen(
        frame.lazy(),
        date(2026, 6, 8),
        [parameters],
        builtin_factors=["rank_alpha", "factor_alpha"],
    )[0]
    selected_screen = screen.filter(pl.col("selected")).sort("code")
    unselected_screen = screen.filter(~pl.col("selected")).sort("code")

    assert selected_screen["code"].to_list() == ["A", "B", "C"]
    assert selected_screen["_weight"].to_list() == [1.0, 2.0, 3.0]
    assert unselected_screen["code"].to_list() == ["D", "E", "F"]
    assert unselected_screen["_weight"].to_list() == [0.0, 0.0, 0.0]

    vector = stk_vector_backtest(
        frame.lazy(), parameters, builtin_factors=["rank_alpha", "factor_alpha"]
    )
    assert vector.height == 2
    assert sorted(vector["hold_securities"].to_list()[1]) == ["A", "B", "C"]

    event, history = stk_event_backtest(
        frame.lazy(),
        parameters,
        record_history=True,
        builtin_factors=["rank_alpha", "factor_alpha"],
    )
    assert event.height == 2
    assert history.height > 0


def test_factor_weight_screen_vector_event_paths_block_all_invalid_values():
    frame = _stock_factor_frame([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    parameters = _stock_parameters(HoldWeight.FACTOR)

    with pytest.raises(ValueError, match="FACTOR 权重因子全部无效"):
        stk_screen(
            frame.lazy(),
            date(2026, 6, 8),
            [parameters],
            builtin_factors=["rank_alpha", "factor_alpha"],
        )
    with pytest.raises(ValueError, match="FACTOR 权重因子全部无效"):
        stk_vector_backtest(
            frame.lazy(), parameters, builtin_factors=["rank_alpha", "factor_alpha"]
        )
    with pytest.raises(ValueError, match="FACTOR 权重因子全部无效"):
        stk_event_backtest(
            frame.lazy(),
            parameters,
            record_history=False,
            builtin_factors=["rank_alpha", "factor_alpha"],
        )


@pytest.mark.parametrize(
    ("factor_values", "error_message"),
    [
        ([1.0, 2.0, None, 0.0, 0.0, 0.0], "FACTOR 权重因子存在空值"),
        ([1.0, 2.0, -1.0, 0.0, 0.0, 0.0], "FACTOR 权重因子存在负值"),
    ],
)
def test_factor_weight_screen_vector_event_paths_block_null_and_negative_values(
    factor_values, error_message
):
    frame = _stock_factor_frame(factor_values)
    parameters = _stock_parameters(HoldWeight.FACTOR)

    with pytest.raises(ValueError, match=error_message):
        stk_screen(
            frame.lazy(),
            date(2026, 6, 8),
            [parameters],
            builtin_factors=["rank_alpha", "factor_alpha"],
        )
    with pytest.raises(ValueError, match=error_message):
        stk_vector_backtest(
            frame.lazy(), parameters, builtin_factors=["rank_alpha", "factor_alpha"]
        )
    with pytest.raises(ValueError, match=error_message):
        stk_event_backtest(
            frame.lazy(),
            parameters,
            record_history=False,
            builtin_factors=["rank_alpha", "factor_alpha"],
        )


@pytest.mark.parametrize(
    "hold_weight",
    [HoldWeight.NOTHOLD, HoldWeight.AVG, HoldWeight.CLOSE, HoldWeight.AMOUNT],
)
def test_old_hold_weight_modes_still_run_screen_vector_event(hold_weight):
    frame = _stock_factor_frame()
    parameters = _stock_parameters(hold_weight)

    screen = stk_screen(
        frame.lazy(),
        date(2026, 6, 8),
        [parameters],
        builtin_factors=["rank_alpha", "factor_alpha"],
    )[0]
    vector = stk_vector_backtest(
        frame.lazy(), parameters, builtin_factors=["rank_alpha", "factor_alpha"]
    )
    event, _ = stk_event_backtest(
        frame.lazy(),
        parameters,
        record_history=False,
        builtin_factors=["rank_alpha", "factor_alpha"],
    )

    assert screen.height == 6
    assert vector.height == 2
    assert event.height == 2
