#!/usr/bin/env python3
"""
获取159516 ETF净值数据
"""

import os
import pandas as pd
import tushare as ts
from datetime import datetime, timedelta

# 设置Tushare token
def get_tushare_token():
    # 首先尝试从环境变量获取
    token = os.environ.get("TUSHARE_TOKEN")
    if token:
        return token
    
    # 尝试从.env文件获取
    try:
        from dotenv import load_dotenv
        load_dotenv()
        token = os.environ.get("TUSHARE_TOKEN")
        if token:
            return token
    except ImportError:
        pass
    
    # 手动读取.env文件
    env_path = os.path.join(os.path.dirname(__file__), '.env')
    if os.path.exists(env_path):
        with open(env_path, 'r') as f:
            for line in f:
                if line.startswith('TUSHARE_TOKEN='):
                    token = line.strip().split('=', 1)[1].strip().strip('"').strip("'")
                    if token:
                        return token
    
    raise ValueError("请设置 TUSHARE_TOKEN 环境变量或在.env文件中配置")


def fetch_etf_nav(etf_code="159516.SZ", days=30):
    """
    获取ETF净值数据
    
    Args:
        etf_code: ETF代码，如 "159516.SZ"
        days: 获取天数
    
    Returns:
        DataFrame: 包含日期和净值
    """
    pro = ts.pro_api(get_tushare_token())
    
    # 计算日期范围
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days + 5)
    
    start_str = start_date.strftime('%Y%m%d')
    end_str = end_date.strftime('%Y%m%d')
    
    # 转换代码格式 (tushare使用.SH而不是.SS)
    ts_code = etf_code.replace(".SS", ".SH")
    
    print(f"获取 {etf_code} 净值数据...")
    print(f"日期范围: {start_str} ~ {end_str}")
    
    try:
        # 获取ETF净值数据
        nav_df = pro.fund_nav(
            ts_code=ts_code,
            start_date=start_str,
            end_date=end_str
        )
        
        if nav_df is None or len(nav_df) == 0:
            print("未获取到净值数据")
            return None
        
        # 排序并处理数据
        nav_df = nav_df.sort_values('nav_date')
        
        # 转换日期格式
        nav_df['date'] = pd.to_datetime(nav_df['nav_date'])
        nav_df = nav_df.set_index('date')
        
        print(f"\n获取到 {len(nav_df)} 条净值数据")
        print(f"最新净值日期: {nav_df.index.max().strftime('%Y-%m-%d')}")
        print(f"最新净值: {nav_df['unit_nav'].iloc[-1]}")
        
        # 显示最近10条数据
        print(f"\n最近10条净值数据:")
        print(nav_df[['unit_nav']].tail(10).to_string())
        
        return nav_df
        
    except Exception as e:
        print(f"获取净值数据失败: {e}")
        return None


if __name__ == "__main__":
    # 获取159516的净值数据
    result = fetch_etf_nav("159516.SZ", days=30)
    
    if result is not None:
        # 保存到CSV文件
        output_file = "159516_nav_data.csv"
        result[['unit_nav']].to_csv(output_file)
        print(f"\n数据已保存到: {output_file}")