import pandas as pd

from .data_source import get_tushare_pro

# 导入当站目录的config文件
try:
    # 尝试相对导入（作为包的一部分）
    from .config import TS_TOKEN, START_DATE, END_DATE 
    from .stock_utils import dataCorrect
except (ImportError, SystemError):  
    # 失败则使用绝对导入（直接运行脚本）
    from config import TS_TOKEN, START_DATE, END_DATE 
    from stock_utils import dataCorrect

def getStockParam(TS_CODE,START_DATE=START_DATE,END_DATE=END_DATE):
    pro = get_tushare_pro()
    """
    从tushare daily_basic接口获取单只股票的所有基础数据,并进行数据修正。

    Parameters:
        TS_CODE (str): 股票代码,格式为 '股票代码.SZ' 或 '股票代码.SH',例如 '000001.SZ'

    Returns:
        pd.DataFrame: 包含股票基础数据的DataFrame,字段说明如下：
            ts_code: 股票代码
            trade_date: 交易日期
            close: 收盘价
            turnover_rate: 换手率（%）
            turnover_rate_f: 换手率（自由流通股）
            volume_ratio: 量比
            pe: 市盈率（总市值/净利润,亏损的PE为空）
            pe_ttm: 市盈率（TTM,亏损的PE为空）
            pb: 市净率（总市值/净资产）
            ps: 市销率（总市值/营业收入）
            ps_ttm: 市销率（TTM）
            dv_ratio: 股息率（%）
            dv_ttm: 股息率（TTM）
            total_share: 总股本（万股）
            float_share: 流通股本（万股）
            free_share: 自由流通股本（万股）
            total_mv: 总市值（万元）
            circ_mv: 流通市值（万元）

    Raises:
        Exception: 如果从Tushare接口获取数据时发生错误。
    """
    try:
        # 获取股票基础数据
        df = pro.daily_basic(ts_code=TS_CODE, start_date=START_DATE, end_date=END_DATE)

        # 需要修正的列
        cols = ['close', 'turnover_rate', 'turnover_rate_f', 'volume_ratio', 'pe', 'pe_ttm',
                'pb', 'ps', 'ps_ttm', 'dv_ratio', 'dv_ttm', 'total_share', 'float_share',
                'free_share', 'total_mv', 'circ_mv']

        # 调用dataCorrect函数进行数据修正
        df = dataCorrect(df, cols)

        return df
    except Exception as e:
        # 捕获并处理异常
        print(f"错误: 获取股票 {TS_CODE} 的基础数据时发生错误: {e}")
        return pd.DataFrame()
if __name__ == "__main__":
    # 简单测试    
    import datetime
    yesterday = (datetime.datetime.now() - datetime.timedelta(days=1)).strftime("%Y%m%d")
    test_df = getStockParam("601398.SH", START_DATE=yesterday, END_DATE=yesterday)
    print(test_df)