backend/scripts/sync_data.py

"""
数据同步脚本 - 云端兼容版
支持本地母本模式 (DuckDB) 和 云端无盘模式 (Parquet View)
"""

import os
import sys
import logging
import time
from datetime import datetime, timedelta
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Optional, Dict, Any
from pathlib import Path

import pandas as pd
import akshare as ak
import yfinance as yf
from huggingface_hub import hf_hub_download, upload_file

# 添加父目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.database import get_db
from app.database_user import get_beijing_time

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 配置
YEARS_OF_DATA = 10
MAX_WORKERS = 5  # 降低并发数，减少超时
SYNC_LIMIT = -1
USE_YAHOO_FALLBACK = 0

def get_stock_list() -> pd.DataFrame:
    """获取全市场标的列表"""
    logger.info("Fetching all-market target list...")
    all_lists = []
    try:
        df_a = ak.stock_zh_a_spot_em()[['代码', '名称']]
        df_a.columns = ['code', 'name']
        df_a['market'] = df_a['code'].apply(lambda x: '主板' if x.startswith(('60', '00')) else ('创业板' if x.startswith('30') else ('科创板' if x.startswith('68') else ('北交所' if x.startswith(('8', '4', '920')) else '其他'))))
        all_lists.append(df_a)
        
        # ETF / LOF / REITs
        try:
            df_etf = ak.fund_etf_spot_em()[['代码', '名称']]
            df_etf.columns = ['code', 'name']
            df_etf['market'] = 'ETF'
            all_lists.append(df_etf)
        except Exception as e:
            logger.warning(f"ETF list fetch failed: {e}")

        try:
            df_lof = ak.fund_lof_spot_em()[['代码', '名称']]
            df_lof.columns = ['code', 'name']
            df_lof['market'] = 'LOF'
            all_lists.append(df_lof)
        except Exception as e:
            logger.warning(f"LOF list fetch failed: {e}")

        try:
            df_reits = ak.reits_realtime_em()[['代码', '名称']]
            df_reits.columns = ['code', 'name']
            df_reits['market'] = 'REITs'
            all_lists.append(df_reits)
        except Exception as e:
            logger.warning(f"REITs list fetch failed: {e}")
        
        try:
            df_cb = ak.bond_zh_hs_cov_spot()
            c_code = '代码' if '代码' in df_cb.columns else 'symbol'
            c_name = '名称' if '名称' in df_cb.columns else 'name'
            df_cb = df_cb[[c_code, c_name]]; df_cb.columns = ['code', 'name']; df_cb['market'] = '可转债'; all_lists.append(df_cb)
        except: pass

    except Exception as e:
        logger.error(f"List fetching error: {e}")

    if not all_lists:
        db = get_db()
        return db.conn.execute("SELECT code, name, market FROM stock_list").df()
        
    df = pd.concat(all_lists).drop_duplicates(subset=['code'])
    df['list_date'] = None
    return df

