#:!/usr/bin/env python
#:  -*- coding: utf-8 -*-
__author__ = 'mayanqiong'

from collections import namedtuple
import datetime
from typing import Callable, Tuple, Optional

import numpy
from pandas import DataFrame, Series, Index
from sgqlc.operation import Operation
from tqsdk.backtest import TqBacktest

from tqsdk.datetime import _get_expire_rest_days, _str_to_timestamp_nano
from tqsdk.ins_schema import ins_schema, _add_all_frags
from tqsdk.objs import Quote
from tqsdk.diff import _get_obj
from tqsdk.utils import _generate_uuid
from tqsdk.tafunc import _get_t_series, get_impv, _get_d1, get_delta, get_theta, get_gamma, get_vega, get_rho

"""
这两个类只在 api 中用到，主要为了支持用户异步中 await 
没有继承 Entity 类
"""


async def ensure_quote(api, quote):
    if quote.price_tick > 0 and quote.datetime != "":
        return quote
    async with api.register_update_notify(quote) as update_chan:
        async for _ in update_chan:
            if quote.price_tick > 0 and quote.datetime != "":
                return quote


async def ensure_quote_with_underlying(api, quote):
    await ensure_quote(api, quote)
    if quote.underlying_symbol:
        await ensure_quote(api, quote.underlying_quote)
    return quote


class QuoteList(list):
    """
    请求合约信息和行情信息，self._task 完成时，所有的合约已经收到了合约信息和行情信息
    """

    def __init__(self, api, quotes):
        self._api = api
        list.__init__(self, quotes)
        self._task = api.create_task(self._ensure_quotes(), _caller_api=True)
        for quote in quotes:
            # 为每个 quote 对象创建 _task
            if not hasattr(quote, '_task'):
                quote._task = api.create_task(ensure_quote_with_underlying(api, quote), _caller_api=True)

    async def _ensure_quotes(self):
        await self._api._ensure_symbol_async([q._path[-1] for q in self])
        self._api._auth._has_md_grants([q._path[-1] for q in self])  # 权限检查
        # 发送的请求会请求到所有字段，如果是期权也会请求标的的合约信息
        underlying_symbols = set([q.underlying_symbol for q in self if q.underlying_symbol])
        need_quotes = set([q._path[-1] for q in self]).union(underlying_symbols)
        if need_quotes - self._api._requests["quotes"] != set():
            self._api._requests["quotes"] = self._api._requests["quotes"].union(need_quotes)
            self._api._send_pack({
                "aid": "subscribe_quote",
                "ins_list": ",".join(self._api._requests["quotes"]),
            })
        if all([q.datetime != "" for q in self]):
            return self
        all_quotes = self + [_get_obj(self._api._data, ["quotes", s], self._api._prototype["quotes"]["#"]) for s in underlying_symbols]
        async with self._api.register_update_notify(self) as update_chan:
            async for _ in update_chan:
                if all([q.datetime != "" for q in all_quotes]):
                    return self

    def __await__(self):
        return self._task.__await__()


async def _query_graphql_async(api, query_id, query):
    api._send_pack({
        "aid": "ins_query",
        "query_id": query_id,
        "query": query
    })
    symbols = _get_obj(api._data, ["symbols"])
    async with api.register_update_notify(symbols) as update_chan:
        async for _ in update_chan:
            s = symbols.get(query_id, {})
            if s.get("query") == query:
                break


class SymbolList(list):
    """
    query 系列函数返回对象
    """

    def __init__(self, api, query_id: str, query: str, filter: Callable[[dict], list]):
        self._api = api
        self._query_id = query_id
        self._query = query
        self._filter = filter
        list.__init__(self, [])
        self._task = api.create_task(self._query_graphql(), _caller_api=True)

    async def _query_graphql(self):
        pack = {"query": self._query}
        symbols = _get_obj(self._api._data, ["symbols"])
        query_result = None
        for symbol in symbols.values():
            if symbol.items() >= pack.items():  # 检查是否发送过相同的请求
                query_result = symbol
        if query_result is None:
            await _query_graphql_async(self._api, self._query_id, self._query)
            query_result = symbols.get(self._query_id)
            if isinstance(self._api._backtest, TqBacktest):  # 回测时，清空缓存的请求
                self._api._send_pack({
                    "aid": "ins_query",
                    "query_id": self._query_id,
                    "query": ""
                })
        self += self._filter(query_result)
        return self

    def __await__(self):
        return self._task.__await__()


