#!/usr/bin/env python3
"""
获取159930 ETF最新10天的收盘价、净值并计算溢价率
"""

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_data(etf_code: str, days: int = 10):
    """
    获取ETF最新N天的价格、净值数据
    
    Args:
        etf_code: ETF代码，如 "159930.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')
    
    # 转换代码格式
    ts_code = etf_code.replace(".SS", ".SH")
    
    print(f"获取 {etf_code} 数据...")
    print(f"日期范围: {start_str} ~ {end_str}")
    
    # 1. 获取ETF价格数据（fund_daily接口）
    print("\n1. 获取ETF价格数据...")
    try:
        price_df = pro.fund_daily(
            ts_code=ts_code,
            start_date=start_str,
            end_date=end_str
        )
        if price_df is not None and len(price_df) > 0:
            price_df = price_df.sort_values('trade_date')
            print(f"   获取到 {len(price_df)} 条价格数据")
            print(f"   最新日期: {price_df['trade_date'].max()}")
        else:
            print("   未获取到价格数据")
            price_df = None
    except Exception as e:
        print(f"   获取价格数据失败: {e}")
        price_df = None
    
    # 2. 获取ETF净值数据（fund_nav接口）
    print("\n2. 获取ETF净值数据...")
    try:
        # 净值通常滞后，多取一天
        nav_end_date = end_date + timedelta(days=1)
        nav_end_str = nav_end_date.strftime('%Y%m%d')
        
        nav_df = pro.fund_nav(
            ts_code=ts_code,
            start_date=start_str,
            end_date=nav_end_str
        )
        if nav_df is not None and len(nav_df) > 0:
            nav_df = nav_df.sort_values('nav_date')
            print(f"   获取到 {len(nav_df)} 条净值数据")
            print(f"   最新日期: {nav_df['nav_date'].max()}")
        else:
            print("   未获取到净值数据")
            nav_df = None
    except Exception as e:
        print(f"   获取净值数据失败: {e}")
        nav_df = None
    
    # 3. 合并数据并计算溢价率
    print("\n3. 合并数据并计算溢价率...")
    
    if price_df is None:
        print("错误: 没有价格数据")
        return None
    
    # 准备价格数据
    price_df['date'] = pd.to_datetime(price_df['trade_date'])
    price_df = price_df.set_index('date')
    price_series = price_df['close']
    
    # 准备净值数据
    if nav_df is not None:
        nav_df['date'] = pd.to_datetime(nav_df['nav_date'])
        nav_df = nav_df.set_index('date')
        nav_series = nav_df['unit_nav']
    else:
        nav_series = pd.Series()
    
    # 创建结果DataFrame
    result = pd.DataFrame({
        '收盘价': price_series
    })
    
    # 对齐净值数据（按日期）
    result = result.join(nav_series.rename('净值'), how='left')
    
    # 计算溢价率
    result['溢价率'] = (result['收盘价'] - result['净值']) / result['净值'] * 100
    
    # 取最新N天
    result = result.tail(days)
    
    # 格式化输出
    result['收盘价'] = result['收盘价'].round(3)
    result['净值'] = result['净值'].round(3)
    result['溢价率'] = result['溢价率'].round(2)
    
    # 重置索引，将日期作为列
    result = result.reset_index()
    result['日期'] = result['date'].dt.strftime('%Y-%m-%d')
    result = result[['日期', '收盘价', '净值', '溢价率']]
    
    return result


def main():
    """主函数"""
    etf_code = "159930.SZ"
    days = 10
    
    print("=" * 60)
    print(f"ETF: {etf_code} (中证能源ETF)")
    print(f"获取最近 {days} 天数据")
    print("=" * 60)
    
    df = fetch_etf_data(etf_code, days)
    
    if df is not None and len(df) > 0:
        print("\n" + "=" * 60)
        print("结果表格:")
        print("=" * 60)
        print(df.to_string(index=False))
        
        # 保存到CSV
        output_file = f"{etf_code.replace('.', '_')}_latest_{days}days.csv"
        df.to_csv(output_file, index=False, encoding='utf-8-sig')
        print(f"\n数据已保存到: {output_file}")
    else:
        print("\n获取数据失败")


if __name__ == "__main__":
    main()