# /backend/app/data_collection/stock_grouper_etl.py

"""
Created on 2025/11/12
---------
@summary: 股票分组脚本 (终极修复版 V5)
          1. 解决 OOM 问题 (按年分批)
          2. 解决 跨批次重复问题 (前序状态合并)
          3. 解决 KeyError: symbol 问题 (手动生成 symbol)
---------
@author: NBR
"""

import os
import time
import gc
from datetime import datetime, timedelta

import pandas as pd
import tushare as ts
from sqlalchemy import create_engine, text

# --- 配置 ---
PROXY_DB_URL = os.environ.get('PROXY_DB_URL')
TUSHARE_TOKEN = os.environ.get('TUSHARE_TOKEN')


class _StockGroupManager:
    """
    股票分组管理器
    """
    REQUEST_INTERVAL = 0.05

    INDEX_MAPPING = {
        '沪深300': '000300.SH',
        '中证500': '000905.SH',
        '中证1000': '000852.SH',
        '创业板': '399006.SZ',
        '科创板': '000688.SH',
        '上证指数': '000001.SH',
        '深证成指': '399001.SZ'
    }

    def __init__(self, token: str, db_engine, test_mode: bool = False):
        self.test_mode = test_mode
        self.pro = ts.pro_api(token)
        if self.pro is None:
            raise Exception("Tushare pro_api initialization failed.")

        self.engine = db_engine
        self.index_constituents_cache = {}
        self.TEST_STOCKS = ['300750.SZ', '688981.SH']
        self._init_database()

    def _init_database(self):
        create_table_sql = """
            CREATE TABLE IF NOT EXISTS gp_stock_category (
                id INT AUTO_INCREMENT PRIMARY KEY,
                ts_code VARCHAR(20) NOT NULL,
                symbol VARCHAR(20) NOT NULL,
                name VARCHAR(100) NOT NULL,
                trade_date DATE NOT NULL,
                industry VARCHAR(50),
                is_st BOOLEAN DEFAULT FALSE,
                index_series TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                UNIQUE KEY unique_stock_date (ts_code, trade_date)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """
        with self.engine.connect() as connection:
            with connection.begin() as trans:
                connection.execute(text(create_table_sql))
                trans.commit()

    # --- 辅助方法 ---
    def _get_month_start_end_dates(self, year_month: str):
        try:
            year = int(year_month[:4])
            month = int(year_month[4:6])
            month_start = datetime(year, month, 1)
            if month == 12:
                next_month = datetime(year + 1, 1, 1)
            else:
                next_month = datetime(year, month + 1, 1)
            month_end = next_month - timedelta(days=1)
            return month_start.strftime('%Y%m%d'), month_end.strftime('%Y%m%d')
        except:
            return year_month + '01', year_month + '31'

    def _get_previous_month_key(self, month_key: str) -> str:
        try:
            year = int(month_key[:4])
            month = int(month_key[4:])
            if month == 1:
                return f"{year - 1:04d}12"
            else:
                return f"{year:04d}{month - 1:02d}"
        except:
            return month_key

    def _cache_month_index_constituents(self, month_key: str):
        if month_key in self.index_constituents_cache:
            return

        self.index_constituents_cache[month_key] = {}
        prev_month_key = self._get_previous_month_key(month_key)
        prev_data = self.index_constituents_cache.get(prev_month_key, {})
        month_start, month_end = self._get_month_start_end_dates(month_key)

        for index_name, index_code in self.INDEX_MAPPING.items():
            current_stocks = set()
            try:
                df = self.pro.index_weight(index_code=index_code, start_date=month_start, end_date=month_end)
                if not df.empty:
                    current_stocks = set(df['con_code'].tolist())
                else:
                    # 前向填充
                    if index_name in prev_data:
                        current_stocks = prev_data[index_name]

                self.index_constituents_cache[month_key][index_name] = current_stocks
                time.sleep(self.REQUEST_INTERVAL)
            except Exception as e:
                print(f"      缓存{index_name}成分股失败: {e}")
                if index_name in prev_data:
                    self.index_constituents_cache[month_key][index_name] = prev_data[index_name]

    def _get_index_series_for_stock(self, ts_code: str, trade_date: str) -> str:
        try:
            index_series = []
            month_key = trade_date[:6]
            if month_key not in self.index_constituents_cache:
                self._cache_month_index_constituents(month_key)

            month_data = self.index_constituents_cache.get(month_key, {})
            for index_name, stock_set in month_data.items():
                if ts_code in stock_set:
                    index_series.append(index_name)
            return ','.join(index_series) if index_series else ''
        except:
            return ''

    def _is_st_stock(self, stock_name: str) -> bool:
        if pd.isna(stock_name): return False
        return 'ST' in stock_name or '*ST' in stock_name

    def _format_code_to_prefix_style(self, ts_code: str) -> str:
        if '.' in ts_code:
            parts = ts_code.split('.')
            if len(parts) == 2:
                return f"{parts[1].upper()}{parts[0]}"
        return ts_code

    def get_latest_states_before(self, before_date_str: str) -> pd.DataFrame:
        """获取前序基准状态"""
        print(f"  🔍 查询 {before_date_str} 之前的基准状态 (用于跨年去重)...")
        query = text("""
            SELECT t1.ts_code, t1.industry, t1.is_st, t1.index_series, t1.trade_date, t1.name
            FROM gp_stock_category t1
            JOIN (
                SELECT ts_code, MAX(trade_date) as max_dt
                FROM gp_stock_category
                WHERE trade_date < :before_date
                GROUP BY ts_code
            ) t2 ON t1.ts_code = t2.ts_code AND t1.trade_date = t2.max_dt
        """)

        try:
            df = pd.read_sql(query, self.engine, params={'before_date': before_date_str})
            if not df.empty:
                df['trade_date'] = pd.to_datetime(df['trade_date'])

                def _revert_code(c):
                    if len(c) > 2 and c[:2] in ['SZ', 'SH']:
                        return f"{c[2:]}.{c[:2]}"
                    return c

                df['ts_code'] = df['ts_code'].apply(_revert_code)

            print(f"  ✅ 获取到 {len(df)} 条基准记录。")
            return df
        except Exception as e:
            print(f"  ⚠️ 获取基准状态失败 (可能是第一次运行): {e}")
            return pd.DataFrame()

    def get_all_bak_basic_in_range(self, start_date: str, end_date: str) -> pd.DataFrame:
        print(f"  正在获取 {start_date} 到 {end_date} 的日线数据...")
        all_dates = pd.date_range(start=start_date, end=end_date, freq='B')
        all_dfs = []

        for i, date in enumerate(all_dates):
            date_str = date.strftime('%Y%m%d')
            try:
                df = self.pro.bak_basic(trade_date=date_str)
                if df is not None and not df.empty:
                    all_dfs.append(df)
                time.sleep(self.REQUEST_INTERVAL)
            except Exception as e:
                print(f"  获取 {date_str} 数据失败: {e}")

        if not all_dfs:
            return pd.DataFrame()
        return pd.concat(all_dfs, ignore_index=True)

    def update_stock_category_data(self, start_date: str, end_date: str):
        """核心处理逻辑"""
        # 1. 基础日线
        daily_data_df = self.get_all_bak_basic_in_range(start_date, end_date)
        if daily_data_df.empty:
            print("  无数据，本批次结束。")
            return

        # =================【关键修复】=================
        # 确保 symbol 列存在，因为后续 merge 需要用到它
        if 'symbol' not in daily_data_df.columns:
            # 从 ts_code 生成 symbol (例如 000001.SZ -> 000001)
            daily_data_df['symbol'] = daily_data_df['ts_code'].apply(lambda x: str(x).split('.')[0])
        # ==============================================

        print(f"  批量计算属性 ({len(daily_data_df)} 行)...")
        daily_data_df['trade_date'] = pd.to_datetime(daily_data_df['trade_date'], format='%Y%m%d')
        daily_data_df['is_st'] = daily_data_df['name'].apply(self._is_st_stock)

        # 2. 预加载指数缓存
        unique_months = sorted(daily_data_df['trade_date'].dt.strftime('%Y%m').unique())
        for ym in unique_months:
            self._cache_month_index_constituents(ym)

        # 3. 计算指数归属
        daily_data_df['index_series'] = daily_data_df.apply(
            lambda row: self._get_index_series_for_stock(row['ts_code'], row['trade_date'].strftime('%Y%m%d')),
            axis=1
        )

        # 4. 合并前序状态
        prev_state_df = self.get_latest_states_before(pd.to_datetime(start_date).strftime('%Y-%m-%d'))

        if not prev_state_df.empty:
            cols = ['ts_code', 'trade_date', 'industry', 'is_st', 'index_series']
            combined_df = pd.concat([prev_state_df[cols], daily_data_df[cols]], ignore_index=True)
        else:
            combined_df = daily_data_df.copy()

        # 5. 变化检测
        combined_df.sort_values(['ts_code', 'trade_date'], inplace=True)
        cols_to_check = ['industry', 'is_st', 'index_series']

        shifted = combined_df.groupby('ts_code')[cols_to_check].shift(1)
        is_changed = (combined_df[cols_to_check].ne(shifted)).any(axis=1)

        # 6. 筛选结果
        current_batch_start_dt = pd.to_datetime(start_date)
        records_to_insert = combined_df[is_changed & (combined_df['trade_date'] >= current_batch_start_dt)].copy()

        # 7. 补回 symbol 和 name
        if not records_to_insert.empty:
            records_to_insert = pd.merge(
                records_to_insert[['ts_code', 'trade_date', 'industry', 'is_st', 'index_series']],
                daily_data_df[['ts_code', 'trade_date', 'symbol', 'name']],  # 这里现在肯定有 symbol 了
                on=['ts_code', 'trade_date'],
                how='left'
            )

        print(f"  筛选出 {len(records_to_insert)} 条真实变更记录。")

        # 8. 写入
        if not records_to_insert.empty:
            records_to_insert['ts_code'] = records_to_insert['ts_code'].apply(self._format_code_to_prefix_style)
            # 双重保险：如果 merge 后 symbol 仍有空值（极少情况），再次生成
            if 'symbol' not in records_to_insert.columns or records_to_insert['symbol'].isnull().any():
                records_to_insert['symbol'] = records_to_insert['ts_code'].apply(lambda x: x[2:] if len(x) > 2 else x)

            final_df = records_to_insert[
                ['ts_code', 'symbol', 'name', 'trade_date', 'industry', 'is_st', 'index_series']]

            with self.engine.connect() as connection:
                with connection.begin() as trans:
                    del_start = pd.to_datetime(start_date).strftime('%Y-%m-%d')
                    del_end = pd.to_datetime(end_date).strftime('%Y-%m-%d')
                    connection.execute(
                        text("DELETE FROM gp_stock_category WHERE trade_date BETWEEN :s AND :e"),
                        {'s': del_start, 'e': del_end}
                    )
                    final_df.to_sql('gp_stock_category', con=connection, if_exists='append', index=False)
                    trans.commit()
            print(f"  ✅ 写入成功。")

        # 9. 清理
        del daily_data_df
        del combined_df
        del records_to_insert
        if not prev_state_df.empty: del prev_state_df
        gc.collect()


