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 == "report_date"

    stk_finance_rule = JOIN_RULES[(2, 10)]
    assert stk_finance_rule.column_suffix == "_fin"
    assert stk_finance_rule.aux_time_key == "report_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"],
            "report_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]
