import os
from datetime import date

import polars as pl

os.environ.setdefault("ENDPOINT_URL", "http://localhost")
os.environ.setdefault("ENDPOINT_ACCESS_ID", "test")
os.environ.setdefault("ENDPOINT_ACCESS_TOKEN", "test")

from ldnode.loader.datasources.cb_quotations import (
    prepare_cb_pre_quotations_for_realtime,
    with_cb_quotation_adj_factor,
    with_cb_quotation_list_days,
    with_cb_quotation_size_min_unit_columns,
    with_cb_quotation_trade_min_unit_columns,
)


def test_cb_quotation_trade_builtin_columns_keep_min_units_before_scaling() -> None:
    result = (
        pl.DataFrame(
            {
                "vol": [12.0],
                "amount": [34.0],
            }
        )
        .lazy()
        .pipe(with_cb_quotation_trade_min_unit_columns)
        .collect()
    )

    assert result["volume"].to_list() == [12.0]
    assert result["amt"].to_list() == [34.0]


def test_cb_quotation_size_builtin_columns_convert_yi_to_min_currency_unit() -> None:
    result = (
        pl.DataFrame(
            {
                "remain_size": [56.0],
                "issue_size": [78.0],
                "remain_cap": [90.0],
            }
        )
        .lazy()
        .pipe(with_cb_quotation_size_min_unit_columns)
        .collect()
    )

    assert result["remaining_size"].to_list() == [5600000000.0]
    assert result["issuance_size"].to_list() == [7800000000.0]
    assert result["remaining_market_value"].to_list() == [9000000000.0]


def test_prepare_cb_pre_quotations_keeps_pre_adj_factor_and_drops_realtime_columns() -> None:
    result = prepare_cb_pre_quotations_for_realtime(
        pl.DataFrame(
            {
                "code": ["111010.SH"],
                "name": ["Z Li"],
                "adj_factor": [2.0],
                "close": [120.0],
                "pre_close": [119.0],
                "close_stk": [30.0],
                "remain_size": [1.0],
            }
        )
    )

    assert "_pre_adj_factor" in result.columns
    assert "adj_factor" not in result.columns
    assert "close" not in result.columns
    assert "pre_close" not in result.columns
    assert "close_stk" not in result.columns
    assert result["remain_size"].to_list() == [1.0]


def test_cb_quotation_adj_factor_uses_stock_realtime_formula() -> None:
    result = (
        pl.DataFrame(
            {
                "_pre_adj_factor": [2.0],
                "close": [121.0],
                "pre_close": [110.0],
            }
        )
        .lazy()
        .pipe(with_cb_quotation_adj_factor)
        .collect()
    )

    assert result["adj_factor"].to_list() == [2.2]


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

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