def get_target_daily(code: str, start_date: str, market: str) -> Optional[pd.DataFrame]:
    """抓取单只标的数据"""
    max_retries = 3  # 增加重试次数
    for attempt in range(max_retries):
        try:
            end_date = get_beijing_time().strftime('%Y%m%d')
            fetch_start = start_date.replace('-', '')
            df = None
            if market == 'INDEX':
                df = ak.stock_zh_index_daily_em(symbol=f"sh{code}" if code.startswith('000') else f"sz{code}")
            elif market == 'ETF':
                df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=fetch_start, end_date=end_date, adjust="hfq")
            elif market == 'LOF':
                df = ak.fund_lof_hist_em(symbol=code, period="daily", start_date=fetch_start, end_date=end_date, adjust="hfq")
            elif market == '可转债':
                # 可转债接口通常使用 6 位数字代码，兼容 sh110xxx / sz12xxxx / bj81xxxx
                cov_symbol = code[-6:] if len(code) > 6 else code
                try:
                    df = ak.bond_zh_hs_cov_daily(symbol=cov_symbol)
                except Exception:
                    df = ak.bond_zh_hs_cov_daily(symbol=code)
            elif market == 'REITs':
                # REITs 接口仅支持 symbol 参数，返回全量历史后再按日期过滤
                df = ak.reits_hist_em(symbol=code)
            else:
                df = ak.stock_zh_a_hist(symbol=code, period="daily", start_date=fetch_start, end_date=end_date, adjust="hfq")
            
            if df is not None and not df.empty:
                # 标准化列名
                rename_map = {
                    '日期': 'trade_date', 'date': 'trade_date', 'Date': 'trade_date',
                    '开盘': 'open', '今开': 'open', 'Open': 'open',
                    '最高': 'high', 'High': 'high',
                    '最低': 'low', 'Low': 'low',
                    '收盘': 'close', '最新价': 'close', 'Close': 'close',
                    '成交量': 'volume', 'Volume': 'volume',
                    '成交额': 'amount', 'Amount': 'amount',
                    '涨跌幅': 'pct_chg',
                    '换手率': 'turnover_rate', '换手': 'turnover_rate'
                }
                df = df.rename(columns=rename_map)
                
                # 如果还是没有 trade_date，尝试将索引转为列
                if 'trade_date' not in df.columns:
                    df = df.reset_index().rename(columns={'index': 'trade_date', 'date': 'trade_date'})
                
                df['trade_date'] = pd.to_datetime(df['trade_date'])
                df = df[df['trade_date'] >= pd.to_datetime(start_date)]
                if 'amount' not in df.columns: df['amount'] = 0
                if 'pct_chg' not in df.columns: df['pct_chg'] = df['close'].pct_change() * 100
                if 'turnover_rate' not in df.columns: df['turnover_rate'] = 0
                df['code'] = code
                return df[['code', 'trade_date', 'open', 'high', 'low', 'close', 'volume', 'amount', 'pct_chg', 'turnover_rate']]
        except Exception as e:
            if attempt == max_retries - 1:
                logger.warning(f"Failed to fetch {code} ({market}): {str(e)}")
            time.sleep(1)
    return None

def get_last_trading_day() -> str:
    """获取最近一个交易日（用于增量截止日）"""
    try:
        df = ak.stock_zh_index_daily_em(symbol="sh000300")
        if df is not None and not df.empty:
            date_col = 'date' if 'date' in df.columns else ('日期' if '日期' in df.columns else None)
            if date_col:
                return pd.to_datetime(df[date_col].iloc[-1]).strftime('%Y-%m-%d')
    except Exception as e:
        logger.warning(f"Failed to get last trading day from index data: {e}")

    # 回退：按工作日估算
    d = get_beijing_time()
    while d.weekday() >= 5:  # 5=周六, 6=周日
        d -= timedelta(days=1)
    return d.strftime('%Y-%m-%d')


def get_index_daily(code: str) -> Optional[pd.DataFrame]:
    """抓取指数日线（默认用于沪深300）"""
    try:
        symbol = f"sh{code}" if code.startswith('000') else f"sz{code}"
        df = ak.stock_zh_index_daily_em(symbol=symbol)
        if df is None or df.empty:
            return None

        rename_map = {
            'date': 'trade_date', '日期': 'trade_date',
            'open': 'open', '开盘': 'open',
            'high': 'high', '最高': 'high',
            'low': 'low', '最低': 'low',
            'close': 'close', '收盘': 'close',
            'volume': 'volume', '成交量': 'volume',
            'amount': 'amount', '成交额': 'amount',
            'pct_chg': 'pct_chg', '涨跌幅': 'pct_chg'
        }
        df = df.rename(columns=rename_map)
        if 'trade_date' not in df.columns:
            return None

        df['trade_date'] = pd.to_datetime(df['trade_date'])
        if 'amount' not in df.columns:
            df['amount'] = 0
        if 'pct_chg' not in df.columns:
            df['pct_chg'] = df['close'].pct_change() * 100
        if 'volume' not in df.columns:
            df['volume'] = 0
        df['turnover_rate'] = 0
        df['code'] = code

        return df[['code', 'trade_date', 'open', 'high', 'low', 'close', 'volume', 'amount', 'pct_chg', 'turnover_rate']]
    except Exception as e:
        logger.warning(f"Failed to fetch index {code}: {e}")
        return None