class SymbolLevelList(namedtuple('SymbolLevel', ['in_money_options', 'at_money_options', 'out_of_money_options'])):
    """
    query 系列函数返回对象
    """

    def __new__(cls, *args, **kwargs):
        return super(SymbolLevelList, cls).__new__(cls, in_money_options=[], at_money_options=[], out_of_money_options=[])

    def __init__(self, api, query_id: str, query: str, filter: Callable[[dict], Tuple[list, list, list]]):
        self._api = api
        self._query_id = query_id
        self._query = query
        self._filter = filter
        self._task = api.create_task(self._query_graphql(), _caller_api=True)

    async def _query_graphql(self):
        pack = {"query": self._query}
        symbols = _get_obj(self._api._data, ["symbols"])
        query_result = None
        for symbol in symbols.values():
            if symbol.items() >= pack.items():  # 检查是否发送过相同的请求
                query_result = symbol
        if query_result is None:
            await _query_graphql_async(self._api, self._query_id, self._query)
            query_result = symbols.get(self._query_id)
            if isinstance(self._api._backtest, TqBacktest):  # 回测时，清空缓存的请求
                self._api._send_pack({
                    "aid": "ins_query",
                    "query_id": self._query_id,
                    "query": ""
                })
        l0, l1, l2 = self._filter(query_result)
        self[0].extend(l0)
        self[1].extend(l1)
        self[2].extend(l2)
        return self

    def __await__(self):
        return self._task.__await__()


class TqDataFrame(DataFrame):

    def __init__(self, api, *args, **kwargs):
        super(TqDataFrame, self).__init__(*args, **kwargs)
        self.__dict__["_api"] = api
        self.__dict__["_task"] = api.create_task(self.async_update(), _caller_api=True)

    async def async_update(self):
        async with self._api.register_update_notify(self) as update_chan:
            async for _ in update_chan:
                if self._api._serials.get(id(self))["init"]:
                    return self

    def __await__(self):
        return self.__dict__["_task"].__await__()


def _get_col_dtype(col):
    if col == "expired":
        return bool
    if col in [
        "price_tick", "volume_multiple", "strike_price", "upper_limit", "lower_limit", "pre_settlement", "pre_close",
        "pre_open_interest", "max_limit_order_volume", "max_market_order_volume", "expire_datetime", "expire_rest_days",
        "delivery_year", "delivery_month", "last_exercise_datetime", "exercise_year", "exercise_month",
    ]:
        return float
    return object


