from datetime import date, time

import easyquotation
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 ETF_PRE_QUOTATION, ETF_QUOTATION
from ldnode.loader.datasources.sina import SinaWithTimeoutPatch


class Model(pt.Model):
    trade_date: date # 交易日期
    time: time # 交易时间
    code: str # 标的代码
    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 # 是否为最后一个交易日


class ETFQuotationsDataSource(DataSource):
    def load(self, key: str) -> list[SyncResult]:
        """ 加载 ETF 实时截面数据 """
        try:
            if not ETF_PRE_QUOTATION.get_local_path().exists():
                return [SyncResult(key, SyncStatus.FAILED, "数据源不存在")]
            quotation = SinaWithTimeoutPatch()  # 使用新浪数据源
            # 读取股票盘前数据
            pre_quotations = pl.scan_parquet(ETF_PRE_QUOTATION.get_local_path())
            # 下载行情数据
            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"), # 上市交易日天数
                    (pl.col("trade_date") - pl.col("list_date")).dt.total_days().cast(pl.Int16).alias("list_days"),  # 上市自然日天数
                    (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) # 数据验证
            )
            # 保存ETF实时截面数据
            quotations.write_parquet(ETF_QUOTATION.get_local_path())
            return [SyncResult(key, SyncStatus.SUCCESS, date.today().strftime("%Y-%m-%d"))]
        except Exception as e:
            logger.loader.error(f"加载 ETF 实时截面数据失败: {e}")
            return [SyncResult(key, SyncStatus.FAILED, str(e))]
