# coding:utf-8

"""
港股日线数据采集工具
功能：从Tushare获取港股日线行情数据，存储到gp_day_data表（与A股共用）
API文档: https://tushare.pro/document/2?doc_id=192
注意：
  1. 该接口需要120积分权限
  2. 单次最大返回8000条数据
  3. 每日收盘后更新（港股交易时间：9:30-12:00, 13:00-16:00）
  4. 可按日期一次性获取全市场数据
数据存储说明：
  - 存储表：gp_day_data（与A股共用）
  - symbol格式：纯数字5位（如01810），不带市场编码
  - 删除策略：只删除LENGTH(symbol)=5的记录，不影响A股数据
  - PB/PE/PS：港股接口无此数据，保持为NULL
采集策略：
  - 每日增量：按日期获取全市场数据（推荐）
  - 历史数据：按股票逐个采集历史数据
  17点以后执行
"""
from time import sleep

import tushare as ts
import pandas as pd
from sqlalchemy import create_engine, text
from datetime import datetime, timedelta
import time
import sys
import os

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

from src.scripts.config import TUSHARE_TOKEN


class HKDailyDataCollector:
    """港股日线数据采集器"""

    def __init__(self, db_url, tushare_token, table_name='gp_day_data'):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            table_name: 数据表名，默认为gp_day_data（与A股共用）
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        self.table_name = table_name
        
        # 初始化Tushare
        ts.set_token(tushare_token)
        self.pro = ts.pro_api()
        
        print("=" * 60)
        print("港股日线数据采集工具")
        print(f"目标数据表: {self.table_name}")
        print("=" * 60)

    def convert_tushare_code_to_db(self, ts_code):
        """
        将Tushare格式的港股代码转换为数据库格式
        Tushare格式：01810.HK、00700.HK
        数据库格式：01810、00700（不带市场编码，保留前导零，长度为5位）
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            数据库格式的股票代码（纯数字，5位）
        """
        if '.HK' in ts_code:
            return ts_code.replace('.HK', '')
        else:
            return ts_code

    def convert_db_code_to_tushare(self, symbol):
        """
        将数据库格式的港股代码转换为Tushare格式
        数据库格式：01810、00700
        Tushare格式：01810.HK、00700.HK
        
        Args:
            symbol: 数据库格式的股票代码
            
        Returns:
            Tushare格式的股票代码
        """
        if not symbol.endswith('.HK'):
            return f"{symbol}.HK"
        return symbol

    def fetch_hk_stock_list(self):
        """
        从数据库获取港股代码列表
        
        Returns:
            港股代码列表（数据库格式）
        """
        try:
            query = """
                SELECT DISTINCT gp_code 
                FROM gp_code_hk 
                WHERE gp_code IS NOT NULL 
                ORDER BY gp_code
            """
            df = pd.read_sql(query, self.engine)
            stock_codes = df['gp_code'].tolist()
            print(f"从数据库获取到{len(stock_codes)}只港股代码")
            return stock_codes
        except Exception as e:
            print(f"从数据库获取港股代码失败: {e}")
            return []

    def fetch_hk_daily_data(self, trade_date=None, ts_code=None, start_date=None, end_date=None):
        """
        获取港股日线数据
        
        Args:
            trade_date: 交易日期，格式YYYYMMDD（单日查询）
            ts_code: 股票代码，Tushare格式（单股查询）
            start_date: 开始日期，格式YYYYMMDD（范围查询）
            end_date: 结束日期，格式YYYYMMDD（范围查询）
            
        Returns:
            港股日线数据DataFrame
        """
        try:
            # 构建API参数
            params = {}
            if trade_date:
                params['trade_date'] = trade_date
            if ts_code:
                params['ts_code'] = ts_code
            if start_date:
                params['start_date'] = start_date
            if end_date:
                params['end_date'] = end_date
            
            df = self.pro.hk_daily(**params)
            time.sleep(0.21)  # API限频控制，每分钟不超过280次
            return df
        except Exception as e:
            print(f"获取港股日线数据失败: {e}")
            return pd.DataFrame()

    def transform_data(self, df):
        """
        转换数据格式为数据库格式（与gp_day_data表结构一致）
        
        Args:
            df: 原始数据DataFrame
            
        Returns:
            转换后的DataFrame，存储到gp_day_data表
            
        字段说明：
            - symbol: 纯数字5位（如01810），不带市场编码
            - timestamp: 交易日期
            - volume: 成交量（手）
            - open/high/low/close: 价格（转为字符串）
            - chg: 涨跌额（转为字符串）
            - percent: 涨跌幅（%）
            - turnoverrate: 换手率（港股接口无此数据，设为NULL）
            - amount: 成交额（港币）
            - pb/pe/ps: 估值指标（港股接口无此数据，设为NULL）
            - pre_close: 昨收价
        """
        if df.empty:
            return pd.DataFrame()
        
        # 创建符合gp_day_data表结构的DataFrame
        result_df = pd.DataFrame()
        
        # 字段映射和转换（与A股V3版本保持一致）
        result_df['symbol'] = df['ts_code'].apply(self.convert_tushare_code_to_db)  # 纯数字5位
        result_df['timestamp'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
        result_df['volume'] = df['vol']  # 港股vol单位是手，保持不变
        result_df['open'] = df['open'].astype(str)  # 转为字符串与A股表结构一致
        result_df['high'] = df['high'].astype(str)
        result_df['low'] = df['low'].astype(str)
        result_df['close'] = df['close'].astype(str)
        result_df['chg'] = df['change'].astype(str)
        result_df['percent'] = df['pct_chg']
        result_df['turnoverrate'] = None  # 港股日线接口没有换手率数据
        result_df['amount'] = df['amount'] * 1000  # Tushare的amount单位是千港币，转换为港币
        result_df['pb'] = None  # 港股日线接口没有PB数据，保持NULL
        result_df['pe'] = None  # 港股日线接口没有PE数据，保持NULL
        result_df['ps'] = None  # 港股日线接口没有PS数据，保持NULL
        result_df['pre_close'] = df['pre_close']
        
        # 按时间排序
        result_df = result_df.sort_values('timestamp').reset_index(drop=True)
        
        return result_df

    def check_ex_rights_before_save(self, df, date_str):
        """
        在保存数据库之前检查除权情况，返回除权股票列表和除权日志数据
        
        Args:
            df: 待保存的日线数据DataFrame
            date_str: 当前采集日期，格式YYYY-MM-DD
            
        Returns:
            tuple: (除权股票列表, 除权日志数据列表)
        """
        print("正在检查港股除权情况...")
        
        ex_rights_stocks = []
        ex_rights_log_data = []
        
        from tqdm import tqdm
        
        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 {table_name}
                WHERE symbol = :symbol 
                ORDER BY `timestamp` DESC 
                LIMIT 2
            """.format(table_name=self.table_name))
            
            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 == date_str and len(results) > 1:
                        # 如果最新记录是今天的，且有第二条记录，则用第二条记录比较
                        db_last_close = float(results[1][0])
                    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({
                            'stock_code': symbol,
                            'change_date': date_str,
                            'before_price': db_last_close,
                            'after_price': current_pre_close,
                            'update_time': datetime.now()
                        })
            except Exception as 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):
        """
        将除权日志保存到数据库
        
        Args:
            log_data: 除权日志数据列表
        """
        if not log_data:
            return
        
        print(f"正在保存 {len(log_data)} 条港股除权日志到gp_ex_rights_log表...")
        try:
            df = pd.DataFrame(log_data)
            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, days=180):
        """
        处理除权港股：保存日志、删除历史数据并重新获取
        
        Args:
            ex_rights_stocks: 除权股票列表
            ex_rights_log_data: 除权日志数据列表
            days: 重新获取的历史天数，默认180天
        """
        if not ex_rights_stocks:
            return
            
        print("=" * 60)
        print("开始处理除权港股...")
        print("=" * 60)
        
        # 1. 保存除权日志
        if ex_rights_log_data:
            self.save_ex_rights_log(ex_rights_log_data)
        
        # 2. 重新获取历史数据
        print(f"开始处理 {len(ex_rights_stocks)} 只除权港股，重新获取历史数据...")
        
        success_count = 0
        fail_count = 0
        
        from tqdm import tqdm
        
        for symbol in tqdm(ex_rights_stocks, desc="处理除权港股"):
            try:
                # 转换为Tushare格式
                ts_code = self.convert_db_code_to_tushare(symbol)
                
                # 删除该港股的所有历史数据（限制LENGTH(symbol)=5确保只删港股）
                delete_query = text(
                    f"DELETE FROM {self.table_name} "
                    f"WHERE symbol = :symbol "
                    f"AND LENGTH(symbol) = 5"
                )
                with self.engine.begin() as conn:
                    result = conn.execute(delete_query, {"symbol": symbol})
                    print(f"已删除 {symbol} 的 {result.rowcount} 条历史数据")
                
                # 计算日期范围
                end_date_obj = datetime.now()
                start_date_obj = end_date_obj - timedelta(days=days)
                start_date_str = start_date_obj.strftime('%Y%m%d')
                end_date_str = end_date_obj.strftime('%Y%m%d')
                
                # 重新获取历史数据
                df = self.fetch_hk_daily_data(
                    ts_code=ts_code,
                    start_date=start_date_str,
                    end_date=end_date_str
                )
                
                if df.empty:
                    print(f"未获取到{symbol}的历史数据")
                    fail_count += 1
                    continue
                
                # 转换数据格式
                result_df = self.transform_data(df)
                
                if result_df.empty:
                    print(f"数据转换失败: {symbol}")
                    fail_count += 1
                    continue
                
                # 保存到数据库
                result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
                print(f"成功重新获取 {symbol} 的历史数据，共{len(result_df)}条记录")
                success_count += 1
                    
            except Exception as e:
                print(f"处理除权港股 {symbol} 失败: {e}")
                fail_count += 1
        
        print("\n" + "=" * 60)
        print("除权港股处理完成")
        print("=" * 60)
        print(f"成功: {success_count}只")
        print(f"失败: {fail_count}只")
        print("=" * 60)

    def run_daily_collection(self, date=None):
        """
        执行每日增量采集
        按日期获取全市场港股日线数据（推荐方式）
        
        Args:
            date: 指定日期，格式YYYY-MM-DD，默认为今天
        """
        print("=" * 60)
        print("开始执行每日增量采集（全市场）")
        print("=" * 60)
        
        try:
            # 确定采集日期
            if date is None:
                date_str = datetime.now().strftime('%Y-%m-%d')
            else:
                date_str = date
            
            # 转换为YYYYMMDD格式
            date_obj = datetime.strptime(date_str, '%Y-%m-%d')
            trade_date = date_obj.strftime('%Y%m%d')
            
            print(f"采集日期: {date_str} ({trade_date})")
            
            # 删除指定日期的旧港股数据（港股代码长度为5位）
            with self.engine.begin() as conn:
                delete_query = text(
                    f"DELETE FROM {self.table_name} "
                    f"WHERE DATE(`timestamp`) = :date_str "
                    f"AND LENGTH(symbol) = 5"
                )
                result = conn.execute(delete_query, {"date_str": date_str})
                print(f"已删除{self.table_name}表中{date_str}的旧港股数据（LENGTH(symbol)=5），共{result.rowcount}条")
            
            # 获取全市场港股日线数据
            print(f"正在获取{date_str}全市场港股日线数据...")
            df = self.fetch_hk_daily_data(trade_date=trade_date)
            
            if df.empty:
                print(f"{date_str}没有港股日线数据（可能是非交易日）")
                return
            
            print(f"成功获取港股日线数据，共{len(df)}只股票")
            
            # 转换数据格式
            result_df = self.transform_data(df)
            
            if result_df.empty:
                print("数据转换失败")
                return
            
            # 检查除权（保存前）
            ex_rights_stocks, ex_rights_log_data = self.check_ex_rights_before_save(result_df, date_str)
            
            # 保存到数据库
            print(f"正在保存数据到{self.table_name}表...")
            result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
            print(f"成功保存{len(result_df)}条记录")
            
            # 处理除权股票（保存后）
            self.handle_ex_rights_stocks(ex_rights_stocks, ex_rights_log_data)
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("每日增量采集完成")
            print("=" * 60)
            print(f"采集日期: {date_str}")
            print(f"股票数量: {len(result_df)}只")
            print(f"除权处理: {len(ex_rights_stocks)}只")
            print(f"数据表: {self.table_name}")
            print("=" * 60)
            
        except Exception as e:
            print(f"每日增量采集失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()

    def run_full_collection(self, start_date=None, end_date=None):
        """
        执行全量历史数据采集（按股票逐个采集）
        
        Args:
            start_date: 开始日期，格式YYYY-MM-DD，默认为180天前
            end_date: 结束日期，格式YYYY-MM-DD，默认为今天
        """
        print("=" * 60)
        print("开始执行全量历史数据采集（按股票逐个采集）")
        print("=" * 60)
        
        try:
            # 确定日期范围
            if end_date is None:
                end_date_obj = datetime.now()
            else:
                end_date_obj = datetime.strptime(end_date, '%Y-%m-%d')
            
            if start_date is None:
                # 默认采集近180天历史数据
                start_date_obj = end_date_obj - timedelta(days=180)
            else:
                start_date_obj = datetime.strptime(start_date, '%Y-%m-%d')
            
            start_date_str = start_date_obj.strftime('%Y%m%d')
            end_date_str = end_date_obj.strftime('%Y%m%d')
            
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print("=" * 60)
            
            # 获取港股代码列表
            stock_codes = self.fetch_hk_stock_list()
            
            if not stock_codes:
                print("没有找到港股代码，请先在gp_code_hk表中添加港股代码")
                return
            
            print(f"开始采集{len(stock_codes)}只港股的历史数据...")
            print("=" * 60)
            
            # 按股票逐个采集
            success_count = 0
            fail_count = 0
            total_records = 0
            
            for i, symbol in enumerate(stock_codes, 1):
                try:
                    # 转换为Tushare格式
                    ts_code = self.convert_db_code_to_tushare(symbol)
                    
                    # 获取该股票历史数据
                    df = self.fetch_hk_daily_data(
                        ts_code=ts_code,
                        start_date=start_date_str,
                        end_date=end_date_str
                    )
                    
                    if not df.empty:
                        # 删除该港股在该日期范围的旧数据（限制LENGTH(symbol)=5确保只删港股）
                        with self.engine.begin() as conn:
                            delete_query = text(
                                f"DELETE FROM {self.table_name} "
                                f"WHERE symbol = :symbol "
                                f"AND LENGTH(symbol) = 5 "
                                f"AND `timestamp` BETWEEN :start_date AND :end_date"
                            )
                            conn.execute(delete_query, {
                                "symbol": symbol,
                                "start_date": start_date_obj.strftime('%Y-%m-%d'),
                                "end_date": end_date_obj.strftime('%Y-%m-%d')
                            })
                        
                        # 转换数据格式
                        result_df = self.transform_data(df)
                        
                        if not result_df.empty:
                            # 保存到数据库
                            result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
                            total_records += len(result_df)
                            success_count += 1
                    
                    # 每10只股票显示一次进度
                    if i % 10 == 0:
                        print(f"进度: {i}/{len(stock_codes)} ({success_count}成功, {fail_count}失败, 共{total_records}条记录)")
                    
                except Exception as e:
                    fail_count += 1
                    if i % 10 == 0:
                        print(f"采集 {symbol} 失败: {e}")
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("全量历史数据采集完成")
            print("=" * 60)
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print(f"成功采集: {success_count}只")
            print(f"失败: {fail_count}只")
            print(f"总记录数: {total_records}条")
            print(f"数据表: {self.table_name}")
            print("=" * 60)
            
        except Exception as e:
            print(f"全量历史数据采集失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()

    def run_date_range_collection(self, start_date=None, end_date=None):
        """
        按日期范围逐日采集全市场数据（推荐用于历史数据采集）
        
        Args:
            start_date: 开始日期，格式YYYY-MM-DD，默认为30天前
            end_date: 结束日期，格式YYYY-MM-DD，默认为今天
        """
        print("=" * 60)
        print("开始执行日期范围采集（按日期逐日采集全市场）")
        print("=" * 60)
        
        try:
            # 确定日期范围
            if end_date is None:
                end_date_obj = datetime.now()
            else:
                end_date_obj = datetime.strptime(end_date, '%Y-%m-%d')
            
            if start_date is None:
                # 默认采集近30天历史数据
                start_date_obj = end_date_obj - timedelta(days=30)
            else:
                start_date_obj = datetime.strptime(start_date, '%Y-%m-%d')
            
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print("=" * 60)
            
            # 按日期逐日采集
            current_date = start_date_obj
            total_days = (end_date_obj - start_date_obj).days + 1
            processed_days = 0
            success_days = 0
            total_records = 0
            
            while current_date <= end_date_obj:
                date_str = current_date.strftime('%Y-%m-%d')
                trade_date = current_date.strftime('%Y%m%d')
                
                print(f"\n[{processed_days + 1}/{total_days}] 采集日期: {date_str}")
                sleep(31)
                try:
                    # 获取全市场数据
                    df = self.fetch_hk_daily_data(trade_date=trade_date)
                    
                    if not df.empty:
                        # 转换数据格式
                        result_df = self.transform_data(df)
                        
                        if not result_df.empty:
                            # 删除该日期旧港股数据（LENGTH(symbol)=5）
                            with self.engine.begin() as conn:
                                delete_query = text(
                                    f"DELETE FROM {self.table_name} "
                                    f"WHERE DATE(`timestamp`) = :date_str "
                                    f"AND LENGTH(symbol) = 5"
                                )
                                conn.execute(delete_query, {"date_str": date_str})
                            
                            # 保存新数据
                            result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
                            total_records += len(result_df)
                            success_days += 1
                            print(f"  成功: {len(result_df)}条记录")
                    else:
                        print(f"  无数据（可能是非交易日）")
                    
                except Exception as e:
                    print(f"  采集失败: {e}")
                
                # 移到下一天
                current_date += timedelta(days=1)
                processed_days += 1
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("日期范围采集完成")
            print("=" * 60)
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print(f"总天数: {total_days}")
            print(f"成功天数: {success_days}")
            print(f"总记录数: {total_records}条")
            print(f"数据表: {self.table_name}")
            print("=" * 60)
            
        except Exception as e:
            print(f"日期范围采集失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()


def collect_hk_daily_data(db_url, tushare_token, mode='daily', date=None, start_date=None, end_date=None):
    """
    港股日线数据采集入口函数
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare API Token
        mode: 采集模式
            - 'daily': 增量采集当日数据（按日期获取全市场，推荐）
            - 'full': 全量覆盖采集（按股票逐个采集）
            - 'range': 日期范围采集（按日期逐日采集全市场，推荐用于历史数据）
        date: 指定日期（仅mode='daily'时有效），格式YYYY-MM-DD
        start_date: 开始日期（mode='full'或'range'时有效），格式YYYY-MM-DD
        end_date: 结束日期（mode='full'或'range'时有效），格式YYYY-MM-DD
    
    注意：
        - 港股数据会保存到gp_day_data表（与A股共用）
        - symbol格式为纯数字5位（如01810），不带市场编码
        - 删除数据时只删除LENGTH(symbol)=5的记录，不影响A股数据
    """
    collector = HKDailyDataCollector(db_url, tushare_token)
    
    if mode == 'daily':
        collector.run_daily_collection(date)
    elif mode == 'full':
        collector.run_full_collection(start_date, end_date)
    elif mode == 'range':
        collector.run_date_range_collection(start_date, end_date)
    else:
        print(f"未知的采集模式: {mode}")


if __name__ == "__main__":
    # 数据库配置
    db_url = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    
    # Tushare Token（从config.py读取）
    tushare_token = TUSHARE_TOKEN
    
    # 使用方式示例：
    
    # 1. 每日增量采集（推荐每日17:00后使用，一次性获取全市场数据）
    # collect_hk_daily_data(db_url, tushare_token, mode='daily')
    
    # 2. 日期范围采集（推荐用于历史数据，按日期逐日采集全市场）
    # collect_hk_daily_data(db_url, tushare_token, mode='range',
    #                       start_date='2022-01-01', end_date='2025-11-03')
    # 港股的采集到了11月4日的数据（待采集的）
    # 3. 采集指定日期的数据
    collect_hk_daily_data(db_url, tushare_token, mode='daily', date='2025-11-04')
    
    # 4. 全量覆盖采集（按股票逐个采集，较慢，不推荐）
    # collect_hk_daily_data(db_url, tushare_token, mode='full',
    #                       start_date='2024-01-01', end_date='2024-12-31')
    # import akshare as ak
    # stock_zh_a_spot_df = ak.stock_zh_a_spot()
    # print(stock_zh_a_spot_df)