class TqSymbolDataFrame(DataFrame):

    def __init__(self, api, symbol_list, backtest_timestamp, *args, **kwargs):
        self.__dict__["_api"] = api
        self.__dict__["_symbol_list"] = symbol_list
        self.__dict__["_backtest_timestamp"] = backtest_timestamp
        self.__dict__["_columns"] = [
            "ins_class",
            "instrument_id",
            "instrument_name",
            "price_tick",
            "volume_multiple",
            "max_limit_order_volume",
            "max_market_order_volume",
            "underlying_symbol",
            "strike_price",
            "exchange_id",
            "product_id",
            "expired",
            "expire_datetime",
            "expire_rest_days",
            "delivery_year",
            "delivery_month",
            "last_exercise_datetime",
            "exercise_year",
            "exercise_month",
            "option_class",
            "upper_limit",
            "lower_limit",
            "pre_settlement",
            "pre_open_interest",
            "pre_close",
            "trading_time_day",
            "trading_time_night"
        ]
        default_quote = Quote(None)
        data = [{k: (s if k == "instrument_id" else default_quote.get(k, None)) for k in self.__dict__["_columns"]} for s in symbol_list]
        super(TqSymbolDataFrame, self).__init__(data=data, columns=self.__dict__["_columns"], *args, **kwargs)
        self.__dict__["_task"] = api.create_task(self.async_update(), _caller_api=True)

    async def async_update(self):
        query_id = _generate_uuid("PYSDK_api")
        op = Operation(ins_schema.rootQuery)
        variables = {"instrument_id": self.__dict__["_symbol_list"]}
        if self.__dict__["_backtest_timestamp"]:
            variables["timestamp"] = self.__dict__["_backtest_timestamp"]
        query = op.multi_symbol_info(**variables)
        _add_all_frags(query)
        self.__dict__["_api"]._send_pack({
            "aid": "ins_query",
            "query_id": query_id,
            "query": op.__to_graphql__()
        })
        symbols = _get_obj(self.__dict__["_api"]._data, ["symbols"])
        async with self.__dict__["_api"].register_update_notify(symbols) as update_chan:
            async for _ in update_chan:
                query_result = symbols.get(query_id, {})
                if query_result:
                    all_keys = set(self.__dict__["_columns"])
                    all_keys.add('trading_time')
                    quotes = self.__dict__["_api"]._symbols_to_quotes(query_result, keys=all_keys)
                    self._quotes_to_dataframe(quotes)
                    if self.__dict__["_backtest_timestamp"]:
                        # 回测时这些字段应该为 nan
                        self.loc[:, ["upper_limit", "lower_limit", "pre_settlement", "pre_open_interest", "pre_close"]] = numpy.nan
                        # 回测时清空请求，不缓存请求内容
                        self.__dict__["_api"]._send_pack({
                            "aid": "ins_query",
                            "query_id": query_id,
                            "query": ""
                        })
                    return self

    def _get_trading_time(self, quotes, symbol, key):
        v = quotes[symbol].get('trading_time', {'day': [], 'night': []}).get(key, [])
        return v if v else None

    def _quotes_to_dataframe(self, quotes):
        default_quote = Quote(None)
        for col in self.__dict__["_columns"]:
            if col == "expire_rest_days":
                current_dt = self._api._get_current_datetime().timestamp()
                self[col] = [_get_expire_rest_days(quotes[s]['expire_datetime'], current_dt)
                                    if quotes[s].get('expire_datetime') else float('nan')
                                    for s in self.__dict__["_symbol_list"]]
            elif col == "trading_time_day" or col == "trading_time_night":
                k = 'day' if col == "trading_time_day" else 'night'
                self[col] = Series(
                    [self._get_trading_time(quotes, s, k) for s in self.__dict__["_symbol_list"]],
                    dtype=_get_col_dtype(col)
                )
            else:
                self[col] = Series(
                    [quotes[s].get(col, default_quote[col]) for s in self.__dict__["_symbol_list"]],
                    dtype=_get_col_dtype(col)
                )

    def __await__(self):
        return self.__dict__["_task"].__await__()


