from datetime import date, time

import polars as pl
import patito as pt

from ldnode import logger
from ldnode.loader.schemas.models import DataSource, SyncResult
from ldnode.loader.schemas.enumerations import SyncStatus
from ldnode.loader.files import FUND_PRE_QUOTATION, FUND_QUOTATION
from ldnode.loader.datasources.sina import SinaWithTimeoutPatch


class Model(pt.Model):
    trade_date: date # 交易日期
    time: time # 交易时间
    code: str # 标的代码
    fund_type: str # 基金类型
    delist_date: date | None # 退市公告日期
    name: str # 标的名称
    pre_close: float # 前收盘价(元)
    open: float # 开盘价(元)
    high: float # 最高价(元)
    low: float # 最低价(元)
    close: float # 收盘价(元)
    vol: int # 成交量(份)
    amount: float # 成交额(元)
    suspension: bool # 停牌标记

    cash_div: float = pt.Field(dtype=pl.Float32) # 现金分红(元/份)
    stk_div: float = pt.Field(dtype=pl.Float32) # 送股分红
    adj_factor: float # 复权因子
    adj_back: float # 后复权因子
    list_trade_days: int = pt.Field(dtype=pl.Int16) # 上市以来交易日天数
    
    list_days: int = pt.Field(dtype=pl.Int16) # 上市以来自然日天数
    pct_chg: float # 涨跌幅
    open_pct_chg: float # 开盘涨跌幅
    high_pct_chg: float # 最高涨跌幅
    low_pct_chg: float # 最低涨跌幅
    limit: int = pt.Field(dtype=pl.Int8) # 涨跌停标记
    
    last_trade_day: bool # 是否为最后一个交易日


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


class ETFQuotationsDataSource(DataSource):
    def load(self, key: str) -> list[SyncResult]:
        """ 加载基金实时截面数据 """
        try:
            logger.loader.info(
                f"[Fund] 实时行情计算开始: output_key={key}, pre_quote_key={FUND_PRE_QUOTATION.key}, "
                f"pre_quote_path={FUND_PRE_QUOTATION.get_local_path()}"
            )
            if not FUND_PRE_QUOTATION.get_local_path().exists():
                logger.loader.warning(
                    f"[Fund] 实时行情计算缺少盘前数据: pre_quote_key={FUND_PRE_QUOTATION.key}, "
                    f"pre_quote_path={FUND_PRE_QUOTATION.get_local_path()}"
                )
                return [SyncResult(key, SyncStatus.FAILED, "数据源不存在")]
            quotation = SinaWithTimeoutPatch()  # 使用新浪数据源
            pre_quotations = pl.scan_parquet(FUND_PRE_QUOTATION.get_local_path())
            pre_quote_schema = pre_quotations.collect_schema()
            logger.loader.info(
                f"[Fund] 盘前数据已读取: pre_quote_key={FUND_PRE_QUOTATION.key}, "
                f"columns={pre_quote_schema.names()}"
            )
            # 下载行情数据
            quotation_data = quotation.stocks(
                [*map(
                    lambda code: code.split('.')[1].lower() + code.split('.')[0], 
                    pre_quotations.select(pl.col('code')).collect().to_series(),
                )],
                prefix=True
            )
            quotations = (
                pl.DataFrame(quotation_data)
                .transpose(include_header=True, header_name='code').lazy()
                .unnest('column_0')
                .rename({
                    'close': 'pre_close',
                    'turnover': 'vol',
                    'volume': 'amount',
                    'date': 'trade_date'
                })
                .with_columns(
                    pl.col('code').str.slice(2) + '.' + pl.col('code').str.slice(0, 2).str.to_uppercase(),
                    pl.col('trade_date').str.to_date(),
                    pl.col('time').str.to_time(),
                    pl.when((pl.col('bid1')==0) | (pl.col('ask1')==0)) # 修正涨跌停
                    .then(pl.col('bid1') + pl.col('ask1'))
                    .otherwise((pl.col('bid1') + pl.col('ask1')) / 2)
                    .alias('close'),
                )
                .with_columns(
                    # 修正无挂单close
                    pl.when(pl.col('close')==0)
                    .then(
                        pl.when(pl.col('now')==0).then(pl.col('pre_close'))# 开盘无挂单修正
                        .otherwise(pl.col('now'))
                    ).otherwise(pl.col('close')).alias('close'),
                )
                .join(
                    pre_quotations.rename({'close': '_pre_close'}), 
                    on='code', how='right'
                )
                .with_columns(
                    # 涨跌幅指标
                    (pl.col('close') / pl.col('pre_close') - 1).alias('pct_chg'),  # 涨跌幅
                    (pl.col("open") / pl.col("pre_close") - 1).alias("open_pct_chg"),  # 开盘涨跌幅
                    (pl.col("high") / pl.col("pre_close") - 1).alias("high_pct_chg"),  # 最高涨跌幅
                    (pl.col("low") / pl.col("pre_close") - 1).alias("low_pct_chg"),  # 最低涨跌幅
                    # 其他指标
                    (pl.col('_pre_close') - pl.col('pre_close')).fill_null(0).cast(pl.Float32).alias('cash_div'), # 现金分红(元/份)
                    pl.lit(0).cast(pl.Float32).alias('stk_div'), # 送股分红(元/份)
                    (pl.col('list_trade_days') + 1).fill_null(1).alias("list_trade_days"), # 上市交易日天数
                    fund_realtime_list_days_expr(),  # 上市自然日天数
                    (pl.col('adj_factor') * pl.col('close') / pl.col('pre_close')).fill_null(1).alias('adj_factor'), # 复权因子
                    (pl.col('_pre_close') / pl.col('pre_close') * pl.col('adj_back')).fill_null(1).alias('adj_back'), # 后复权因子
                    # 停牌标记
                    pl.when(
                        (pl.col("open") == 0) & (pl.col("bid1") == 0) & (pl.col("ask1") == 0)
                    ).then(True).otherwise(False).alias('suspension'),
                    pl.lit(0).cast(pl.Int8).alias('limit'), # 涨跌停标记
                    pl.lit(False).alias('last_trade_day'), # 是否为最后一个交易日
                )
                .select(Model.columns) # 保留所需列
                .collect() # 生成dataframe
                .pipe(Model.validate) # 数据验证
            )
            quotations.write_parquet(FUND_QUOTATION.get_local_path())
            fund_type_counts = quotations.group_by("fund_type").len().to_dicts()
            logger.loader.info(
                f"[Fund] 实时行情计算完成: output_key={FUND_QUOTATION.key}, "
                f"output_path={FUND_QUOTATION.get_local_path()}, rows={quotations.height}, "
                f"fund_type_counts={fund_type_counts}"
            )
            return [SyncResult(key, SyncStatus.SUCCESS, date.today().strftime("%Y-%m-%d"))]
        except Exception as e:
            logger.loader.error(
                f"[Fund] 实时行情计算失败: output_key={key}, pre_quote_key={FUND_PRE_QUOTATION.key}, error={e}"
            )
            return [SyncResult(key, SyncStatus.FAILED, str(e))]
