#!/usr/bin/env python3
"""
批量更新数据库中缺少财务指标的基本面数据
"""
import akshare as ak
import pandas as pd
import psycopg2
from datetime import date, datetime
import time

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'port': 5432,
    'dbname': 'stock_app',
    'user': 'postgres',
    'password': 'xypg5432'
}

def get_db():
    """获取数据库连接"""
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        return conn
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return None

def get_stocks_need_update():
    """获取需要更新财务指标的股票列表"""
    conn = get_db()
    if not conn:
        return []
    
    try:
        cur = conn.cursor()
        cur.execute("""
            SELECT code, name 
            FROM stock_fundamental 
            WHERE roe IS NULL OR eps IS NULL
        """)
        return cur.fetchall()
    finally:
        conn.close()

def fetch_financial_indicators(stock_code):
    """从akshare获取股票财务指标"""
    try:
        df = ak.stock_financial_analysis_indicator(symbol=stock_code, start_year='2024')
        if df is None or df.empty:
            return None
        
        latest = df.iloc[-1]
        
        def safe_float(val):
            if pd.isna(val):
                return None
            try:
                return float(val)
            except:
                return None
        
        return {
            'eps': safe_float(latest.get('摊薄每股收益(元)')),
            'bps': safe_float(latest.get('每股净资产_调整后(元)')),
            'roe': safe_float(latest.get('净资产收益率(%)')),
            'gross_margin': safe_float(latest.get('销售毛利率(%)')),
            'net_margin': safe_float(latest.get('销售净利率(%)')),
            'revenue_yoy': safe_float(latest.get('主营业务收入增长率(%)')),
            'profit_yoy': safe_float(latest.get('净利润增长率(%)')),
        }
    except Exception as e:
        print(f"  获取 {stock_code} 财务指标失败: {e}")
        return None

def update_stock_fundamental(stock_code, data):
    """更新股票基本面数据"""
    conn = get_db()
    if not conn:
        return False
    
    try:
        cur = conn.cursor()
        cur.execute("""
            UPDATE stock_fundamental 
            SET roe = %s, eps = %s, bps = %s, 
                revenue_yoy = %s, profit_yoy = %s, 
                gross_margin = %s, net_margin = %s,
                updated_at = NOW()
            WHERE code = %s
        """, (
            data.get('roe'),
            data.get('eps'),
            data.get('bps'),
            data.get('revenue_yoy'),
            data.get('profit_yoy'),
            data.get('gross_margin'),
            data.get('net_margin'),
            stock_code
        ))
        conn.commit()
        return True
    except Exception as e:
        conn.rollback()
        print(f"  更新 {stock_code} 失败: {e}")
        return False
    finally:
        conn.close()

def main():
    print("=" * 50)
    print("批量更新基本面财务指标")
    print("=" * 50)
    
    # 获取需要更新的股票
    stocks = get_stocks_need_update()
    print(f"\n需要更新的股票数量: {len(stocks)}")
    
    if not stocks:
        print("所有股票财务指标已完整，无需更新")
        return
    
    success_count = 0
    fail_count = 0
    
    for i, (code, name) in enumerate(stocks, 1):
        print(f"\n[{i}/{len(stocks)}] 正在更新 {code} {name}...")
        
        # 获取财务指标
        data = fetch_financial_indicators(code)
        
        if data:
            # 更新数据库
            if update_stock_fundamental(code, data):
                print(f"  ✓ 更新成功: EPS={data.get('eps')}, ROE={data.get('roe')}%")
                success_count += 1
            else:
                fail_count += 1
        else:
            print(f"  ✗ 无法获取财务指标")
            fail_count += 1
        
        # 避免请求过快
        time.sleep(0.5)
    
    print("\n" + "=" * 50)
    print(f"更新完成! 成功: {success_count}, 失败: {fail_count}")
    print("=" * 50)

if __name__ == '__main__':
    main()