class TqMdSettlementDataFrame(DataFrame):

    def __init__(self, api, symbol, days, start_dt):
        self.__dict__["_api"] = api
        params = {'days': days, 'symbols': symbol if isinstance(symbol, str) else ",".join(symbol)}
        if start_dt is not None:
            params['start_date'] = start_dt.strftime("%Y%m%d")
        else:
            params['days'] += 1
        self.__dict__["_params"] = params
        self.__dict__["_days"] = days
        self.__dict__["_symbol"] = symbol
        self.__dict__["_columns"] = [
            "datetime",
            "symbol",
            "settlement"
        ]
        super(TqMdSettlementDataFrame, self).__init__(data=[], columns=self.__dict__["_columns"])
        self.__dict__["_task"] = api.create_task(self.async_update(), _caller_api=True)

    async def _get_settlement_data(self, settlement_id):
        # 下载结算价数据，并将数据发回到 api.recv_chan
        session = self.__dict__["_api"]._http_session
        url = "https://md-settlement-system-fc-api.shinnytech.com/mss"
        async with session.get(url, params=self.__dict__["_params"]) as response:
                response.raise_for_status()
                content = await response.json()
                await self.__dict__["_api"]._ws_md_recv_chan.send({
                    "aid": "rtn_data",
                    "data": [{
                        "_settlement": {
                            settlement_id: content
                        },
                        "_settlement_finished": {
                            settlement_id: True
                        }
                    }]
                })

    async def async_update(self):
        await self.__dict__["_api"]._ensure_symbol_async(self.__dict__["_symbol"])
        settlement_id = _generate_uuid("PYSDK_settlement")
        self.__dict__["_api"].create_task(self._get_settlement_data(settlement_id), _caller_api=True)  # 错误会抛给 api 处理
        settlement_finished = _get_obj(self.__dict__["_api"]._data, ["_settlement_finished"])
        async with self.__dict__["_api"].register_update_notify(settlement_finished) as update_chan:
            async for _ in update_chan:
                if not settlement_finished.get(settlement_id, False):
                    continue
                content = self.__dict__["_api"]._data.get("_settlement", {}).get(settlement_id, {})
                data = self._content_to_list(content)
                for i, d in enumerate(data):
                    self.loc[i] = d
                self.sort_values(by=['datetime', 'symbol'], inplace=True, ignore_index=True)
                # 读完数据，清空数据
                await self.__dict__["_api"]._ws_md_recv_chan.send({
                    "aid": "rtn_data",
                    "data": [{
                        "_settlement": {
                            settlement_id: None
                        },
                        "_settlement_finished": {
                            settlement_id: None
                        }
                    }]
                })
                return self


    def _content_to_list(self, content):
        sorted_content = sorted(
            content.items(),
            key=lambda item: item[0],
            reverse=True
        )
        # 取最近的 N 天数据
        days = self.__dict__["_days"]
        recent_content = sorted_content[:days]
        return [
            {"datetime": dt, "symbol": symbol, "settlement": settlement}
            for dt, symbols in recent_content
            for symbol, settlement in symbols.items()
            if settlement is not None 
        ]

    def __await__(self):
        return self.__dict__["_task"].__await__()