# --- 运行入口 ---
def run_stock_grouper_etl(start_date: str = None, end_date: str = None, test_mode: bool = False,
                          is_full_run: bool = False):
    if not PROXY_DB_URL or not TUSHARE_TOKEN:
        raise ValueError("Environment variables required.")

    engine = create_engine(PROXY_DB_URL, pool_pre_ping=True)
    manager = _StockGroupManager(TUSHARE_TOKEN, engine, test_mode)

    if is_full_run:
        print("🚀 [Full Run] 启动分批处理模式 (按年)...")
        current_year = 2016
        end_year = datetime.now().year

        while current_year <= end_year:
            batch_start = f"{current_year}0101"
            batch_end = f"{current_year}1231"

            if current_year == end_year:
                batch_end = datetime.now().strftime('%Y%m%d')
                if batch_start > batch_end: break

            print(f"\n>>> 处理批次: {batch_start} - {batch_end}")
            manager.update_stock_category_data(batch_start, batch_end)
            current_year += 1
            time.sleep(1)

        print("🎉 全量重跑完成。")

    else:
        if not start_date:
            with engine.connect() as conn:
                res = conn.execute(text("SELECT MAX(trade_date) FROM gp_stock_category")).scalar()
            if res:
                start_date = (res + timedelta(days=1)).strftime('%Y%m%d')
            else:
                start_date = '20160101'
            end_date = datetime.now().strftime('%Y%m%d')

        if start_date > datetime.now().strftime('%Y%m%d'):
            return {"message": "No new data."}

        print(f"🚀 [Incremental] 增量更新: {start_date} - {end_date}")
        manager.update_stock_category_data(start_date, end_date)

    return {"message": "Success"}
