"""
使用 Tushare 获取 AU9999 黄金数据
"""

import os
import pandas as pd
import tushare as ts
from datetime import datetime, timedelta
from typing import Optional
from dotenv import load_dotenv


def get_tushare_token() -> str:
    """从环境变量获取 Tushare token"""
    load_dotenv()
    token = os.environ.get('TUSHARE_TOKEN')
    if not token:
        raise ValueError("未设置 TUSHARE_TOKEN 环境变量")
    return token


def fetch_au9999(start_date: str, end_date: str) -> Optional[pd.DataFrame]:
    """
    获取 AU9999 黄金现货数据（使用上期所黄金主力合约 AU.SHF）
    
    Args:
        start_date: 开始日期 (YYYY-MM-DD)
        end_date: 结束日期 (YYYY-MM-DD)
    
    Returns:
        DataFrame with columns: date, open, high, low, close, vol
    """
    try:
        # 初始化 Tushare
        pro = ts.pro_api(get_tushare_token())
        
        # 转换日期格式
        start_str = start_date.replace('-', '')
        end_str = end_date.replace('-', '')
        
        print(f"从 Tushare 获取 AU9999 数据...")
        print(f"时间范围: {start_date} ~ {end_date}")
        
        # 获取黄金期货主力合约数据
        # ts_code: AU.SHF 表示上海期货交易所黄金主力合约
        df = pro.fut_daily(ts_code='AU.SHF', start_date=start_str, end_date=end_str)
        
        if df is None or df.empty:
            print("未获取到数据")
            return None
        
        # 标准化列名
        df = df.rename(columns={
            'trade_date': 'date',
            'open': 'open',
            'high': 'high',
            'low': 'low',
            'close': 'close',
            'vol': 'volume',
        })
        
        # 转换日期格式
        df['date'] = pd.to_datetime(df['date'])
        df = df.set_index('date')
        df = df.sort_index()
        
        # 选择需要的列
        df = df[['open', 'high', 'low', 'close', 'volume']]
        
        print(f"✓ 获取成功: {len(df)} 条数据")
        print(f"时间范围: {df.index[0].strftime('%Y-%m-%d')} ~ {df.index[-1].strftime('%Y-%m-%d')}")
        
        return df
        
    except Exception as e:
        print(f"获取数据失败: {e}")
        return None


def print_au9999_data(df: pd.DataFrame):
    """打印 AU9999 数据"""
    print("\n" + "="*80)
    print("AU9999 黄金数据 (上期所主力合约 AU.SHF)")
    print("="*80)
    print(f"{'日期':<15} {'开盘价':>12} {'最高价':>12} {'最低价':>12} {'收盘价':>12} {'成交量':>12}")
    print("-"*80)
    
    for date, row in df.iterrows():
        date_str = date.strftime('%Y-%m-%d')
        print(f"{date_str:<15} {row['open']:>12.2f} {row['high']:>12.2f} {row['low']:>12.2f} {row['close']:>12.2f} {row['volume']:>12.2f}")
    
    print("="*80)


def main():
    """主函数"""
    # 计算最近30天的日期范围
    end_date = datetime.now()
    start_date = end_date - timedelta(days=30)
    
    start_str = start_date.strftime('%Y-%m-%d')
    end_str = end_date.strftime('%Y-%m-%d')
    
    # 获取数据
    df = fetch_au9999(start_str, end_str)
    
    if df is not None:
        print_au9999_data(df)
        
        # 保存到CSV
        output_file = "au9999_data.csv"
        df.to_csv(output_file)
        print(f"\n数据已保存: {output_file}")
    else:
        print("获取数据失败")


if __name__ == "__main__":
    main()