class TqSymbolRankingDataFrame(DataFrame):

    def __init__(self, api, symbol, ranking_type, days, start_dt, broker):
        self.__dict__["_api"] = api
        params = {'symbol': symbol}
        if days is not None:
            params['days'] = days
        if start_dt is not None:
            params['start_date'] = start_dt.strftime("%Y%m%d")
        if broker is not None:
            params['broker'] = broker
        self.__dict__["_params"] = params
        self.__dict__["_symbol"] = symbol
        self.__dict__["_ranking_type"] = f"{ranking_type.lower()}_ranking"
        self.__dict__["_columns"] = [
            "datetime",
            "symbol",
            "exchange_id",
            "instrument_id",
            "broker",
            "volume",
            "volume_change",
            "volume_ranking",
            "long_oi",
            "long_change",
            "long_ranking",
            "short_oi",
            "short_change",
            "short_ranking"
        ]
        super(TqSymbolRankingDataFrame, self).__init__(data=[], columns=self.__dict__["_columns"])
        self.__dict__["_task"] = api.create_task(self.async_update(), _caller_api=True)

    async def _get_ranking_data(self, ranking_id):
        # 下载持仓排名数据，并将数据发回到 api.recv_chan
        session = self.__dict__["_api"]._http_session
        url = "https://symbol-ranking-system-fc-api.shinnytech.com/srs"
        async with session.get(url, params=self.__dict__["_params"]) as response:
                response.raise_for_status()
                content = await response.json()
                await self.__dict__["_api"]._ws_md_recv_chan.send({
                    "aid": "rtn_data",
                    "data": [{
                        "_symbol_rankings": {
                            ranking_id: content
                        },
                        "_symbol_rankings_finished": {
                            ranking_id: True
                        }
                    }]
                })

    async def async_update(self):
        await self.__dict__["_api"]._ensure_symbol_async(self.__dict__["_symbol"])
        ranking_id = _generate_uuid("PYSDK_rank")
        self.__dict__["_api"].create_task(self._get_ranking_data(ranking_id), _caller_api=True)  # 错误会抛给 api 处理
        symbol_rankings_finished = _get_obj(self.__dict__["_api"]._data, ["_symbol_rankings_finished"])
        async with self.__dict__["_api"].register_update_notify(symbol_rankings_finished) as update_chan:
            async for _ in update_chan:
                if not symbol_rankings_finished.get(ranking_id, False):
                    continue
                content = self.__dict__["_api"]._data.get("_symbol_rankings", {}).get(ranking_id, {})
                data = self._content_to_list(content)
                for i, d in enumerate(data):
                    self.loc[i] = d
                self.dropna(subset=[self.__dict__["_ranking_type"]], inplace=True)
                self.sort_values(by=['datetime', self.__dict__["_ranking_type"]], inplace=True, ignore_index=True)
                # 读完数据，清空数据
                await self.__dict__["_api"]._ws_md_recv_chan.send({
                    "aid": "rtn_data",
                    "data": [{
                        "_symbol_rankings": {
                            ranking_id: None
                        },
                        "_symbol_rankings_finished": {
                            ranking_id: None
                        }
                    }]
                })
                return self

    def _content_to_list(self, content):
        data = {}
        for dt in content.keys():
            for symbol in content[dt].keys():
                if content[dt][symbol] is None:
                    continue
                for data_type, rankings in content[dt][symbol].items():
                    for broker, rank_item in rankings.items():
                        item = data.setdefault((dt, symbol, broker), self._get_default_item(dt, symbol, broker))
                        if data_type == 'volume_ranking':
                            item['volume'] = rank_item['volume']
                            item['volume_change'] = rank_item['varvolume']
                            item['volume_ranking'] = rank_item['ranking']
                        elif data_type == 'long_ranking':
                            item['long_oi'] = rank_item['volume']
                            item['long_change'] = rank_item['varvolume']
                            item['long_ranking'] = rank_item['ranking']
                        elif data_type == 'short_ranking':
                            item['short_oi'] = rank_item['volume']
                            item['short_change'] = rank_item['varvolume']
                            item['short_ranking'] = rank_item['ranking']
        return data.values()

    def _get_default_item(self, dt, symbol, broker):
        return {
            "datetime": dt,
            "symbol": symbol,
            "exchange_id": symbol.split(".", maxsplit=1)[0],
            "instrument_id": symbol.split(".", maxsplit=1)[1],
            "broker": broker,
            "volume": float('nan'),
            "volume_change": float('nan'),
            "volume_ranking": float('nan'),
            "long_oi": float('nan'),
            "long_change": float('nan'),
            "long_ranking": float('nan'),
            "short_oi": float('nan'),
            "short_change": float('nan'),
            "short_ranking": float('nan')
        }

    def __await__(self):
        return self.__dict__["_task"].__await__()



