# coding:utf-8
"""
同花顺板块数据转存脚本
功能：将 ths_industry_member 和 ths_concept_member 表的数据转存到 gp_hybk 和 gp_gnbk 表。但是提前需要处理ths_concept_member_collector.py和ths_industry_member_collector.py里面的数据。
说明：全量清理目标表后重新插入，保持数据结构对齐
"""

import os
import sys
import re
from typing import Optional

import pandas as pd
from sqlalchemy import create_engine, text

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


class THSToGPBKTransfer:
    """同花顺板块数据转存器"""

    def __init__(self, db_url: str):
        """
        Args:
            db_url: 数据库连接 URL
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600,
        )

        print("=" * 60)
        print("同花顺板块数据转存工具")
        print("=" * 60)

    @staticmethod
    def extract_bk_code(code: str) -> Optional[int]:
        """
        从板块代码中提取数字部分（如：885800.TI -> 885800）
        
        Args:
            code: 板块代码，如 '885800.TI' 或 '885800'
            
        Returns:
            提取的数字，如果无法提取则返回 None
        """
        if not code:
            return None
        # 提取数字部分
        match = re.search(r'(\d+)', str(code))
        if match:
            try:
                return int(match.group(1))
            except ValueError:
                return None
        return None

    def transfer_industry_data(self) -> dict:
        """
        将 ths_industry_member 表数据转存到 gp_hybk 表
        
        Returns:
            转存结果统计
        """
        print("\n" + "=" * 60)
        print("开始转存行业板块数据 (ths_industry_member -> gp_hybk)")
        print("=" * 60)

        result = {
            'success': False,
            'source_count': 0,
            'transferred_count': 0,
            'error': None
        }

        try:
            # 1. 读取源表数据
            print("正在读取 ths_industry_member 表数据...")
            source_query = text("""
                SELECT 
                    industry_code,
                    industry_name,
                    stock_symbol,
                    stock_name
                FROM ths_industry_member
                WHERE industry_code IS NOT NULL 
                  AND stock_symbol IS NOT NULL
            """)
            
            with self.engine.connect() as conn:
                source_df = pd.read_sql(source_query, conn)
            
            result['source_count'] = len(source_df)
            print(f"源表共 {result['source_count']} 条记录")

            if source_df.empty:
                print("源表无数据，转存终止")
                result['success'] = True
                return result

            # 2. 数据转换
            print("正在转换数据格式...")
            target_df = pd.DataFrame()
            target_df['bk_code'] = source_df['industry_code'].apply(self.extract_bk_code)
            target_df['bk_name'] = source_df['industry_name']
            target_df['gp_code'] = source_df['stock_symbol']
            target_df['gp_name'] = source_df['stock_name']
            
            # 过滤掉 bk_code 为 None 的记录
            target_df = target_df[target_df['bk_code'].notna()]

            print(f"转换后共 {len(target_df)} 条有效记录")

            # 3. 清空目标表
            print("正在清空 gp_hybk 表...")
            with self.engine.begin() as conn:
                conn.execute(text("TRUNCATE TABLE gp_hybk"))
            print("gp_hybk 表已清空")

            # 4. 插入数据
            print("正在插入数据到 gp_hybk 表...")
            target_df.to_sql(
                'gp_hybk',
                self.engine,
                if_exists='append',
                index=False,
                method='multi',
                chunksize=1000
            )
            
            result['transferred_count'] = len(target_df)
            result['success'] = True
            print(f"成功转存 {result['transferred_count']} 条记录到 gp_hybk 表")

        except Exception as e:
            result['error'] = str(e)
            print(f"转存行业板块数据失败: {e}")
            import traceback
            traceback.print_exc()

        return result

    def transfer_concept_data(self) -> dict:
        """
        将 ths_concept_member 表数据转存到 gp_gnbk 表
        
        Returns:
            转存结果统计
        """
        print("\n" + "=" * 60)
        print("开始转存概念板块数据 (ths_concept_member -> gp_gnbk)")
        print("=" * 60)

        result = {
            'success': False,
            'source_count': 0,
            'transferred_count': 0,
            'error': None
        }

        try:
            # 1. 读取源表数据
            print("正在读取 ths_concept_member 表数据...")
            source_query = text("""
                SELECT 
                    concept_code,
                    concept_name,
                    stock_symbol,
                    stock_name
                FROM ths_concept_member
                WHERE concept_code IS NOT NULL 
                  AND stock_symbol IS NOT NULL
            """)
            
            with self.engine.connect() as conn:
                source_df = pd.read_sql(source_query, conn)
            
            result['source_count'] = len(source_df)
            print(f"源表共 {result['source_count']} 条记录")

            if source_df.empty:
                print("源表无数据，转存终止")
                result['success'] = True
                return result

            # 2. 数据转换
            print("正在转换数据格式...")
            target_df = pd.DataFrame()
            target_df['bk_code'] = source_df['concept_code'].apply(self.extract_bk_code)
            target_df['bk_name'] = source_df['concept_name']
            target_df['gp_code'] = source_df['stock_symbol']
            target_df['gp_name'] = source_df['stock_name']
            
            # 过滤掉 bk_code 为 None 的记录
            target_df = target_df[target_df['bk_code'].notna()]

            print(f"转换后共 {len(target_df)} 条有效记录")

            # 3. 清空目标表
            print("正在清空 gp_gnbk 表...")
            with self.engine.begin() as conn:
                conn.execute(text("TRUNCATE TABLE gp_gnbk"))
            print("gp_gnbk 表已清空")

            # 4. 插入数据
            print("正在插入数据到 gp_gnbk 表...")
            target_df.to_sql(
                'gp_gnbk',
                self.engine,
                if_exists='append',
                index=False,
                method='multi',
                chunksize=1000
            )
            
            result['transferred_count'] = len(target_df)
            result['success'] = True
            print(f"成功转存 {result['transferred_count']} 条记录到 gp_gnbk 表")

        except Exception as e:
            result['error'] = str(e)
            print(f"转存概念板块数据失败: {e}")
            import traceback
            traceback.print_exc()

        return result

    def run_full_transfer(self) -> dict:
        """
        执行全量转存：行业 + 概念
        
        Returns:
            转存结果汇总
        """
        print("\n" + "=" * 60)
        print("开始执行全量转存")
        print("=" * 60)

        summary = {
            'industry': None,
            'concept': None,
            'all_success': False
        }

        # 转存行业数据
        industry_result = self.transfer_industry_data()
        summary['industry'] = industry_result

        # 转存概念数据
        concept_result = self.transfer_concept_data()
        summary['concept'] = concept_result

        # 汇总结果
        summary['all_success'] = industry_result.get('success') and concept_result.get('success')

        # 打印汇总
        print("\n" + "=" * 60)
        print("转存完成汇总")
        print("=" * 60)
        print(f"行业板块转存: {'成功' if industry_result.get('success') else '失败'}")
        if industry_result.get('success'):
            print(f"  - 源表记录数: {industry_result.get('source_count', 0)}")
            print(f"  - 转存记录数: {industry_result.get('transferred_count', 0)}")
        if industry_result.get('error'):
            print(f"  - 错误: {industry_result.get('error')}")

        print(f"\n概念板块转存: {'成功' if concept_result.get('success') else '失败'}")
        if concept_result.get('success'):
            print(f"  - 源表记录数: {concept_result.get('source_count', 0)}")
            print(f"  - 转存记录数: {concept_result.get('transferred_count', 0)}")
        if concept_result.get('error'):
            print(f"  - 错误: {concept_result.get('error')}")

        print("=" * 60)

        return summary


def transfer_ths_to_gp_bk(db_url: str):
    """
    转存入口函数
    
    Args:
        db_url: 数据库连接 URL
    """
    transfer = THSToGPBKTransfer(db_url)
    transfer.run_full_transfer()


if __name__ == "__main__":
    # DB_URL = "mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj"
    DB_URL = "mysql+pymysql://fac_pattern:Chlry$%.8_app@192.168.16.153:3307/my_quant_db"
    
    # 执行全量转存
    transfer_ths_to_gp_bk(DB_URL)

