# src/data_processing.py

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

# 导入我们刚刚定义的 Bar 类
from .core_data import Bar 

def load_raw_data(file_path: str) -> pd.DataFrame:
    """
    从 CSV 文件加载原始数据，并进行初步的数据类型处理。
    假设 datetime 列已经是北京时间，无需额外时区转换或本地化。

    Args:
        file_path (str): CSV 文件的完整路径。

    Returns:
        pd.DataFrame: 包含已处理时间戳的原始数据DataFrame。
                      时间戳作为索引，列包括 'open', 'high', 'low', 'close', 
                      'volume', 'open_oi', 'close_oi'。
    
    Raises:
        FileNotFoundError: 如果文件不存在。
        KeyError: 如果CSV中缺少必要的列。
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"数据文件未找到: {file_path}")

    # 定义期望的列名，用于检查和选择
    expected_cols = ['datetime', 'open', 'high', 'low', 'close', 'volume', 'open_oi', 'close_oi', 'underlying_symbol']
    
    try:
        # 使用 pandas.read_csv 直接解析 datetime 列
        # 'datetime' 列的格式是 'YYYY-MM-DD HH:MM:SS'，pandas 能够很好地自动识别
        df = pd.read_csv(
            file_path,
            index_col='datetime',       # 将 datetime 列设置为索引
            parse_dates=True            # 自动解析索引为 datetime 类型
        )
        
        # 检查所有必需的列是否存在
        missing_cols = [col for col in expected_cols[1:] if col not in df.columns]
        if missing_cols:
            print(f"CSV文件中缺少以下列: {', '.join(missing_cols)}")
        expected_cols = [col for col in expected_cols if col in df.columns]

        # 确保数据按时间排序 (这是回测的基础)
        df = df.sort_index()

        print(f"数据加载成功: {file_path}")
        print(f"数据范围从 {df.index.min()} 到 {df.index.max()}")
        print(f"总计 {len(df)} 条记录。")

        return df[expected_cols] # 返回包含核心数据的DataFrame
    except Exception as e:
        print(f"加载数据时发生错误: {e}")
        raise

def df_to_bar_stream(df: pd.DataFrame, symbol: str) -> Iterator[Bar]:
    """
    将 Pandas DataFrame 转换为 Bar 对象的迭代器（数据流）。
    这符合函数式编程的理念，按需生成，不一次性加载所有Bar到内存。

    Args:
        df (pd.DataFrame): 包含 K 线数据的 DataFrame。

    Yields:
        Bar: 逐个生成的 Bar 对象。
    """
    print("开始将 DataFrame 转换为 Bar 对象流...")
    for index, row in df.iterrows():
        try:
            if 'underlying_symbol' in df.columns and row['underlying_symbol'] != '':
                bar = Bar(
                    symbol=row['underlying_symbol'],
                    datetime=index,
                    open=row['open'],
                    high=row['high'],
                    low=row['low'],
                    close=row['close'],
                    volume=int(row['volume']),      # 确保成交量为整数
                    open_oi=int(row['open_oi']),    # 确保持仓量为整数
                    close_oi=int(row['close_oi'])
                )
            else:
                bar = Bar(
                    symbol=symbol,
                    datetime=index,
                    open=row['open'],
                    high=row['high'],
                    low=row['low'],
                    close=row['close'],
                    volume=int(row['volume']),      # 确保成交量为整数
                    open_oi=int(row['open_oi']),    # 确保持仓量为整数
                    close_oi=int(row['close_oi'])
                )
            yield bar
        except (ValueError, TypeError) as e:
            print(f"警告: 无法为 {index} 时间创建Bar对象，跳过。错误: {e}")
            continue
    print("Bar 对象流生成完毕。")