class TqOptionGreeksDataFrame(DataFrame):

    def __init__(self, api, symbol_list, v_list, r):
        self.__dict__["_api"] = api
        self.__dict__["_symbol_list"] = symbol_list
        self.__dict__["_v_list"] = v_list
        self.__dict__["_r"] = r
        self.__dict__["_columns"] = [
            "instrument_id",
            "instrument_name",
            "option_class",
            "expire_rest_days",
            "expire_datetime",
            "underlying_symbol",
            "strike_price",
            "delta",
            "gamma",
            "theta",
            "vega",
            "rho"
        ]
        super(TqOptionGreeksDataFrame, self).__init__(data=[], columns=self.__dict__["_columns"])
        self.__dict__["_task"] = api.create_task(self.async_update(), _caller_api=True)

    async def async_update(self):
        symbol_list = self.__dict__["_symbol_list"]
        quotes = await self.__dict__["_api"].get_quote_list(symbol_list)
        if not all([q.ins_class.endswith("OPTION") for q in quotes]):
            raise Exception("quote 参数列表中元素必须是期权类型")
        for i, q in enumerate(quotes):
            self.loc[i] = {k: q.get(k, float('nan')) for k in self.__dict__["_columns"]}
        self._get_greeks(quotes)
        return self

    def _get_greeks(self, quotes):
        series_close = Series(data=[q.last_price for q in quotes])  # 期权最新价
        series_close1 = Series(data=[q.underlying_quote.last_price for q in quotes])  # 标的最新价
        series_o = Series(data=[q.option_class for q in quotes])
        series_datetime = Series(data=[_str_to_timestamp_nano(q.datetime) for q in quotes])
        series_expire_datetime = Series(data=[q.expire_datetime for q in quotes])
        series_t = _get_t_series(series_datetime, 0, series_expire_datetime)  # 到期时间
        if self.__dict__["_v_list"] is None:
            series_v = get_impv(series_close1, series_close, self["strike_price"], self.__dict__["_r"], 0.3, series_t, series_o)
        else:
            series_v = Series(data=self.__dict__["_v_list"])
        series_d1 = _get_d1(series_close1, self["strike_price"], self.__dict__["_r"], series_v, series_t)
        self["delta"] = get_delta(series_close1, self["strike_price"], self.__dict__["_r"], series_v, series_t, series_o, series_d1)
        self["theta"] = get_theta(series_close1, self["strike_price"], self.__dict__["_r"], series_v, series_t, series_o, series_d1)
        self["gamma"] = get_gamma(series_close1, self["strike_price"], self.__dict__["_r"], series_v, series_t, series_d1)
        self["vega"] = get_vega(series_close1, self["strike_price"], self.__dict__["_r"], series_v, series_t, series_d1)
        self["rho"] = get_rho(series_close1, self["strike_price"], self.__dict__["_r"], series_v, series_t, series_o, series_d1)

    def __await__(self):
        return self.__dict__["_task"].__await__()


