# coding:utf-8

import requests
import pandas as pd
from sqlalchemy import create_engine, text
from datetime import datetime, timedelta
from tqdm import tqdm
import sys
import os
import gc

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

from src.scripts.config import XUEQIU_HEADERS
from src.scripts.ProxyIP import EnhancedProxyManager
from src.quantitative_analysis.batch_stock_price_collector import fetch_and_store_stock_data
from src.valuation_analysis.stock_price_collector import StockPriceCollector
from src.scripts.stock_daily_data_collector import StockDailyDataCollector


class StockDailyDataCollectorV2:
    """股票日线数据采集器V2版本 - 整合雪球和东方财富数据"""

    def __init__(self, db_url):
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        self.headers = XUEQIU_HEADERS
        self.proxy_manager = EnhancedProxyManager()
        # 创建东方财富数据采集器
        self.eastmoney_collector = StockPriceCollector(db_url)
        # 创建原版采集器用于单只股票历史数据获取
        self.original_collector = StockDailyDataCollector(db_url)

    def convert_symbol_format(self, symbol):
        """
        将雪球格式的股票代码转换为东方财富格式
        雪球格式：SZ300177 -> 东方财富格式：300177.SZ
        """
        if symbol.startswith('SZ'):
            return f"{symbol[2:]}.SZ"
        elif symbol.startswith('SH'):
            return f"{symbol[2:]}.SH"
        else:
            return symbol

    def convert_eastmoney_to_xueqiu_format(self, stock_code):
        """
        将东方财富格式的股票代码转换为雪球格式
        东方财富格式：300177.SZ -> 雪球格式：SZ300177
        """
        if '.SZ' in stock_code:
            return f"SZ{stock_code.replace('.SZ', '')}"
        elif '.SH' in stock_code:
            return f"SH{stock_code.replace('.SH', '')}"
        else:
            return stock_code

    def fetch_eastmoney_data(self):
        """获取东方财富的实时数据"""
        print("正在获取东方财富数据...")
        df = self.eastmoney_collector.fetch_all_data()
        if not df.empty:
            # 转换股票代码格式为雪球格式
            df['symbol'] = df['stock_code'].apply(self.convert_eastmoney_to_xueqiu_format)
            print(f"成功获取东方财富数据，共 {len(df)} 条记录")
        return df

    def merge_data(self, xueqiu_df, eastmoney_df):
        """合并雪球和东方财富数据"""
        print("正在合并雪球和东方财富数据...")

        # 基于symbol进行合并
        merged_df = pd.merge(
            xueqiu_df,
            eastmoney_df[['symbol', 'high_price', 'low_price', 'open_price', 'pre_close', 'list_date']],
            on='symbol',
            how='left'
        )
        
        print(f"数据合并完成，共 {len(merged_df)} 条记录")
        return merged_df

    def transform_to_gp_day_data(self, merged_df):
        """将合并后的数据转换为gp_day_data表结构"""
        print("正在转换数据格式...")
        
        # 创建符合gp_day_data表结构的DataFrame
        gp_day_df = pd.DataFrame()
        
        # 映射字段
        gp_day_df['symbol'] = merged_df['symbol']
        # 将timestamp设置为当天的00:00:00格式
        today = datetime.now().date()
        gp_day_df['timestamp'] = pd.to_datetime(today)
        gp_day_df['volume'] = merged_df['volume']
        gp_day_df['open'] = merged_df['open_price']
        gp_day_df['high'] = merged_df['high_price']
        gp_day_df['low'] = merged_df['low_price']
        gp_day_df['close'] = merged_df['current']
        gp_day_df['chg'] = merged_df['chg']
        gp_day_df['percent'] = merged_df['percent']
        gp_day_df['turnoverrate'] = merged_df['turnover_rate']
        gp_day_df['amount'] = merged_df['amount']
        gp_day_df['pb'] = merged_df['pb']
        gp_day_df['pe'] = merged_df['pe_ttm']
        gp_day_df['ps'] = merged_df['ps']
        
        # 添加pre_close字段（用于除权检查）
        gp_day_df['pre_close'] = merged_df['pre_close']
        
        print(f"数据转换完成，共 {len(gp_day_df)} 条记录")
        return gp_day_df

    def save_to_database(self, df):
        """保存数据到数据库"""
        if df.empty:
            print("没有数据需要保存")
            return
        
        print(f"正在保存数据到数据库，共 {len(df)} 条记录...")
        
        # 删除今日数据
        today_str = datetime.now().strftime('%Y-%m-%d')
        delete_query = text("DELETE FROM gp_day_data WHERE `timestamp` LIKE :date_str")
        
        try:
            with self.engine.begin() as conn:
                conn.execute(delete_query, {"date_str": f"{today_str}%"})
                print(f"已删除今日 {today_str} 的旧数据")
        except Exception as e:
            print(f"删除今日数据失败: {e}")
        
        # 分批保存数据
        batch_size = 1000
        for i in range(0, len(df), batch_size):
            batch = df.iloc[i:i+batch_size]
            try:
                batch.to_sql('gp_day_data', self.engine, if_exists='append', index=False)
                print(f"已保存第 {i//batch_size + 1} 批数据")
            except Exception as e:
                print(f"保存第 {i//batch_size + 1} 批数据失败: {e}")
        
        print("数据保存完成")

    def check_ex_rights_before_save(self, df):
        """在保存数据库之前检查除权情况，返回除权股票列表和除权日志数据"""
        print("步骤5.1: 检查除权情况（保存前）...")
        
        ex_rights_stocks = []
        ex_rights_log_data = []
        today_str = datetime.now().strftime('%Y-%m-%d')
        
        for _, row in tqdm(df.iterrows(), total=len(df), desc="检查除权"):
            symbol = row['symbol']
            current_pre_close = row['pre_close']
            
            # 如果pre_close为空，跳过
            if pd.isna(current_pre_close):
                continue
            
            # 查询数据库中该股票的最近两条收盘价记录
            query = text("""
                SELECT `close`, `timestamp` FROM gp_day_data 
                WHERE symbol = :symbol 
                ORDER BY `timestamp` DESC 
                LIMIT 2
            """)
            
            try:
                with self.engine.connect() as conn:
                    results = conn.execute(query, {"symbol": symbol}).fetchall()
                    
                if results:
                    # 检查最新记录是否为今天的数据
                    latest_record = results[0]
                    latest_timestamp = latest_record[1]
                    latest_date_str = latest_timestamp.strftime('%Y-%m-%d')
                    
                    if latest_date_str == today_str and len(results) > 1:
                        # 如果最新记录是今天的，且有第二条记录，则用第二条记录比较
                        db_last_close = float(results[1][0])
                        # print(f"股票 {symbol}: 检测到今日已有数据，使用昨日收盘价 {db_last_close} 进行比较")
                    else:
                        # 如果最新记录不是今天的，或者只有一条记录，则用最新记录比较
                        db_last_close = float(results[0][0])
                    
                    # 比较pre_close和数据库中的收盘价
                    if abs(db_last_close - current_pre_close) > 0.001:
                        print(f"发现除权股票: {symbol}, 数据库收盘价: {db_last_close}, 当前昨收价: {current_pre_close}")
                        ex_rights_stocks.append(symbol)
                        
                        # 收集除权日志数据
                        ex_rights_log_data.append({
                            'symbol': symbol,
                            'date': today_str,
                            'db_price': db_last_close,
                            'api_price': current_pre_close,
                            'log_time': datetime.now()
                        })
            except Exception as e:
                print(f"查询 {symbol} 历史数据失败: {e}")
                continue
        
        if ex_rights_stocks:
            print(f"检测到 {len(ex_rights_stocks)} 只除权股票: {ex_rights_stocks}")
        else:
            print("未发现除权股票")
            
        return ex_rights_stocks, ex_rights_log_data

    def save_ex_rights_log(self, log_data: list):
        """将除权日志保存到数据库"""
        if not log_data:
            return
        
        print(f"正在保存 {len(log_data)} 条除权日志到gp_ex_rights_log表...")
        try:
            df = pd.DataFrame(log_data)
            # 确保列名与数据库字段匹配
            df = df.rename(columns={
                'symbol': 'stock_code',
                'date': 'change_date',
                'db_price': 'before_price',
                'api_price': 'after_price',
                'log_time': 'update_time'
            })
            df.to_sql('gp_ex_rights_log', self.engine, if_exists='append', index=False)
            print("除权日志保存成功")
        except Exception as e:
            print(f"保存除权日志失败: {e}")

    def handle_ex_rights_stocks(self, ex_rights_stocks, ex_rights_log_data):
        """处理除权股票：保存日志、删除历史数据并重新获取1800天数据"""
        if not ex_rights_stocks:
            return
            
        print("步骤6: 处理除权股票...")
        
        # 6.1 保存除权日志
        if ex_rights_log_data:
            self.save_ex_rights_log(ex_rights_log_data)
        
        # 6.2 重新获取历史数据
        print(f"开始处理 {len(ex_rights_stocks)} 只除权股票，重新获取历史数据...")
        
        for symbol in tqdm(ex_rights_stocks, desc="处理除权股票"):
            try:
                # 删除该股票的所有历史数据
                delete_query = text("DELETE FROM gp_day_data WHERE symbol = :symbol")
                with self.engine.begin() as conn:
                    conn.execute(delete_query, {"symbol": symbol})
                print(f"已删除 {symbol} 的历史数据")
                
                # 重新获取1800天的历史数据
                success = self.original_collector.fetch_single_stock_history(symbol, 1800)
                if success:
                    print(f"成功重新获取 {symbol} 的历史数据")
                else:
                    print(f"重新获取 {symbol} 的历史数据失败")
                    
            except Exception as e:
                print(f"处理除权股票 {symbol} 失败: {e}")

    def run_daily_collection(self):
        """执行每日数据采集"""
        print("=" * 60)
        print("股票日线数据采集器V2 - 开始运行")
        print("=" * 60)
        
        try:
            # 1. 获取雪球数据
            print("步骤1: 获取雪球数据...")
            xueqiu_df = fetch_and_store_stock_data()
            if xueqiu_df.empty:
                print("雪球数据获取失败，终止运行")
                return

            # 2. 获取东方财富数据
            print("步骤2: 获取东方财富数据...")
            eastmoney_df = self.fetch_eastmoney_data()
            if eastmoney_df.empty:
                print("东方财富数据获取失败，终止运行")
                return

            # 3. 合并数据
            print("步骤3: 合并数据...")
            merged_df = self.merge_data(xueqiu_df, eastmoney_df)

            # 4. 转换数据格式
            print("步骤4: 转换数据格式...")
            gp_day_df = self.transform_to_gp_day_data(merged_df)

            # 5. 检查除权（保存前）
            ex_rights_stocks, ex_rights_log_data = self.check_ex_rights_before_save(gp_day_df)

            # 5.2. 保存到数据库
            print("步骤5.2: 保存到数据库...")
            self.save_to_database(gp_day_df)

            # 6. 处理除权股票（保存后）
            self.handle_ex_rights_stocks(ex_rights_stocks, ex_rights_log_data)
            
            print("=" * 60)
            print("股票日线数据采集完成")
            print("=" * 60)
            
        except Exception as e:
            print(f"采集过程中发生错误: {e}")
        finally:
            # 清理资源
            self.engine.dispose()
            gc.collect()


def collect_stock_daily_data_v2(db_url):
    """V2版本的股票日线数据采集入口函数"""
    collector = StockDailyDataCollectorV2(db_url)
    collector.run_daily_collection()


if __name__ == "__main__":
    db_url = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    collect_stock_daily_data_v2(db_url) 