import polars as pl


RISK_CONTROL_COLUMN = "risk_control"
DELIST_DATE_COLUMN = "delist_date"
NAME_COLUMN = "name"


def empty_risk_control_expr() -> pl.Expr:
    return pl.lit([], dtype=pl.List(pl.String)).alias(RISK_CONTROL_COLUMN)


def _risk_control_list_expr(risk_items: list[pl.Expr]) -> pl.Expr:
    return pl.concat_list(risk_items).list.drop_nulls().alias(RISK_CONTROL_COLUMN)


def delist_date_risk_item_expr() -> pl.Expr:
    return pl.when(pl.col(DELIST_DATE_COLUMN).is_not_null()).then(
        pl.concat_str(
            [
                pl.lit("已公告于"),
                pl.col(DELIST_DATE_COLUMN).dt.strftime("%Y-%m-%d"),
                pl.lit("退市"),
            ]
        )
    ).otherwise(None)


def name_contains_delist_risk_item_expr() -> pl.Expr:
    return (
        pl.when(pl.col(NAME_COLUMN).str.contains("退", literal=True).fill_null(False))
        .then(pl.lit("已公告退市"))
        .otherwise(None)
    )


def name_contains_z_risk_item_expr(name_column: str = NAME_COLUMN) -> pl.Expr:
    return (
        pl.when(pl.col(name_column).str.contains("Z", literal=True).fill_null(False))
        .then(pl.lit("退市前最后交易日"))
        .otherwise(None)
    )


def stock_fund_risk_control_expr(has_delist_date: bool = True) -> pl.Expr:
    risk_items = [name_contains_delist_risk_item_expr()]
    if has_delist_date:
        risk_items.insert(0, delist_date_risk_item_expr())
    return _risk_control_list_expr(risk_items)


def cb_risk_control_expr(name_column: str = NAME_COLUMN) -> pl.Expr:
    return _risk_control_list_expr([name_contains_z_risk_item_expr(name_column)])


def delist_risk_control_expr() -> pl.Expr:
    return _risk_control_list_expr([delist_date_risk_item_expr()])


def stock_fund_risk_control_source_columns(has_delist_date: bool) -> list[str]:
    columns = [NAME_COLUMN]
    if has_delist_date:
        columns.append(DELIST_DATE_COLUMN)
    return columns


def cb_risk_control_source_columns(name_column: str = NAME_COLUMN) -> list[str]:
    return [name_column]


def realtime_risk_control_fixed_columns(fixed_columns: list[str]) -> list[str]:
    return [*fixed_columns, RISK_CONTROL_COLUMN]


def count_risk_control_rows(result: pl.DataFrame) -> int:
    if RISK_CONTROL_COLUMN not in result.columns:
        return 0
    return result.select(
        pl.col(RISK_CONTROL_COLUMN).list.len().gt(0).sum()
    ).item()
