from datetime import datetime

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
from .realtime_derived_plan import (
    CB_REALTIME_DERIVED_REQUIRED_COLUMNS,
    plan_cb_realtime_derived_factors,
)


def with_cb_quotation_trade_min_unit_columns(lf: pl.LazyFrame) -> pl.LazyFrame:
    return lf.with_columns(
        col("vol").alias("volume"),
        col("amount").alias("amt"),
    )


def with_cb_quotation_size_min_unit_columns(lf: pl.LazyFrame) -> pl.LazyFrame:
    min_currency_unit = pl.lit(100000000)
    return lf.with_columns(
        (col("remain_size") * min_currency_unit).alias("remaining_size"),
        (col("issue_size") * min_currency_unit).alias("issuance_size"),
        (col("remain_cap") * min_currency_unit).alias("remaining_market_value"),
    )


def prepare_cb_pre_quotations_for_realtime(pre_quotations: pl.DataFrame) -> pl.DataFrame:
    realtime_columns_from_sina = [
        "trade_date",
        "date",
        "time",
        "pre_close",
        "open",
        "high",
        "low",
        "close",
        "vol",
        "volume",
        "amount",
        "amt",
        "bid1",
        "ask1",
        "pre_close_stk",
        "open_stk",
        "high_stk",
        "low_stk",
        "close_stk",
        "pct_chg_stk",
        "vol_stk",
        "amount_stk",
    ]
    return (
        pre_quotations
        .rename({"adj_factor": "_pre_adj_factor"})
        .drop(realtime_columns_from_sina, strict=False)
    )


def with_cb_quotation_adj_factor(lf: pl.LazyFrame) -> pl.LazyFrame:
    return lf.with_columns(
        (col("_pre_adj_factor") * col("close") / col("pre_close")).alias("adj_factor")
    )


def with_cb_quotation_list_days(lf: pl.LazyFrame) -> pl.LazyFrame:
    return lf.with_columns(
        ((col("trade_date") - col("list_date")).dt.total_days() + 1).cast(pl.Int16).alias("list_days")
    )


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()
            # 读取转债盘前数据
            pre_quotations = prepare_cb_pre_quotations_for_realtime(
                pl.read_parquet(CB_PRE_QUOTATION.get_local_path())
            )
            logger.loader.info(
                f"[CBQuotation] 盘前数据已读取: pre_quote_key={CB_PRE_QUOTATION.key}, "
                f"pre_quote_path={CB_PRE_QUOTATION.get_local_path()}, "
                f"columns={pre_quotations.columns}"
            )
            missing_realtime_derived_columns = sorted(
                CB_REALTIME_DERIVED_REQUIRED_COLUMNS - set(pre_quotations.columns)
            )
            logger.loader.info(
                "[CBQuotation] 实时现算原材料检查 "
                f"missing_columns={missing_realtime_derived_columns}"
            )
            missing_pay_calendar_columns = sorted(
                {"pay_interest_day_list"} - set(pre_quotations.columns)
            )
            logger.loader.info(
                "[CBQuotation] 付息日历现金流原材料检查 "
                f"missing_columns={missing_pay_calendar_columns}"
            )
            # 下载行情数据
            ### 计算需要下载的标的代码列表
            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)
            # 计算可转债实时截面数据
            _quotations = (
                pl.DataFrame(quotation_data).transpose(include_header=True, header_name='code').lazy()
                .unnest('column_0')
                .pipe(plan_rename)
                .pipe(plan_correct)
            )
            logger.loader.info("[CBQuotation] 转债名称口径采用新浪行情 name source=name_right fallback=name")
            quotations = (
                pre_quotations.lazy()
                .join(
                    _quotations,
                    left_on='code',
                    right_on='code',
                    how='left'
                )
                .with_columns(
                    pl.coalesce(col('name_right').cast(pl.String), col('name').cast(pl.String)).alias('name')
                )
                .drop('name_right')
                .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))
                )
                .pipe(with_cb_quotation_trade_min_unit_columns)
                .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"),  # 转股价值
                )
                .pipe(with_cb_quotation_adj_factor)
                .pipe(with_cb_quotation_list_days)
                .pipe(plan_cb_realtime_derived_factors, today)
                .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")  # 剩余市值
                )
                .pipe(with_cb_quotation_size_min_unit_columns)
                .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()
            cb_name_contains_z_count = quotations.select(
                col('name').cast(pl.String).str.contains('Z', literal=True).fill_null(False).sum()
            ).item()
            logger.loader.info(
                f"[CBQuotation] 转债名称口径生成完成 rows={quotations.height} "
                f"name_contains_z_count={cb_name_contains_z_count} "
                f"has_name_right={'name_right' in quotations.columns}"
            )
            logger.loader.info(
                "[CBQuotation] 内置字段兼容列生成完成 "
                "columns=volume,amt,remaining_size,issuance_size,remaining_market_value,adj_factor,list_days "
                f"rows={quotations.height}"
            )
            logger.loader.info(
                "[CBQuotation] 实时现算字段生成完成 "
                f"rows={quotations.height} "
                f"pure_value_nulls={quotations['pure_value'].null_count() if 'pure_value' in quotations.columns else 'missing'} "
                f"is_call_nulls={quotations['is_call'].null_count() if 'is_call' in quotations.columns else 'missing'} "
                f"redeem_real_days_nulls={quotations['redeem_real_days'].null_count() if 'redeem_real_days' in quotations.columns else 'missing'} "
                f"redeem_remain_days_nulls={quotations['redeem_remain_days'].null_count() if 'redeem_remain_days' in quotations.columns else 'missing'}"
            )
            # 对齐 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))]
