from typing import Any


import time
from datetime import date

import easyquotation
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


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()
            start_at = time.time()
            # 读取股票盘前数据
            pre_quotations = pl.read_parquet(STK_PRE_QUOTATION.get_local_path())
            # 下载行情数据
            download_at = time.time()
            quotation_data = quotation.stocks(
                [*map(
                    lambda code: code.split('.')[1].lower() + code.split('.')[0], 
                    pre_quotations['code'],
                )],
                prefix=True
            )
            # logger.loader.info(f"下载行情数据耗时: {humanize.naturaldelta(time.time() - download_at)}")
            # 计算股票实时截面数据
            calculate_at = time.time()
            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"), # 上市交易日天数
                    (col("trade_date") - col("list_date")).dt.total_days().cast(pl.Int16).alias("list_days"),  # 上市自然日天数
                    # 计算涨停价格
                    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())
            # logger.loader.info(f"生成股票实时截面数据总耗时: {humanize.naturaldelta(time.time() - start_at)}")
            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))]
