from datetime import date

import polars as pl

from ldnode.calculator.plans.polars.join_resolver import resolve_joins
from ldnode.calculator.plans.polars.join_rules import JOIN_RULES, SCOPE_TO_FILE
from ldnode.loader.files import STK_FINA_IND


def test_finance_scope_uses_stock_fina_ind_dataset() -> None:
    assert SCOPE_TO_FILE[10] is STK_FINA_IND

    cb_finance_rule = JOIN_RULES[(1, 10)]
    assert cb_finance_rule.column_suffix == "_fin"
    assert cb_finance_rule.main_join_key == "code_stk"
    assert cb_finance_rule.aux_join_key == "code"
    assert cb_finance_rule.time_key == "trade_date"
    assert cb_finance_rule.aux_time_key == "trade_date"

    stk_finance_rule = JOIN_RULES[(2, 10)]
    assert stk_finance_rule.column_suffix == "_fin"
    assert stk_finance_rule.aux_time_key == "trade_date"


def test_cb_finance_join_generates_finance_suffix_columns() -> None:
    main_lf = pl.DataFrame(
        {
            "code": ["CB001", "CB002"],
            "code_stk": ["000001.SZ", "000002.SZ"],
            "trade_date": [date(2024, 4, 10), date(2024, 4, 10)],
            "close": [101.0, 102.0],
        }
    ).lazy()
    finance_lf = pl.DataFrame(
        {
            "code": ["000001.SZ", "000002.SZ"],
            "trade_date": [date(2024, 3, 31), date(2024, 3, 31)],
            "bs_nwc": [1200, 2400],
        }
    ).lazy()

    result_lf, builtin_factors = resolve_joins(
        main_lf,
        strategy_scope=1,
        required_datasets=[
            {
                "security_type": 10,
                "columns": ["bs_nwc"],
                "column_suffix": "_fin",
                "pre_join_dsl": [],
            }
        ],
        builtin_factors=["code", "code_stk", "trade_date", "close"],
        aux_data_overrides={10: finance_lf},
    )

    result = result_lf.collect().sort("code")
    assert "bs_nwc_fin" in result.columns
    assert "bs_nwc_fin" in builtin_factors
    assert result["bs_nwc_fin"].to_list() == [1200, 2400]


def test_join_resolver_aliases_requested_aux_join_key_and_time_key() -> None:
    main_lf = pl.DataFrame(
        {
            "code": ["CB001"],
            "code_stk": ["000001.SZ"],
            "trade_date": [date(2024, 4, 10)],
            "close": [101.0],
        }
    ).lazy()
    stk_lf = pl.DataFrame(
        {
            "code": ["000001.SZ"],
            "trade_date": [date(2024, 4, 10)],
            "factor_a": [12.5],
        }
    ).lazy()

    result_lf, builtin_factors = resolve_joins(
        main_lf,
        strategy_scope=1,
        required_datasets=[
            {
                "security_type": 2,
                "columns": ["code", "trade_date", "factor_a"],
                "column_suffix": "_astk",
                "pre_join_dsl": [],
            }
        ],
        builtin_factors=["code", "code_stk", "trade_date", "close"],
        aux_data_overrides={2: stk_lf},
    )

    result = result_lf.collect()
    assert "factor_a_astk" in result.columns
    assert "factor_a_astk" in builtin_factors
    assert "code_astk" in result.columns
    assert "trade_date_astk" in result.columns
    assert "code_astk" in builtin_factors
    assert "trade_date_astk" in builtin_factors
    assert result["code_astk"].to_list() == ["000001.SZ"]
    assert result["trade_date_astk"].to_list() == [date(2024, 4, 10)]
    assert result["factor_a_astk"].to_list() == [12.5]


def test_cb_stock_nested_finance_join_computes_stock_scope_factor() -> None:
    main_lf = pl.DataFrame(
        {
            "code": ["CB001", "CB002"],
            "code_stk": ["000001.SZ", "000002.SZ"],
            "trade_date": [date(2024, 4, 3), date(2024, 4, 3)],
            "close": [101.0, 102.0],
        }
    ).lazy()
    stock_lf = pl.DataFrame(
        {
            "code": ["000001.SZ", "000001.SZ", "000002.SZ", "000002.SZ"],
            "trade_date": [
                date(2024, 4, 2),
                date(2024, 4, 3),
                date(2024, 4, 2),
                date(2024, 4, 3),
            ],
            "close": [11.0, 12.0, 21.0, 22.0],
        }
    ).lazy()
    finance_lf = pl.DataFrame(
        {
            "code": ["000001.SZ", "000002.SZ"],
            "trade_date": [date(2024, 3, 31), date(2024, 3, 31)],
            "bs_nwc": [100.0, 200.0],
        }
    ).lazy()

    result_lf, builtin_factors = resolve_joins(
        main_lf,
        strategy_scope=1,
        required_datasets=[
            {
                "security_type": 2,
                "columns": ["stock_factor_s"],
                "column_suffix": "_stk",
                "nested_required_datasets": [
                    {
                        "security_type": 10,
                        "columns": ["bs_nwc"],
                        "column_suffix": "_fin",
                        "pre_join_dsl": [],
                    }
                ],
                "scoped_computed_factors": [
                    [
                        {
                            "factor": "stock_factor_s",
                            "expr": "ts_shift(close, 1) + bs_nwc_fin",
                            "expr_type": 2,
                            "is_intermediate": False,
                        }
                    ]
                ],
            }
        ],
        builtin_factors=["code", "code_stk", "trade_date", "close"],
        aux_data_overrides={2: stock_lf, 10: finance_lf},
    )

    result = result_lf.collect().sort("code")
    assert "stock_factor_s_stk" in result.columns
    assert "stock_factor_s_stk" in builtin_factors
    assert result["stock_factor_s_stk"].to_list() == [111.0, 221.0]


def test_nested_join_rejects_unsupported_scope() -> None:
    main_lf = pl.DataFrame(
        {
            "code": ["CB001"],
            "code_stk": ["000001.SZ"],
            "trade_date": [date(2024, 4, 3)],
        }
    ).lazy()
    stock_lf = pl.DataFrame(
        {
            "code": ["000001.SZ"],
            "trade_date": [date(2024, 4, 3)],
        }
    ).lazy()

    try:
        resolve_joins(
            main_lf,
            strategy_scope=1,
            required_datasets=[
                {
                    "security_type": 10,
                    "columns": ["x"],
                    "nested_required_datasets": [
                        {"security_type": 2, "columns": ["close"]}
                    ],
                }
            ],
            builtin_factors=["code", "code_stk", "trade_date"],
            aux_data_overrides={2: stock_lf},
        )
    except ValueError as exc:
        assert "P0 仅支持 CB -> STK -> FIN/FINA" in str(exc)
    else:
        raise AssertionError("unsupported nested scope should fail")
