from datetime import date

import polars as pl
from polars import col

from ldnode import logger
from ldnode.loader.schemas.models import DataSource, SyncResult
from ldnode.loader.schemas.enumerations import SyncStatus
from ldnode.loader.files import STK_PRE_QUOTATION, STK_QUOTATION
from ldnode.loader.datasources.common_plans import plan_rename, plan_correct
from ldnode.loader.datasources.common_plans.limit_exprs import (
    third_board_upper_limit_expr,
    third_board_lower_limit_expr,
    second_board_upper_limit_expr,
    second_board_lower_limit_expr,
    main_board_upper_limit_expr,
    main_board_lower_limit_expr,
    main_board_st_upper_limit_expr,
    main_board_st_lower_limit_expr,
)
from ldnode.loader.datasources.sina import SinaWithTimeoutPatch


def stock_realtime_list_days_expr() -> pl.Expr:
    return ((col("trade_date") - col("list_date")).dt.total_days() + 1).cast(pl.Int16).alias("list_days")


class StkQuotationsDataSource(DataSource):
    def load(self, key: str) -> list[SyncResult]:
        """ 生成股票实时截面数据 """
        try:
            if not STK_PRE_QUOTATION.get_local_path().exists():
                return [SyncResult(key, SyncStatus.FAILED, "数据源不存在")]
            # logger.loader.info(f"开始生成股票实时截面数据")
            quotation = SinaWithTimeoutPatch()
            # 读取股票盘前数据
            pre_quotations = pl.read_parquet(STK_PRE_QUOTATION.get_local_path())
            # 下载行情数据
            quotation_data = quotation.stocks(
                [*map(
                    lambda code: code.split('.')[1].lower() + code.split('.')[0], 
                    pre_quotations['code'],
                )],
                prefix=True
            )
            # 计算股票实时截面数据
            lf = pl.DataFrame(quotation_data).transpose(include_header=True, header_name='code').lazy()
            quotations = (
                lf
                .unnest('column_0')
                .pipe(plan_rename)
                .pipe(plan_correct)
                .drop('name', strict=False)  # 新浪名称不带 -U/-W 后缀，丢弃后由盘前数据提供
                .join(
                    pre_quotations.lazy(),
                    on='code',
                    how='right'
                )
                .with_columns(
                    # # 涨跌幅指标
                    (col('close') / col('pre_close') - 1).alias('pct_chg'),  # 涨跌幅
                    (col("open") / col("pre_close") - 1).alias("open_pct_chg"),  # 开盘涨跌幅
                    (col("high") / col("pre_close") - 1).alias("high_pct_chg"),  # 最高涨跌幅
                    (col("low") / col("pre_close") - 1).alias("low_pct_chg"),  # 最低涨跌幅
                    # 市值相关指标
                    (col("float_share") * col("close")).alias("circ_mv"),  # 流通市值（单位：元）
                    (col("total_share") * col("close")).alias("total_mv"),  # 总市值（单位：元）
                    (col("vol") / col("float_share")).alias("turnover"),  # 换手率
                    # 其他指标
                    (col('adj_factor') * col('close') / col('pre_close')).alias('adj_factor'), # 复权因子
                    (col('list_trade_days') + 1).alias("list_trade_days"), # 上市交易日天数
                    stock_realtime_list_days_expr(),  # 上市自然日天数
                    # 计算涨停价格
                    pl.when(col('code').str.ends_with('BJ')).then(
                        third_board_upper_limit_expr # 北交所
                    ).otherwise(
                        pl.when(col('code').str.starts_with('688') | col('code').str.starts_with('30')).then(
                            second_board_upper_limit_expr # 双创
                        ).otherwise(
                            pl.when(col('name').str.to_uppercase().str.contains('ST')).then(
                                main_board_st_upper_limit_expr # 沪深主板ST
                            ).otherwise(
                                main_board_upper_limit_expr # 沪深主板
                            )
                        )
                    ).alias('upper_limit'),
                    # 计算跌停价格
                    pl.when(col('code').str.ends_with('BJ')).then(
                        third_board_lower_limit_expr # 北交所
                    ).otherwise(
                        pl.when(col('code').str.starts_with('688') | col('code').str.starts_with('30')).then(
                            second_board_lower_limit_expr # 双创
                        ).otherwise(
                            pl.when(col('name').str.to_uppercase().str.contains('ST')).then(
                                main_board_st_lower_limit_expr # 沪深主板ST
                            ).otherwise(
                                main_board_lower_limit_expr # 沪深主板
                            )
                        )
                    ).alias('lower_limit'), 
                )
                .with_columns(
                    # 计算涨跌停标志
                    pl.when(col('close') == col('upper_limit')).then(pl.lit(1)).otherwise(
                        pl.when(col('close') == col('lower_limit')).then(pl.lit(-1)).otherwise(pl.lit(0))
                    ).cast(pl.Int8).alias('limit'),
                )
                # 计算财务指标
                .with_columns(
                    (col('total_mv') / col('npatoshopc_ttm')).alias('pe_ttm'), # 市盈率TTM
                    (col('npatoshopc_ttm') / col('total_mv')).alias('ep_ttm'), # 市盈率TTM倒数
                    (col('total_mv') / col('npadnrpatoshaopc_ttm')).alias('d_pe_ttm'), # 扣非市盈率TTM
                    (col('npadnrpatoshaopc_ttm') / col('total_mv')).alias('d_ep_ttm'), # 扣非市盈率TTM倒数
                    (col('total_mv') / col('tetoshopc')).alias('pb'), # 市净率
                    (col('tetoshopc') / col('total_mv')).alias('bp'), # 市净率倒数
                    (col('total_mv') / col('tetoshopc_wo_gw')).alias('pb_wo_gw'), # 扣除商誉市净率
                    (col('tetoshopc_wo_gw') / col('total_mv')).alias('bp_wo_gw'), # 扣除商誉市净率倒数
                    (col('total_mv') / col('toi')).alias('ps'), # 市销率
                    (col('total_mv') / col('toi_ttm')).alias('ps_ttm'), # 市销率TTM
                    (col('total_mv') / col('ncffoa')).alias('pcf'), # 市现率
                    (col('total_mv') / col('ncffoa_ttm')).alias('pcf_ttm'), # 市现率TTM
                    (col('dy_ttm') / col('total_mv')).alias('dyr'), # 股息率
                )
                # 排除停牌
                .filter(
                    ~((col("open") == 0) & (col("bid1") == 0) & (col("ask1") == 0))
                )
                # 清理无用列
                .select(pl.exclude('now'))
            ).collect()
            # 保存股票实时截面数据
            quotations.write_parquet(STK_QUOTATION.get_local_path())
            return [SyncResult(key, SyncStatus.SUCCESS, date.today().strftime("%Y-%m-%d"))]
        except Exception as e:
            logger.loader.error(f"生成股票实时截面数据失败: {e}")
            return [SyncResult(key, SyncStatus.FAILED, str(e))]
