# coding:utf-8

"""
A股股票列表更新工具
功能：从Tushare获取A股股票列表，更新gp_code_all表
API文档: https://tushare.pro/document/2?doc_id=25
"""

import tushare as ts
import pandas as pd
from sqlalchemy import create_engine, text
from datetime import datetime
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 StockListUpdater:
    """A股股票列表更新器"""

    def __init__(self, db_url, tushare_token):
        """
        初始化更新器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        
        # 初始化Tushare
        ts.set_token(tushare_token)
        self.pro = ts.pro_api()
        
        print("=" * 60)
        print("A股股票列表更新工具")
        print("=" * 60)

    def convert_tushare_code_to_db(self, ts_code):
        """
        将Tushare格式的股票代码转换为数据库格式
        Tushare格式：600000.SH、000001.SZ
        数据库格式：SH600000、SZ000001
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            数据库格式的股票代码
        """
        if '.SH' in ts_code:
            return f"SH{ts_code.replace('.SH', '')}"
        elif '.SZ' in ts_code:
            return f"SZ{ts_code.replace('.SZ', '')}"
        elif '.BJ' in ts_code:
            return f"BJ{ts_code.replace('.BJ', '')}"
        else:
            return ts_code

    def fetch_stock_list_from_tushare(self):
        """
        从Tushare获取A股股票列表
        API: stock_basic
        文档: https://tushare.pro/document/2?doc_id=25
        
        Returns:
            股票列表DataFrame
        """
        try:
            print("正在从Tushare获取A股股票列表...")
            
            # 获取A股股票列表（上市状态：L=上市，D=退市，P=暂停上市）
            # 只获取上市状态的股票
            df = self.pro.stock_basic(
                exchange='',  # 空值表示获取所有交易所
                list_status='L',  # L=上市
                fields='ts_code,symbol,name,area,industry,market,list_date,fullname,enname,cnspell,exchange,curr_type,list_status,delist_date,is_hs,act_name,act_ent_type'
            )
            
            if df.empty:
                print("未获取到股票数据")
                return pd.DataFrame()
            
            print(f"成功获取{len(df)}只A股股票信息")
            return df
            
        except Exception as e:
            print(f"获取股票列表失败: {e}")
            return pd.DataFrame()

    def get_existing_stocks_from_db(self):
        """
        从数据库获取现有的股票列表
        
        Returns:
            现有股票代码列表（数据库格式）
        """
        try:
            query = "SELECT gp_code, gp_name FROM gp_code_all"
            df = pd.read_sql(query, self.engine)
            
            print(f"数据库现有股票数量: {len(df)}")
            return df
            
        except Exception as e:
            print(f"查询数据库失败: {e}")
            return pd.DataFrame()

    def transform_to_db_format(self, df_tushare):
        """
        将Tushare数据转换为数据库格式
        
        Args:
            df_tushare: Tushare返回的股票列表DataFrame
            
        Returns:
            转换后的DataFrame
        """
        if df_tushare.empty:
            return pd.DataFrame()
        
        # 创建新的DataFrame
        df_result = pd.DataFrame()
        
        # 字段映射
        df_result['gp_name'] = df_tushare['name']  # 股票名称
        df_result['gp_code'] = df_tushare['ts_code'].apply(self.convert_tushare_code_to_db)  # 主代码（如SH600000）
        df_result['gp_code_two'] = df_tushare['ts_code']  # Tushare格式代码（如300918.SZ）
        df_result['gp_code_three'] = None  # 不需要填充
        
        # market_cap字段为空，需要通过其他接口获取（这里暂不处理）
        df_result['market_cap'] = None
        
        # mark1和mark2字段可用于标记市场类型
        # 根据交易所代码设置标记
        df_result['market'] = df_tushare['market'].apply(lambda x: '主板' if x == '主板' else ('科创板' if x == '科创板' else ('创业板' if x == '创业板' else '其他')))
        df_result['area'] = df_tushare['area']
        df_result['industry'] = df_tushare['industry']
        df_result['fullname'] = df_tushare['fullname']
        df_result['enname'] = df_tushare['enname']
        df_result['cnspell'] = df_tushare['cnspell']
        df_result['exchange'] = df_tushare['exchange']
        df_result['curr_type'] = df_tushare['curr_type']
        df_result['list_status'] = df_tushare['list_status']
        df_result['list_date'] = df_tushare['list_date']
        df_result['delist_date'] = df_tushare['delist_date']
        df_result['is_hs'] = df_tushare['is_hs']
        df_result['act_name'] = df_tushare['act_name']
        df_result['act_ent_type'] = df_tushare['act_ent_type']

        return df_result

    def update_stock_list(self):
        """
        更新股票列表到数据库
        主流程：
        1. 从Tushare获取最新股票列表
        2. 从数据库获取现有股票列表
        3. 对比找出新增和需要更新的股票
        4. 执行插入和更新操作
        """
        try:
            # 1. 从Tushare获取股票列表
            df_tushare = self.fetch_stock_list_from_tushare()
            if df_tushare.empty:
                print("未获取到Tushare数据，终止更新")
                return
            
            # 2. 转换为数据库格式
            df_new = self.transform_to_db_format(df_tushare)
            print(f"转换数据完成，共{len(df_new)}条记录")
            
            # 3. 从数据库获取现有股票列表
            df_existing = self.get_existing_stocks_from_db()
            
            if df_existing.empty:
                # 数据库为空，全部插入
                print("数据库为空，将插入所有股票...")
                df_new.to_sql('gp_code_all', self.engine, if_exists='append', index=False)
                print(f"成功插入{len(df_new)}只股票")
            else:
                # 4. 对比数据，找出新增、更新和需要删除的股票
                existing_codes = set(df_existing['gp_code'].tolist())
                new_codes = set(df_new['gp_code'].tolist())
                
                # 新增的股票代码
                codes_to_insert = new_codes - existing_codes
                # 需要更新的股票代码（已存在的）
                codes_to_update = new_codes & existing_codes
                # 需要删除的股票代码（数据库有但Tushare没有，可能已退市）
                codes_to_delete = existing_codes - new_codes
                
                print(f"\n数据对比结果:")
                print(f"  数据库现有: {len(existing_codes)}只")
                print(f"  Tushare获取: {len(new_codes)}只")
                print(f"  新增股票: {len(codes_to_insert)}只")
                print(f"  需更新: {len(codes_to_update)}只")
                print(f"  需删除: {len(codes_to_delete)}只（退市或其他原因）")
                
                # 5. 插入新股票
                if codes_to_insert:
                    df_insert = df_new[df_new['gp_code'].isin(codes_to_insert)]
                    df_insert.to_sql('gp_code_all', self.engine, if_exists='append', index=False)
                    print(f"\n新增股票列表:")
                    for _, row in df_insert.iterrows():
                        print(f"  {row['gp_code']} - {row['gp_name']}")
                    print(f"成功插入{len(df_insert)}只新股票")
                else:
                    print("\n没有新增股票")
                
                # 6. 更新已有股票信息
                if codes_to_update:
                    print(f"\n正在更新{len(codes_to_update)}只股票的信息...")
                    update_count = 0
                    
                    df_update = df_new[df_new['gp_code'].isin(codes_to_update)]
                    
                    with self.engine.begin() as conn:
                        for _, row in df_update.iterrows():
                            # 只更新gp_code_all表中有的字段
                            update_query = text("""
                                UPDATE gp_code_all 
                                SET gp_name = :gp_name,
                                    gp_code_two = :gp_code_two,
                                    gp_code_three = :gp_code_three,
                                    market = :market,
                                    area = :area,
                                    fullname = :fullname,
                                    enname = :enname,
                                    cnspell = :cnspell,
                                    exchange = :exchange,
                                    curr_type = :curr_type,
                                    list_status = :list_status,
                                    list_date = :list_date,
                                    delist_date = :delist_date,
                                    is_hs = :is_hs,
                                    act_name = :act_name,
                                    act_ent_type = :act_ent_type
                                WHERE gp_code = :gp_code
                            """)
                            
                            conn.execute(update_query, {
                                'gp_code': row['gp_code'],
                                'gp_name': row['gp_name'],
                                'gp_code_two': row['gp_code_two'],
                                'gp_code_three': row['gp_code_three'],
                                'market': row['market'],
                                'area': row['area'],
                                'fullname': row['fullname'],
                                'enname': row['enname'],
                                'cnspell': row['cnspell'],
                                'exchange': row['exchange'],
                                'curr_type': row['curr_type'],
                                'list_status': row['list_status'],
                                'list_date': row['list_date'],
                                'delist_date': row['delist_date'],
                                'is_hs': row['is_hs'],
                                'act_name': row['act_name'],
                                'act_ent_type': row['act_ent_type']
                            })
                            update_count += 1
                    
                    print(f"成功更新{update_count}只股票信息")
                
                # 7. 删除退市股票
                if codes_to_delete:
                    print(f"\n正在删除{len(codes_to_delete)}只退市股票...")
                    
                    # 显示要删除的股票
                    df_delete = df_existing[df_existing['gp_code'].isin(codes_to_delete)]
                    print(f"\n退市股票列表:")
                    for _, row in df_delete.iterrows():
                        print(f"  {row['gp_code']} - {row['gp_name']}")
                    
                    # 执行删除
                    with self.engine.begin() as conn:
                        for code in codes_to_delete:
                            delete_query = text("DELETE FROM gp_code_all WHERE gp_code = :gp_code")
                            conn.execute(delete_query, {'gp_code': code})
                    
                    print(f"成功删除{len(codes_to_delete)}只退市股票")
                else:
                    print("\n没有需要删除的股票")
            
            # 8. 输出最终统计
            print("\n" + "=" * 60)
            print("股票列表更新完成")
            print("=" * 60)
            
            # 查询更新后的总数
            query = "SELECT COUNT(*) as count FROM gp_code_all"
            result = pd.read_sql(query, self.engine)
            total_count = result['count'].iloc[0]
            print(f"数据库当前股票总数: {total_count}")
            print("=" * 60)
            
        except Exception as e:
            print(f"更新股票列表失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()

    def show_recent_stocks(self, limit=10):
        """
        显示最近上市的股票
        
        Args:
            limit: 显示数量，默认10只
        """
        try:
            query = """
                SELECT gp_code, gp_name, market, area 
                FROM gp_code_all 
                ORDER BY id DESC 
                LIMIT :limit
            """
            df = pd.read_sql(text(query), self.engine, params={'limit': limit})
            
            if not df.empty:
                print(f"\n最近添加的{limit}只股票:")
                print("-" * 60)
                for _, row in df.iterrows():
                    print(f"{row['gp_code']} - {row['gp_name']} ")
            
        except Exception as e:
            print(f"查询失败: {e}")


def update_stock_list(db_url, tushare_token):
    """
    更新A股股票列表的入口函数
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare API Token
    """
    updater = StockListUpdater(db_url, tushare_token)
    updater.update_stock_list()
    updater.show_recent_stocks(10)


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
    
    # 执行更新
    update_stock_list(db_url, tushare_token)

