from datetime import date, datetime

import polars as pl

from ldnode.calculator.handlers.risk_control import (
    RISK_CONTROL_COLUMN,
    cb_risk_control_expr,
    delist_risk_control_expr,
    empty_risk_control_expr,
    realtime_risk_control_fixed_columns,
    stock_fund_risk_control_expr,
)
from ldnode.calculator.operators import cb_screen, etf_screen, stk_screen
from ldnode.calculator.schemas import (
    CB_SCREEN_RESULT_FIXED_COLUMNS,
    ETF_SCREEN_RESULT_FIXED_COLUMNS,
    HoldRangeType,
    HoldWeight,
    STK_SCREEN_RESULT_FIXED_COLUMNS,
    SortOrder,
)
from ldnode.loader.datasources.etf_quotations import Model as ETFQuotationModel


def _screen_parameters() -> dict:
    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_range_type": HoldRangeType.NUM,
            "hold_range_value": {"start": 1, "end": 3},
            "hold_threshold": 0,
            "hold_weight": HoldWeight.AVG,
            "max_hold_weight": 1,
        },
    }


def _screen_frame() -> pl.DataFrame:
    trade_date = date(2026, 6, 10)
    rows = [
        {
            "code": "A",
            "trade_date": trade_date,
            "time": datetime.combine(trade_date, datetime.min.time()),
            "name": "退市股-A",
            "close": 10.0,
            "pct_chg": 0.01,
            "amount": 1000.0,
            "limit": 0,
            "suspension": False,
            "rank_alpha": 3.0,
            "delist_date": date(2026, 7, 1),
        },
        {
            "code": "B",
            "trade_date": trade_date,
            "time": datetime.combine(trade_date, datetime.min.time()),
            "name": "stock-B",
            "close": 9.0,
            "pct_chg": 0.0,
            "amount": 900.0,
            "limit": 0,
            "suspension": False,
            "rank_alpha": 2.0,
            "delist_date": None,
        },
    ]
    return pl.DataFrame(rows)


def test_delist_risk_control_expr_generates_list_string_without_row_loop():
    result = _screen_frame().with_columns(delist_risk_control_expr())

    assert result[RISK_CONTROL_COLUMN].to_list() == [
        ["已公告于2026-07-01退市"],
        [],
    ]


def test_stock_fund_risk_control_expr_generates_multiple_list_items():
    result = _screen_frame().with_columns(stock_fund_risk_control_expr())

    assert result[RISK_CONTROL_COLUMN].to_list() == [
        ["已公告于2026-07-01退市", "已公告退市"],
        [],
    ]


def test_stock_fund_risk_control_expr_uses_name_when_delist_date_missing():
    result = pl.DataFrame(
        {
            "name": ["退市基金", "正常基金"],
        }
    ).with_columns(stock_fund_risk_control_expr(has_delist_date=False))

    assert result[RISK_CONTROL_COLUMN].to_list() == [["已公告退市"], []]


def test_cb_risk_control_expr_generates_last_trade_day_item_from_name():
    result = pl.DataFrame(
        {
            "name": ["ZZ转债", "正常转债"],
        }
    ).with_columns(cb_risk_control_expr())

    assert result[RISK_CONTROL_COLUMN].to_list() == [["退市前最后交易日"], []]


def test_cb_risk_control_expr_uses_sina_realtime_name():
    result = pl.DataFrame(
        {
            "name": ["Z甬矽转", "兴瑞转债"],
        }
    ).with_columns(cb_risk_control_expr())

    assert result["name"].to_list() == ["Z甬矽转", "兴瑞转债"]
    assert result[RISK_CONTROL_COLUMN].to_list() == [["退市前最后交易日"], []]


def test_empty_risk_control_expr_generates_typed_empty_list_column():
    result = pl.DataFrame({"code": ["A", "B"]}).with_columns(empty_risk_control_expr())

    assert result.schema[RISK_CONTROL_COLUMN] == pl.List(pl.String)
    assert result[RISK_CONTROL_COLUMN].to_list() == [[], []]


def test_realtime_stk_screen_returns_risk_control_and_drops_delist_date():
    frame = (
        _screen_frame()
        .lazy()
        .with_columns(stock_fund_risk_control_expr())
        .drop("delist_date")
    )

    result = stk_screen(
        frame,
        date(2026, 6, 10),
        [_screen_parameters()],
        builtin_factors=["rank_alpha"],
        fixed_columns=realtime_risk_control_fixed_columns(STK_SCREEN_RESULT_FIXED_COLUMNS),
    )[0]

    assert RISK_CONTROL_COLUMN in result.columns
    assert "delist_date" not in result.columns
    assert result.sort("code")[RISK_CONTROL_COLUMN].to_list() == [
        ["已公告于2026-07-01退市", "已公告退市"],
        [],
    ]


def test_realtime_etf_screen_returns_multiple_risk_items_and_drops_delist_date():
    frame = (
        _screen_frame()
        .lazy()
        .with_columns(stock_fund_risk_control_expr())
        .drop("delist_date")
    )

    result = etf_screen(
        frame,
        date(2026, 6, 10),
        [_screen_parameters()],
        builtin_factors=["rank_alpha"],
        fixed_columns=realtime_risk_control_fixed_columns(ETF_SCREEN_RESULT_FIXED_COLUMNS),
    )[0]

    assert RISK_CONTROL_COLUMN in result.columns
    assert "delist_date" not in result.columns
    assert result.sort("code")[RISK_CONTROL_COLUMN].to_list() == [
        ["已公告于2026-07-01退市", "已公告退市"],
        [],
    ]


def test_historical_stk_screen_does_not_return_risk_control_by_default():
    result = stk_screen(
        _screen_frame().lazy().drop("delist_date"),
        date(2026, 6, 10),
        [_screen_parameters()],
        builtin_factors=["rank_alpha"],
    )[0]

    assert RISK_CONTROL_COLUMN not in result.columns
    assert "delist_date" not in result.columns


def test_realtime_cb_screen_returns_columnar_name_risk_control():
    frame = (
        _screen_frame()
        .with_columns(
            pl.when(pl.col("code") == "A")
            .then(pl.lit("Z甬矽转"))
            .otherwise(pl.lit("正常转债"))
            .alias("name"),
        )
        .lazy()
        .drop("delist_date")
        .with_columns(cb_risk_control_expr())
    )

    result = cb_screen(
        frame,
        date(2026, 6, 10),
        [_screen_parameters()],
        builtin_factors=["rank_alpha"],
        fixed_columns=realtime_risk_control_fixed_columns(CB_SCREEN_RESULT_FIXED_COLUMNS),
    )[0]

    assert result.sort("code")[RISK_CONTROL_COLUMN].to_list() == [["退市前最后交易日"], []]
    assert result.sort("code")["name"].to_list() == ["Z甬矽转", "正常转债"]


def test_etf_quotation_model_keeps_delist_date_from_pre_quotation():
    assert "delist_date" in ETFQuotationModel.columns
