# src/data_manager.py (修改并添加 get_history_bars 方法)

from datetime import datetime
import pandas as pd
from typing import Iterator, List, Dict, Any, Optional
import os

# 导入我们刚刚定义的 Bar 类
from .core_data import Bar
from .data_processing import load_raw_data, df_to_bar_stream


class DataManager:  # DataManager 现在是一个类，以便维护内部索引和提供历史数据
    """
    负责从外部数据源加载数据，并将其转换为统一的、可供回测引擎使用的 Bar 对象。
    并提供获取历史Bar的能力。
    """

    def __init__(self, file_path: str, symbol: str, tz='Asia/Shanghai', start_time: Optional[datetime] = None, end_time: Optional[datetime] = None):
        self.file_path = file_path
        self.tz = tz
        self.raw_df = load_raw_data(self.file_path)  # 调用函数式加载数据

        self.start_time = start_time
        self.end_time = end_time
        if start_time is not None and end_time is not None:
            self.raw_df = self.raw_df[(self.raw_df.index >= start_time) & (self.raw_df.index <= end_time)]

        # self.bars = list(df_to_bar_stream(self.raw_df)) # 一次性转换所有Bar，方便历史数据查找
        # 优化：使用内部迭代器和缓存，避免一次性生成所有Bar，但可以按需提供历史
        self._bar_generator = df_to_bar_stream(self.raw_df, symbol)
        self._bars_cache: List[Bar] = []  # 缓存已经处理过的Bar
        self._last_bar_index = -1  # 用于跟踪当前的Bar在原始df中的索引
        self.symbol = symbol



    def get_next_bar(self) -> Optional[Bar]:
        """
        按顺序返回下一根 Bar 对象。
        """
        try:
            next_bar = next(self._bar_generator)
            self._bars_cache.append(next_bar)
            self._last_bar_index += 1
            return next_bar
        except StopIteration:
            return None  # 数据已读完

    def get_history_bars(self, current_bar_datetime: pd.Timestamp, num_bars: int) -> List[Bar]:
        """
        获取当前Bar之前的历史Bar列表。
        严格防止未来函数：只提供时间戳在 current_bar_datetime 之前的 Bar。

        Args:
            current_bar_datetime (pd.Timestamp): 当前正在处理的Bar的时间戳。
                                                  用于确保获取的历史数据不包含当前及未来的Bar。
            num_bars (int): 需要获取的历史Bar的数量。

        Returns:
            List[Bar]: 历史Bar对象的列表，按时间升序排列。
        """
        # 查找当前 Bar 在缓存中的位置
        # 这是一个查找操作，效率可能不高，但对于简易回测可以接受。
        # 更优方案可以是直接传递当前Bar在DataFrame中的索引。

        # 过滤掉日期时间大于等于 current_bar_datetime 的 Bar
        # 并只取最新的 num_bars 根

        # 确保历史数据是严格在 current_bar_datetime 之前的
        # 考虑到DataFrame索引已经排序，可以直接通过切片查找
        # 找到 current_bar_datetime 在 DataFrame 索引中的位置
        try:
            # 获取所有在当前时间点之前的历史Bar
            # 这里我们依赖raw_df的索引，假设它与_bars_cache是同步的
            # 这种方式更严格地避免未来函数，因为我们直接从原始DataFrame取历史
            # find_loc = self.raw_df.index.get_loc(current_bar_datetime)
            # if isinstance(find_loc, slice): # 处理重复时间戳
            #     find_loc = find_loc.start

            # 使用列表缓存来获取历史数据，避免每次都转换DataFrame
            # 必须确保 _bars_cache 已经包含了 current_bar_datetime 之前的Bar

            # 从缓存中取出所有在 current_bar_datetime 之前的Bar
            # 找到第一个大于或等于 current_bar_datetime 的 Bar 的索引
            end_idx = -1
            for idx, bar in enumerate(self._bars_cache):
                if bar.datetime >= current_bar_datetime:
                    end_idx = idx
                    break
            if end_idx == -1:  # 如果当前bar是最后一个，则所有缓存都是历史
                historical_data = self._bars_cache[:]
            else:
                historical_data = self._bars_cache[:end_idx]

            # 从筛选出的历史数据中，取出最近的 num_bars 根
            return historical_data[-num_bars:]
        except KeyError:
            # 如果 current_bar_datetime 不在索引中，可能还没到该时间点
            return self._bars_cache[-num_bars:] if self._bars_cache else []
        except Exception as e:
            print(f"获取历史Bar时发生错误: {e}")
            return []

    def reset(self):
        """重置数据管理器，以便重新运行回测。"""
        self._bar_generator = df_to_bar_stream(self.raw_df, self.symbol)
        self._bars_cache = []
        self._last_bar_index = -1
        print("DataManager 已重置。")