def sync_stock_daily(targets: List[Dict[str, str]], last_trade_day: str) -> int:
    """增量同步逻辑 - 兼容云端视图"""
    db = get_db()
    
    # 1. 获取现状 (无论是表还是视图)
    existing_latest = db.conn.execute("SELECT code, CAST(MAX(trade_date) AS VARCHAR) FROM stock_daily GROUP BY code").fetchall()
    latest_map = {row[0]: row[1] for row in existing_latest}
    
    pending = []
    for t in targets:
        code = t['code']
        if code in latest_map:
            if latest_map[code] >= last_trade_day: continue
            start_dt = (pd.to_datetime(latest_map[code]) + timedelta(days=1)).strftime('%Y-%m-%d')
        else:
            start_dt = (get_beijing_time() - timedelta(days=YEARS_OF_DATA * 365)).strftime('%Y-%m-%d')
        t['start_dt'] = start_dt
        pending.append(t)

    if not pending: return 0
    logger.info(f"Syncing {len(pending)} targets...")

    all_new_data = []
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {executor.submit(get_target_daily, t['code'], t['start_dt'], t['market']): t['code'] for t in pending}
        for i, future in enumerate(as_completed(futures), 1):
            res = future.result()
            if res is not None: all_new_data.append(res)
            if i % 500 == 0: logger.info(f"Progress: {i}/{len(pending)}")

    if all_new_data:
        inc_df = pd.concat(all_new_data, ignore_index=True)
        # 识别变动月份
        changed = inc_df.assign(yr=inc_df['trade_date'].dt.year, mo=inc_df['trade_date'].dt.month)[['yr', 'mo']].drop_duplicates().values
        
        for yr, mo in changed:
            yr, mo = int(yr), int(mo)
            filename = f"{yr}-{mo:02d}.parquet"
            local_path = Path(f"backend/data/parquet/{filename}")
            local_path.parent.mkdir(parents=True, exist_ok=True)
            
            # 云端增量核心：如果本地没母本，先从云端拉取旧的月份文件
            if not Path(os.getenv("DUCKDB_PATH", "")).exists() or os.path.getsize(os.getenv("DUCKDB_PATH")) < 1024:
                try:
                    repo_id = os.getenv("DATASET_REPO_ID")
                    if repo_id:
                        old_file = hf_hub_download(repo_id=repo_id, filename=f"data/parquet/{filename}", repo_type="dataset")
                        old_df = pd.read_parquet(old_file)
                        # 合并
                        month_inc = inc_df[(inc_df['trade_date'].dt.year == yr) & (inc_df['trade_date'].dt.month == mo)]
                        final_month_df = pd.concat([old_df, month_inc]).drop_duplicates(subset=['code', 'trade_date'])
                        final_month_df.to_parquet(local_path)
                        logger.info(f"Merged cloud data for {filename}")
                        continue
                except: pass

            # 本地母本模式：直接从母本导出
            db.conn.execute("CREATE TEMP TABLE temp_inc AS SELECT * FROM inc_df")
            db.conn.execute("INSERT OR IGNORE INTO stock_daily SELECT * FROM temp_inc")
            db.conn.execute(f"COPY (SELECT * FROM stock_daily WHERE date_part('year', trade_date) = {yr} AND date_part('month', trade_date) = {mo}) TO '{local_path}' (FORMAT PARQUET)")
            db.conn.execute("DROP TABLE temp_inc")
            
    return len(all_new_data)

def main():
    logger.info("Sync started...")
    db = get_db(); db.init_db()
    
    # 1. 列表同步
    target_list = get_stock_list()
    # 如果是本地母本，更新表；如果是云端，导出 parquet 准备上传
    if Path(os.getenv("DUCKDB_PATH", "")).exists():
        db.conn.execute("DELETE FROM stock_list")
        db.conn.execute("INSERT INTO stock_list SELECT code, name, market, list_date FROM target_list")
    
    list_parquet = Path("backend/data/stock_list.parquet")
    list_parquet.parent.mkdir(parents=True, exist_ok=True)
    target_list.to_parquet(list_parquet)

    # 2. 行情同步
    last_day = get_last_trading_day()
    sync_stock_daily(target_list.to_dict('records'), last_day)
    
    # 3. 指数同步
    idx_df = get_index_daily('000300')
    if idx_df is not None:
        idx_path = Path("backend/data/parquet/index_000300.parquet")
        idx_df.to_parquet(idx_path)

    # 4. 上传
    db.upload_db()
    
    # 5. 热重载：刷新数据库视图和缓存
    try:
        logger.info("Reloading database views to reflect new data...")
        db.init_db() # 重新下载并挂载最新的 Parquet 文件
        
        from app.core import clear_eligible_stocks_cache
        clear_eligible_stocks_cache() # 清除内存缓存
        logger.info("Hot reload completed!")
    except Exception as e:
        logger.error(f"Hot reload failed: {e}")
        
    logger.info("Sync finished!")

if __name__ == "__main__":
    main()