import time
from datetime import datetime

import humanize
import easyquotation
import polars as pl
from polars import col

from ldnode import logger
from ldnode.loader.schemas.enumerations import SyncStatus
from ldnode.loader.schemas.models import DataSource, SyncResult
from ldnode.loader.files import CB_PRE_QUOTATION, CB_QUOTATION, CB_HIST_DATA, STK_HIST_DATA
from ldnode.loader.datasources.common_plans import plan_rename, plan_correct
from ldnode.loader.datasources.common_plans.limit_exprs import (
    cb_lower_limit_expr, 
    cb_upper_limit_expr,
)
from ldnode.loader.datasources.sina import SinaWithTimeoutPatch
from .ytm_plan import plan_ytm
from .conv_prem_plan import plan_conv_prem
from .history_plan import plan_history
from .option_plan import plan_option


class CBQuotationsDataSource(DataSource):
    def load(self, key: str) -> list[SyncResult]:
        """ 生成可转债实时截面数据 """
        try:
            if not CB_PRE_QUOTATION.get_local_path().exists() or not CB_HIST_DATA.get_local_path().exists() or not STK_HIST_DATA.get_local_path().exists():
                return [SyncResult(key, SyncStatus.FAILED, "数据源不存在")]
            # logger.loader.info(f"开始生成可转债实时截面数据")
            quotation = SinaWithTimeoutPatch()
            today = datetime.now().date()
            start_at = time.time()
            # 读取股票盘前数据
            pre_quotations = pl.read_parquet(CB_PRE_QUOTATION.get_local_path())
            # 下载行情数据
            download_at = time.time()
            ### 计算需要下载的标的代码列表
            codes = [
                *map(
                    lambda code: code.split('.')[1].lower() + code.split('.')[0], 
                    pre_quotations['code'],
                ),
                *map(
                    lambda code: code.split('.')[1].lower() + code.split('.')[0], 
                    pre_quotations['code_stk'],
                ),
            ]
            quotation_data = quotation.stocks(codes, prefix=True)
            # logger.loader.info(f"下载行情数据耗时: {humanize.naturaldelta(time.time() - download_at)}")
            # 计算可转债实时截面数据
            calculate_at = time.time()
            _quotations = (
                pl.DataFrame(quotation_data).transpose(include_header=True, header_name='code').lazy()
                .unnest('column_0')
                .pipe(plan_rename)
                .pipe(plan_correct)
            )
            quotations = (
                pre_quotations.lazy()
                .join(
                    _quotations,
                    left_on='code',
                    right_on='code',
                    how='left'
                )
                .drop('name_stk')
                .join(
                    _quotations,
                    left_on='code_stk',
                    right_on='code',
                    how='left',
                    suffix='_stk',
                )
                .pipe(plan_ytm, today)
                # 排除停牌
                .filter(
                    ~((col("open") == 0) & (col("bid1") == 0) & (col("ask1") == 0))
                )
                .with_columns(
                    (col('close') / col('pre_close') - 1).alias('pct_chg'),  # 转债涨跌幅
                    (col('close_stk') / col('pre_close_stk') - 1).alias('pct_chg_stk'),  # 正股涨跌幅
                    (col("amount") / 10000).alias("amount"),  # 成交额（单位：万）
                    (col("vol") / 10).alias("vol"),  # 成交量（单位：手）
                    (col("amount_stk") / 10000).alias("amount_stk"),  # 成交额（单位：万）
                    (col("vol_stk") / 100).alias("vol_stk"),  # 成交量（单位：手）
                    (100 / col("conv_price") * col("close_stk")).alias("conv_value"),  # 转股价值
                )
                .with_columns(
                    (col("close") / col("conv_value") - 1).alias("conv_prem"),  # 转股溢价率
                    (col("close") / col("pure_value") - 1).alias("bond_prem"),
                )
                .with_columns(  
                    # 纯债溢价率
                    (col("close") + col("conv_prem") * 100).alias("dblow"),  # 双低值
                    (col("pe_ttm") / col("pre_close_stk") * col("close_stk")).alias("pe_ttm"),
                    (col("pb") / col("pre_close_stk") * col("close_stk")).alias("pb"),
                    (col("ps_ttm") / col("pre_close_stk") * col("close_stk")).alias("ps_ttm"),
                    (col("total_mv") / col("pre_close_stk") * col("close_stk") / 1e4).alias("total_mv"),  # 总市值（单位：亿元）
                    (col("circ_mv") / col("pre_close_stk") * col("close_stk") / 1e4).alias("circ_mv"), # 流通市值（单位：亿元）
                    (col("remain_size") / 100 * col("close")).alias("remain_cap")  # 剩余市值
                )
                .with_columns(
                    (col("amount") / col("remain_cap") / 10000).alias("turnover"), # 换手率（注意：amount已被修改）
                    (col("remain_cap") / col("total_mv")).alias("cap_mv_rate"),  # 市占比
                    cb_upper_limit_expr.alias('upper_limit'), # 涨停价格
                    cb_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')
                )
                .pipe(plan_conv_prem)
                .pipe(plan_history)
                .pipe(plan_option)
            ).collect()
            # 对齐 S3 数据范式：S3 上的 cb_quotations.pq 使用 date（Datetime）而非 trade_date（Date），
            # 内部计算管道（plan_history 等）依赖 trade_date 列名，因此在输出前重命名回 date
            if 'trade_date' in quotations.columns:
                quotations = quotations.rename({'trade_date': 'date'})
            # 保存可转债实时截面数据
            quotations.write_parquet(CB_QUOTATION.get_local_path())
            # logger.loader.info(f"生成可转债实时截面数据总耗时: {humanize.naturaldelta(time.time() - start_at)}")
            return [SyncResult(key, SyncStatus.SUCCESS, today.strftime("%Y-%m-%d"))]
        except Exception as e:
            logger.loader.error(f"生成可转债实时截面数据失败: {e}")
            return [SyncResult(key, SyncStatus.FAILED, str(e))]
