from datetime import date

import polars as pl

from ldnode.loader.datasources.cb_quotations import with_cb_quotation_list_days
from ldnode.loader.datasources.etf_quotations import fund_realtime_list_days_expr
from ldnode.loader.datasources.stk_quotations import stock_realtime_list_days_expr


def collect_realtime_list_days(expr: pl.Expr, list_date: date) -> list[int]:
    return (
        pl.DataFrame(
            {
                "trade_date": [date(2026, 6, 17)],
                "list_date": [list_date],
            }
        )
        .lazy()
        .with_columns(expr)
        .collect()["list_days"]
        .to_list()
    )


def test_stock_realtime_list_days_aligns_tail_close_inclusive_days() -> None:
    assert collect_realtime_list_days(stock_realtime_list_days_expr(), date(2026, 6, 17)) == [1]
    assert collect_realtime_list_days(stock_realtime_list_days_expr(), date(2026, 6, 1)) == [17]


def test_fund_realtime_list_days_aligns_tail_close_inclusive_days() -> None:
    assert collect_realtime_list_days(fund_realtime_list_days_expr(), date(2026, 6, 17)) == [1]
    assert collect_realtime_list_days(fund_realtime_list_days_expr(), date(2026, 6, 1)) == [17]


def test_cb_realtime_list_days_aligns_tail_close_inclusive_days() -> None:
    result = (
        pl.DataFrame(
            {
                "trade_date": [date(2026, 6, 17), date(2026, 6, 17)],
                "list_date": [date(2026, 6, 17), date(2026, 6, 1)],
            }
        )
        .lazy()
        .pipe(with_cb_quotation_list_days)
        .collect()
    )

    assert result["list_days"].to_list() == [1, 17]
    assert result.schema["list_days"] == pl.Int16