class TqEdbIndexDataFrame(DataFrame):
    """
    EDB 指标取值查询结果（不更新）。
    行索引为日期字符串 YYYY-MM-DD，列为指标 id（去重后保持请求顺序 / 以服务端返回为准）。
    """

    def __init__(
            self, api, ids, start: str = None, end: str = None, align: str = None, fill: str = None,
            raw_df: Optional["TqEdbIndexDataFrame"] = None
    ):
        self.__dict__["_api"] = api
        self.__dict__["_ids"] = ids
        self.__dict__["_start"] = start
        self.__dict__["_end"] = end
        self.__dict__["_align"] = align
        self.__dict__["_fill"] = fill
        self.__dict__["_raw_df"] = raw_df
        payload = {"ids": ids}
        if start is not None:
            payload["start"] = start
        if end is not None:
            payload["end"] = end
        self.__dict__["_payload"] = payload
        self.__dict__["_url"] = "https://edb.shinnytech.com/data/index_data"
        super(TqEdbIndexDataFrame, self).__init__(data=[], columns=list(ids))
        self.index.name = "date"
        self.__dict__["_task"] = api.create_task(self.async_update(), _caller_api=True)

    @staticmethod
    def _apply_align_fill(df: DataFrame, start_s: str, end_s: str, align: Optional[str], fill: Optional[str]) -> DataFrame:
        """
        根据 align/fill 对 index 做补齐/填充。
        注意：回测模式下必须先切片窗口再调用，避免引入未来数据。
        """
        if align != "day":
            return df
        start_d = datetime.datetime.strptime(start_s, "%Y-%m-%d").date()
        end_d = datetime.datetime.strptime(end_s, "%Y-%m-%d").date()
        all_dates = [(start_d + datetime.timedelta(days=i)).strftime("%Y-%m-%d")
                     for i in range((end_d - start_d).days + 1)]
        df2 = df.reindex(Index(all_dates, name=df.index.name))
        if fill == "ffill":
            df2 = df2.infer_objects().ffill()
        elif fill == "bfill":
            df2 = df2.infer_objects().bfill()
        return df2

    async def _get_index_data(self, req_id):
        session = self.__dict__["_api"]._http_session
        async with session.post(self.__dict__["_url"], json=self.__dict__["_payload"]) as response:
            response.raise_for_status()
            content = await response.json()

            if content.get("error_code", 0) != 0:
                raise Exception(f"edb 指标取值查询失败: {content.get('error_msg', '')}")

            # 只下发 data 节点，且将 values 改名为 datas，避免覆盖 Entity.values()
            data = content.get("data", {}) or {}
            payload = {
                "ids": data.get("ids", []),
                "datas": data.get("values", {}) or {}
            }
            await self.__dict__["_api"]._ws_md_recv_chan.send({
                "aid": "rtn_data",
                "data": [{
                    "_edb_index_data": {req_id: payload},
                    "_edb_index_data_finished": {req_id: True}
                }]
            })

    async def async_update(self):
        # window 模式：从 raw_df 切片窗口，再对齐/补全（避免引入未来数据）
        raw_df = self.__dict__.get("_raw_df")
        if raw_df is not None:
            await raw_df
            requested_ids = list(self.__dict__.get("_ids", []))
            start_s = self.__dict__.get("_start")
            end_s = self.__dict__.get("_end")
            align = self.__dict__.get("_align", None)
            fill = self.__dict__.get("_fill", None)

            idx = raw_df.index
            if len(idx) == 0:
                df2 = DataFrame([], index=Index([], name="date"), columns=requested_ids)
            else:
                # YYYY-MM-DD 字符串按字典序与时间序一致
                mask = (idx >= start_s) & (idx <= end_s)
                df2 = raw_df.loc[mask, requested_ids].copy()
            df2 = self._apply_align_fill(df2, start_s=start_s, end_s=end_s, align=align, fill=fill)
            self._update_inplace(df2)
            return self

        # raw 模式：raw_df 为空时走原始下载结构（ws->api._data->finished）
        # 在 query_edb_data 的 raw 缓存场景下，会传入 align=None/fill=None，因此不会发生对齐/补全，原始下载结构不变。
        req_id = _generate_uuid("PYSDK_edb_data")
        self.__dict__["_api"].create_task(self._get_index_data(req_id), _caller_api=True)
        finished = _get_obj(self.__dict__["_api"]._data, ["_edb_index_data_finished"])
        async with self.__dict__["_api"].register_update_notify(finished) as update_chan:
            async for _ in update_chan:
                if not finished.get(req_id, False):
                    continue
                content = self.__dict__["_api"]._data.get("_edb_index_data", {}).get(req_id, {})
                ids_from_server = list(content.get("ids", self.__dict__.get("_ids", [])) or [])
                values = content.get("datas", {}) or {}

                # 列顺序：优先保持用户请求顺序（去重后），避免与服务端返回顺序不一致导致列错位
                requested_ids = list(self.__dict__.get("_ids", []))
                self.columns = requested_ids

                for dt in sorted(values.keys()):
                    row = values.get(dt, None)
                    if row is None:
                        continue
                    row_map = {k: v for k, v in zip(ids_from_server, row)}
                    self.loc[dt] = [row_map.get(i, None) for i in requested_ids]

                # 读完数据，清空缓存（保持原有结构不变）
                await self.__dict__["_api"]._ws_md_recv_chan.send({
                    "aid": "rtn_data",
                    "data": [{
                        "_edb_index_data": {req_id: None},
                        "_edb_index_data_finished": {req_id: None}
                    }]
                })
                return self

    def __await__(self):
        return self.__dict__["_task"].__await__()
