from datetime import date

import polars as pl

from ldnode.calculator.operators import etf_screen
from ldnode.calculator.schemas import (
    ETF_BUILTIN_FACTORS,
    HoldRangeType,
    HoldWeight,
    SortOrder,
)


def _fund_screen_frame() -> pl.DataFrame:
    rows = [
        {
            "code": "159001.SZ",
            "trade_date": date(2026, 6, 10),
            "fund_type": "ETF",
            "name": "fund-etf",
            "pre_close": 10.0,
            "open": 10.0,
            "high": 10.2,
            "low": 9.9,
            "close": 10.1,
            "vol": 1000,
            "amount": 10000.0,
            "pct_chg": 0.01,
            "open_pct_chg": 0.0,
            "high_pct_chg": 0.02,
            "low_pct_chg": -0.01,
            "suspension": False,
            "limit": 0,
            "list_trade_days": 100,
            "list_days": 150,
            "adj_factor": 1.0,
            "adj_back": 1.0,
            "cash_div": 0.0,
        },
        {
            "code": "160001.SZ",
            "trade_date": date(2026, 6, 10),
            "fund_type": "LOF",
            "name": "fund-lof",
            "pre_close": 9.0,
            "open": 9.0,
            "high": 9.2,
            "low": 8.9,
            "close": 9.3,
            "vol": 900,
            "amount": 9000.0,
            "pct_chg": 0.033,
            "open_pct_chg": 0.0,
            "high_pct_chg": 0.022,
            "low_pct_chg": -0.011,
            "suspension": False,
            "limit": 0,
            "list_trade_days": 90,
            "list_days": 130,
            "adj_factor": 1.0,
            "adj_back": 1.0,
            "cash_div": 0.0,
        },
        {
            "code": "508001.SH",
            "trade_date": date(2026, 6, 10),
            "fund_type": "REITs",
            "name": "fund-reits",
            "pre_close": 8.0,
            "open": 8.0,
            "high": 8.2,
            "low": 7.9,
            "close": 8.5,
            "vol": 800,
            "amount": 8000.0,
            "pct_chg": 0.0625,
            "open_pct_chg": 0.0,
            "high_pct_chg": 0.025,
            "low_pct_chg": -0.0125,
            "suspension": False,
            "limit": 0,
            "list_trade_days": 80,
            "list_days": 120,
            "adj_factor": 1.0,
            "adj_back": 1.0,
            "cash_div": 0.0,
        },
    ]
    return pl.DataFrame(rows)


def _screen_parameters():
    return {
        "computed_factors": [],
        "filter_factors": {
            "relation": "OR",
            "conditions": [
                {
                    "factor": "fund_type",
                    "type": 0,
                    "operator": "eq",
                    "value": "ETF",
                }
            ],
        },
        "rank_factors": [
            {
                "factor": "close",
                "type": 0,
                "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 test_etf_builtin_factors_keep_fund_type_for_exclude_factors():
    assert "fund_type" in ETF_BUILTIN_FACTORS


def test_etf_screen_excludes_fund_type_after_builtin_factor_projection():
    frame = (
        _fund_screen_frame()
        .lazy()
        .select(ETF_BUILTIN_FACTORS)
        .with_columns(pl.lit(None).cast(pl.String).alias("time"))
    )

    result = etf_screen(
        frame,
        date(2026, 6, 10),
        [_screen_parameters()],
        builtin_factors=ETF_BUILTIN_FACTORS,
    )[0]

    assert "fund_type" in result.columns
    assert result["fund_type"].to_list() == ["REITs", "LOF"]
    assert "ETF" not in result["fund_type"].to_list()
